compiler/IR/Ops.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2019 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include "compiler/IR/Ops.h"

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/SMLoc.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/STLExtras.h"

 namespace mlir {
 namespace iree_compiler {
 namespace IREE {

 //===----------------------------------------------------------------------===//
 // iree.constant
 //===----------------------------------------------------------------------===//

 static ParseResult parseConstantOp(OpAsmParser &parser,
                                    OperationState &result) {
   Attribute valueAttr;
   Type type;
   if (parser.parseLSquare() ||
       parser.parseAttribute(valueAttr, "value", result.attributes) ||
       parser.parseRSquare() ||
       parser.parseOptionalAttributeDict(result.attributes) ||
       parser.parseColonType(type))
     return failure();

   return parser.addTypeToList(type, result.types);
 }

 static void printConstantOp(OpAsmPrinter &p, ConstantOp &op) {
   p << "iree.constant[";
   p.printAttribute(op.getValue());
   p << "] ";
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"value"});

   p << " : ";
   p.printType(op.getType());
 }

 namespace {

 // TODO(gcmn) this is duplicated from MemRefUtils to avoid a circular
 // dependency. Extract op-dependent parts of memref utils to allow reuse.
 MemRefType convertTypeToMemRef(Type type) {
   if (type.isIntOrIndexOrFloat()) {
     return MemRefType::get({}, type, {}, 0);
   } else if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
     return MemRefType::get(tensorType.getShape(), tensorType.getElementType());
   } else if (auto memRefType = type.dyn_cast<MemRefType>()) {
     return MemRefType::get(memRefType.getShape(), memRefType.getElementType());
   } else {
     llvm_unreachable("Unconvertable type");
   }
 }

 }  // namespace

 void ConstantOp::build(Builder *builder, OperationState &state,
                        ElementsAttr value) {
   auto type = convertTypeToMemRef(value.getType());
   return build(builder, state, type, value);
 }

 // TODO(b/134575149): enable folder when we store the correct type.
 // OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
 //   assert(operands.empty() && "constant has no operands");
 //   return getValue();
 // }

 //===----------------------------------------------------------------------===//
 // iree.tensor_to_memref
 //===----------------------------------------------------------------------===//

 static ParseResult parseTensorToMemRefOp(OpAsmParser &parser,
                                          OperationState &state) {
   OpAsmParser::OperandType operand;
   Type operandType;
   Type resultType;
   if (failed(parser.parseLParen()) || failed(parser.parseOperand(operand)) ||
       failed(parser.parseColonType(operandType)) ||
       failed(parser.resolveOperand(operand, operandType, state.operands)) ||
       failed(parser.parseRParen()) ||
       failed(parser.parseColonType(resultType)) ||
       failed(parser.addTypeToList(resultType, state.types))) {
     return failure();
   }
   return success();
 }

 static void printTensorToMemRefOp(OpAsmPrinter &p, TensorToMemRefOp &op) {
   p << "iree.tensor_to_memref(";
   p.printOperand(op.getOperand());
   p << " : ";
   p.printType(op.getOperand()->getType());
   p << ") : ";
   p.printType(op.getType());
 }

 OpFoldResult TensorToMemRefOp::fold(ArrayRef<Attribute> operands) {
   if (auto memrefToTensorOp = dyn_cast_or_null<IREE::MemRefToTensorOp>(
           getOperand()->getDefiningOp())) {
     return memrefToTensorOp.getOperand();
   }

   return {};
 }

 void TensorToMemRefOp::build(Builder *builder, OperationState &state,
                              Value *arg) {
   build(builder, state, convertTypeToMemRef(arg->getType()), arg);
 }

 //===----------------------------------------------------------------------===//
 // iree.memref_to_tensor
 //===----------------------------------------------------------------------===//

 static ParseResult parseMemRefToTensorOp(OpAsmParser &parser,
                                          OperationState &state) {
   OpAsmParser::OperandType operand;
   Type operandType;
   Type resultType;
   if (failed(parser.parseLParen()) || failed(parser.parseOperand(operand)) ||
       failed(parser.parseColonType(operandType)) ||
       failed(parser.resolveOperand(operand, operandType, state.operands)) ||
       failed(parser.parseRParen()) ||
       failed(parser.parseColonType(resultType)) ||
       failed(parser.addTypeToList(resultType, state.types))) {
     return failure();
   }
   return success();
 }

 static void printMemRefToTensorOp(OpAsmPrinter &p, MemRefToTensorOp &op) {
   p << "iree.memref_to_tensor(";
   p.printOperand(op.getOperand());
   p << " : ";
   p.printType(op.getOperand()->getType());
   p << ") : ";
   p.printType(op.getType());
 }

 OpFoldResult MemRefToTensorOp::fold(ArrayRef<Attribute> operands) {
   if (auto tensorToMemRefOp = dyn_cast_or_null<IREE::TensorToMemRefOp>(
           getOperand()->getDefiningOp())) {
     return tensorToMemRefOp.getOperand();
   }

   return {};
 }

 void MemRefToTensorOp::build(Builder *builder, OperationState &state,
                              Value *arg) {
   // TODO(gcmn) Use getTensorType from MemRefUtils when circular dependency can
   // be avoided.
   auto memRefType = arg->getType().cast<MemRefType>();
   auto tensorType =
       RankedTensorType::get(memRefType.getShape(), memRefType.getElementType());
   build(builder, state, tensorType, arg);
 }

 //===----------------------------------------------------------------------===//
 // iree.scalar_to_memref
 //===----------------------------------------------------------------------===//

 static ParseResult parseScalarToMemRefOp(OpAsmParser &parser,
                                          OperationState &state) {
   OpAsmParser::OperandType operand;
   Type operandType;
   Type resultType;
   if (failed(parser.parseLParen()) || failed(parser.parseOperand(operand)) ||
       failed(parser.parseColonType(operandType)) ||
       failed(parser.resolveOperand(operand, operandType, state.operands)) ||
       failed(parser.parseRParen()) ||
       failed(parser.parseColonType(resultType)) ||
       failed(parser.addTypeToList(resultType, state.types))) {
     return failure();
   }
   return success();
 }

 static void printScalarToMemRefOp(OpAsmPrinter &p, ScalarToMemRefOp &op) {
   p << "iree.scalar_to_memref(";
   p.printOperand(op.getOperand());
   p << " : ";
   p.printType(op.getOperand()->getType());
   p << ") : ";
   p.printType(op.getType());
 }

 OpFoldResult ScalarToMemRefOp::fold(ArrayRef<Attribute> operands) {
   if (auto memrefToScalarOp = dyn_cast_or_null<IREE::MemRefToScalarOp>(
           getOperand()->getDefiningOp())) {
     return memrefToScalarOp.getOperand();
   }

   return {};
 }

 void ScalarToMemRefOp::build(Builder *builder, OperationState &state,
                              Value *arg) {
   build(builder, state, convertTypeToMemRef(arg->getType()), arg);
 }

 //===----------------------------------------------------------------------===//
 // iree.memref_to_scalar
 //===----------------------------------------------------------------------===//

 static ParseResult parseMemRefToScalarOp(OpAsmParser &parser,
                                          OperationState &state) {
   OpAsmParser::OperandType operand;
   Type operandType;
   Type resultType;
   if (failed(parser.parseLParen()) || failed(parser.parseOperand(operand)) ||
       failed(parser.parseColonType(operandType)) ||
       failed(parser.resolveOperand(operand, operandType, state.operands)) ||
       failed(parser.parseRParen()) ||
       failed(parser.parseColonType(resultType)) ||
       failed(parser.addTypeToList(resultType, state.types))) {
     return failure();
   }
   return success();
 }

 static void printMemRefToScalarOp(OpAsmPrinter &p, MemRefToScalarOp &op) {
   p << "iree.memref_to_scalar(";
   p.printOperand(op.getOperand());
   p << " : ";
   p.printType(op.getOperand()->getType());
   p << ") : ";
   p.printType(op.getType());
 }

 OpFoldResult MemRefToScalarOp::fold(ArrayRef<Attribute> operands) {
   if (auto scalarToMemRefOp = dyn_cast_or_null<IREE::ScalarToMemRefOp>(
           getOperand()->getDefiningOp())) {
     return scalarToMemRefOp.getOperand();
   }

   return {};
 }

 void MemRefToScalarOp::build(Builder *builder, OperationState &state,
                              Value *arg) {
   build(builder, state, getElementTypeOrSelf(arg), arg);
 }

 //===----------------------------------------------------------------------===//
 // iree.dispatch_region
 //===----------------------------------------------------------------------===//

 void DispatchRegionOp::build(Builder *builder, OperationState &state,
                              ArrayRef<Type> resultTypes, Value *workload,
                              ArrayRef<Value *> operands,
                              ArrayRef<NamedAttribute> attributes) {
   state.addTypes(resultTypes);
   state.addOperands({workload});
   state.addOperands(operands);
   state.addAttributes(attributes);
   state.addRegion();
   state.setOperandListToResizable();
 }

 ParseResult parseDispatchRegionOp(OpAsmParser &parser, OperationState &state) {
   // Parse required workload.
   OpAsmParser::OperandType workloadArg;
   Type workloadArgType;
   if (failed(parser.parseLSquare()) ||
       failed(parser.parseOperand(workloadArg)) ||
       failed(parser.parseColonType(workloadArgType)) ||
       failed(parser.parseRSquare()) ||
       failed(parser.resolveOperand(workloadArg, workloadArgType,
                                    state.operands))) {
     return failure();
   }

   // Parse (optional) args.
   SmallVector<OpAsmParser::OperandType, 16> regionArgs;
   SmallVector<Type, 16> regionArgTypes;
   if (failed(parser.parseLParen())) {
     return failure();
   }
   if (failed(parser.parseOptionalRParen())) {
     SmallVector<OpAsmParser::OperandType, 16> regionOperands;
     auto argsLoc = parser.getCurrentLocation();
     do {
       // Reserve entries in the lists.
       regionArgs.emplace_back();
       regionOperands.emplace_back();
       regionArgTypes.emplace_back();
       if (failed(parser.parseRegionArgument(regionArgs.back())) ||
           failed(parser.parseEqual()) ||
           failed(parser.parseOperand(regionOperands.back())) ||
           failed(parser.parseColonType(regionArgTypes.back()))) {
         return failure();
       }
     } while (succeeded(parser.parseOptionalComma()));
     if (failed(parser.parseRParen()) ||
         failed(parser.resolveOperands(regionOperands, regionArgTypes, argsLoc,
                                       state.operands))) {
       return failure();
     }
   }
   state.setOperandListToResizable();

   // Parse (optional) results.
   if (failed(parser.parseOptionalColonTypeList(state.types))) {
     return failure();
   }

   // Parse region body.
   Region *body = state.addRegion();
   if (failed(parser.parseRegion(*body, regionArgs, regionArgTypes)) ||
       failed(parser.parseOptionalAttributeDict(state.attributes))) {
     return failure();
   }
   return success();
 }

 void printDispatchRegionOp(OpAsmPrinter &p, DispatchRegionOp op) {
   p << "iree.dispatch_region";

   // Print the workload argument.
   p << "[";
   p.printOperand(op.getWorkload());
   p << " : ";
   p.printType(op.getWorkload()->getType());
   p << "]";

   // Print the data argument remapping.
   p << "(";
   interleaveComma(
       llvm::zip(op.getBody().front().getArguments(), op.getArgOperands()), p,
       [&](std::tuple<BlockArgument *, Value *> it) {
         p << *std::get<0>(it) << " = " << *std::get<1>(it);
         p << " : ";
         p << std::get<1>(it)->getType();
       });
   p << ")";

   // Print the result types, if any.
   if (op.getNumResults() > 0) {
     p << " : ";
     interleaveComma(op.getResultTypes(), p);
   }

   p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(op.getAttrs(),
                           /*elidedAttrs=*/{});
 }

 //===----------------------------------------------------------------------===//
 // iree.reduction_region
 //===----------------------------------------------------------------------===//

 void ReductionRegionOp::build(Builder *builder, OperationState &state,
                               ArrayRef<Type> resultTypes, Value *workload,
                               ArrayRef<Value *> operands,
                               ArrayRef<Value *> initialValues,
                               ArrayRef<int64_t> dimensions,
                               ArrayRef<NamedAttribute> attributes) {
   state.addTypes(resultTypes);
   state.addOperands({workload});
   state.addOperands(operands);
   state.addOperands(initialValues);
   state.addAttribute(
       "dimensions",
       DenseIntElementsAttr::get(
           RankedTensorType::get({static_cast<int64_t>(dimensions.size())},
                                 builder->getIntegerType(64)),
           dimensions));
   state.addAttributes(attributes);
   state.addRegion();
   state.setOperandListToResizable();
 }

 void ReductionRegionOp::build(
     Builder *builder, OperationState &state, ArrayRef<Type> resultTypes,
     Value *workload, ArrayRef<Value *> operands,
     ArrayRef<Value *> initialValues, ArrayRef<int64_t> windowDimensions,
     ArrayRef<int64_t> windowStrides, ArrayRef<int64_t> baseDilations,
     ArrayRef<int64_t> windowDilations, PaddingMode paddingMode,
     ArrayRef<NamedAttribute> attributes) {
   state.addTypes(resultTypes);
   state.addOperands({workload});
   state.addOperands(operands);
   state.addOperands(initialValues);
   state.addAttribute(
       "window_dimensions",
       DenseIntElementsAttr::get(
           RankedTensorType::get({static_cast<int64_t>(windowDimensions.size())},
                                 builder->getIntegerType(64)),
           windowDimensions));
   state.addAttribute(
       "window_strides",
       DenseIntElementsAttr::get(
           RankedTensorType::get({static_cast<int64_t>(windowStrides.size())},
                                 builder->getIntegerType(64)),
           windowStrides));
   state.addAttribute(
       "base_dilations",
       DenseIntElementsAttr::get(
           RankedTensorType::get({static_cast<int64_t>(baseDilations.size())},
                                 builder->getIntegerType(64)),
           baseDilations));
   state.addAttribute(
       "window_dilations",
       DenseIntElementsAttr::get(
           RankedTensorType::get({static_cast<int64_t>(windowDilations.size())},
                                 builder->getIntegerType(64)),
           windowDilations));
   state.addAttribute("padding_mode", builder->getI32IntegerAttr(
                                          static_cast<int32_t>(paddingMode)));
   state.addAttributes(attributes);
   state.addRegion();
   state.setOperandListToResizable();
 }

 ParseResult parseReductionRegionOp(OpAsmParser &parser, OperationState &state) {
   OpAsmParser::OperandType workloadArg;
   Type workloadArgType;
   if (failed(parser.parseLSquare()) ||
       failed(parser.parseOperand(workloadArg)) ||
       failed(parser.parseColonType(workloadArgType)) ||
       failed(parser.parseRSquare()) ||
       failed(parser.resolveOperand(workloadArg, workloadArgType,
                                    state.operands))) {
     return failure();
   }

   SmallVector<OpAsmParser::OperandType, 8> reductionOperands;
   Type reductionType;
   auto operandsLoc = parser.getCurrentLocation();
   if (failed(parser.parseLParen()) ||
       failed(parser.parseOperandList(reductionOperands)) ||
       failed(parser.parseRParen()) ||
       failed(parser.parseColonType(reductionType)) ||
       failed(parser.resolveOperands(
           reductionOperands, reductionType.cast<FunctionType>().getInputs(),
           operandsLoc, state.operands))) {
     return failure();
   }
   for (auto type : reductionType.cast<FunctionType>().getResults()) {
     state.types.push_back(type);
   }
   state.setOperandListToResizable();

   SmallVector<OpAsmParser::OperandType, 8> regionArgs;
   SmallVector<Type, 8> regionArgTypes;
   if (failed(parser.parseKeyword("invocation")) ||
       failed(parser.parseLParen())) {
     return failure();
   }
   do {
     Type argType;
     SmallVector<OpAsmParser::OperandType, 2> reductionRegionArgs;
     OpAsmParser::OperandType initialValue;
     if (failed(parser.parseLParen()) ||
         failed(parser.parseOperandList(reductionRegionArgs, 2)) ||
         failed(parser.parseRParen()) || failed(parser.parseEqual()) ||
         failed(parser.parseOperand(initialValue)) ||
         failed(parser.parseColonType(argType)) ||
         failed(parser.resolveOperand(initialValue, argType, state.operands))) {
       return failure();
     }
     regionArgs.push_back(reductionRegionArgs[0]);
     regionArgTypes.push_back(argType);
     regionArgs.push_back(reductionRegionArgs[1]);
     regionArgTypes.push_back(argType);
   } while (succeeded(parser.parseOptionalComma()));
   if (failed(parser.parseRParen())) {
     return failure();
   }

   // Parse region body.
   Region *body = state.addRegion();
   if (failed(parser.parseRegion(*body, regionArgs, regionArgTypes)) ||
       failed(parser.parseOptionalAttributeDict(state.attributes))) {
     return failure();
   }

   return success();
 }

 void printReductionRegionOp(OpAsmPrinter &p, ReductionRegionOp op) {
   p << "iree.reduction_region";

   // Print the workload argument.
   p << "[";
   p.printOperand(op.getWorkload());
   p << " : ";
   p.printType(op.getWorkload()->getType());
   p << "]";

   p << "(";
   p.printOperands(op.getODSOperands(1));
   p << ")";
   if (op.getNumResults() > 0) {
     p << " : (";
     interleaveComma(op.getODSOperands(1), p,
                     [&](Value *operand) { p.printType(operand->getType()); });
     p << ")";
     p << " -> (";
     interleaveComma(op.getResultTypes(), p);
     p << ")";
   }
   p << "\n";

   p << "      invocation(";
   auto &entryBlock = op.getBody().getBlocks().front();
   int regionArgIndex = 0;
   interleaveComma(op.getODSOperands(2), p, [&](Value *operand) {
     p << "(";
     p.printOperand(entryBlock.getArgument(regionArgIndex++));
     p << ", ";
     p.printOperand(entryBlock.getArgument(regionArgIndex++));
     p << ") = ";
     p.printOperand(operand);
     p << " : ";
     p.printType(operand->getType());
   });
   p << ") ";

   p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(op.getAttrs(),
                           /*elidedAttrs=*/{});
 }

 //===----------------------------------------------------------------------===//
 // iree.return
 //===----------------------------------------------------------------------===//

 static ParseResult parseReturnOp(OpAsmParser &parser, OperationState &state) {
   SmallVector<OpAsmParser::OperandType, 2> opInfo;
   SmallVector<Type, 2> types;
   llvm::SMLoc loc = parser.getCurrentLocation();
   return failure(parser.parseOperandList(opInfo) ||
                  (!opInfo.empty() && parser.parseColonTypeList(types)) ||
                  parser.resolveOperands(opInfo, types, loc, state.operands));
 }

 static void printReturnOp(OpAsmPrinter &p, ReturnOp op) {
   p << "iree.return";
   if (op.getNumOperands() > 0) {
     p << ' ';
     p.printOperands(op.operand_begin(), op.operand_end());
     p << " : ";
     interleaveComma(op.getOperandTypes(), p);
   }
 }

 //===----------------------------------------------------------------------===//
 // iree.load_input
 //===----------------------------------------------------------------------===//

 ParseResult parseLoadInputOp(OpAsmParser &parser, OperationState &state) {
   OpAsmParser::OperandType operand;
   Type argType;
   if (parser.parseLParen() || parser.parseOperand(operand) ||
       parser.parseColonType(argType) || parser.parseRParen() ||
       parser.resolveOperand(operand, argType, state.operands)) {
     return failure();
   }
   Type outputType;
   if (parser.parseColonType(outputType) ||
       parser.addTypeToList(outputType, state.types)) {
     return failure();
   }
   return success();
 }

 void printLoadInputOp(OpAsmPrinter &printer, Operation *op) {
   auto *inputValue = op->getOperand(0);
   auto *outputValue = op->getResult(0);
   printer << op->getName() << '(';
   printer.printOperand(inputValue);
   printer << " : ";
   printer.printType(inputValue->getType());
   printer << ") : ";
   printer.printType(outputValue->getType());
 }

 //===----------------------------------------------------------------------===//
 // iree.store_output
 //===----------------------------------------------------------------------===//

 ParseResult parseStoreOutputOp(OpAsmParser &parser, OperationState &state) {
   OpAsmParser::OperandType op0, op1;
   Type argType0, argType1;
   if (parser.parseLParen() || parser.parseOperand(op0) ||
       parser.parseColonType(argType0) || parser.parseComma() ||
       parser.resolveOperand(op0, argType0, state.operands) ||
       parser.parseOperand(op1) || parser.parseColonType(argType1) ||
       parser.parseRParen() ||
       parser.resolveOperand(op1, argType1, state.operands)) {
     return failure();
   }
   return success();
 }

 void printStoreOutputOp(OpAsmPrinter &printer, Operation *op) {
   auto *inputValue = op->getOperand(0);
   auto *outputValue = op->getOperand(1);
   printer << op->getName() << '(';
   printer.printOperand(inputValue);
   printer << " : ";
   printer.printType(inputValue->getType());
   printer << ", ";
   printer.printOperand(outputValue);
   printer << " : ";
   printer.printType(outputValue->getType());
   printer << ")";
 }

 #define GET_OP_CLASSES
 #include "compiler/IR/Ops.cpp.inc"

 }  // namespace IREE
 }  // namespace iree_compiler
 }  // namespace mlir
	// Copyright 2019 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include "compiler/IR/Ops.h"

	#include "llvm/ADT/SmallVector.h"
	#include "llvm/Support/SMLoc.h"
	#include "mlir/IR/Attributes.h"
	#include "mlir/IR/Builders.h"
	#include "mlir/IR/Diagnostics.h"
	#include "mlir/IR/OpImplementation.h"
	#include "mlir/IR/TypeUtilities.h"
	#include "mlir/IR/Value.h"
	#include "mlir/Support/LogicalResult.h"
	#include "mlir/Support/STLExtras.h"

	namespace mlir {
	namespace iree_compiler {
	namespace IREE {

	//===----------------------------------------------------------------------===//
	// iree.constant
	//===----------------------------------------------------------------------===//

	static ParseResult parseConstantOp(OpAsmParser &parser,
	OperationState &result) {
	Attribute valueAttr;
	Type type;
	if (parser.parseLSquare() \|\|
	parser.parseAttribute(valueAttr, "value", result.attributes) \|\|
	parser.parseRSquare() \|\|
	parser.parseOptionalAttributeDict(result.attributes) \|\|
	parser.parseColonType(type))
	return failure();

	return parser.addTypeToList(type, result.types);
	}

	static void printConstantOp(OpAsmPrinter &p, ConstantOp &op) {
	p << "iree.constant[";
	p.printAttribute(op.getValue());
	p << "] ";
	p.printOptionalAttrDict(op.getAttrs(), /elidedAttrs=/{"value"});

	p << " : ";
	p.printType(op.getType());
	}

	namespace {

	// TODO(gcmn) this is duplicated from MemRefUtils to avoid a circular
	// dependency. Extract op-dependent parts of memref utils to allow reuse.
	MemRefType convertTypeToMemRef(Type type) {
	if (type.isIntOrIndexOrFloat()) {
	return MemRefType::get({}, type, {}, 0);
	} else if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
	return MemRefType::get(tensorType.getShape(), tensorType.getElementType());
	} else if (auto memRefType = type.dyn_cast<MemRefType>()) {
	return MemRefType::get(memRefType.getShape(), memRefType.getElementType());
	} else {
	llvm_unreachable("Unconvertable type");
	}
	}

	} // namespace

	void ConstantOp::build(Builder *builder, OperationState &state,
	ElementsAttr value) {
	auto type = convertTypeToMemRef(value.getType());
	return build(builder, state, type, value);
	}

	// TODO(b/134575149): enable folder when we store the correct type.
	// OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
	// assert(operands.empty() && "constant has no operands");
	// return getValue();
	// }

	//===----------------------------------------------------------------------===//
	// iree.tensor_to_memref
	//===----------------------------------------------------------------------===//

	static ParseResult parseTensorToMemRefOp(OpAsmParser &parser,
	OperationState &state) {
	OpAsmParser::OperandType operand;
	Type operandType;
	Type resultType;
	if (failed(parser.parseLParen()) \|\| failed(parser.parseOperand(operand)) \|\|
	failed(parser.parseColonType(operandType)) \|\|
	failed(parser.resolveOperand(operand, operandType, state.operands)) \|\|
	failed(parser.parseRParen()) \|\|
	failed(parser.parseColonType(resultType)) \|\|
	failed(parser.addTypeToList(resultType, state.types))) {
	return failure();
	}
	return success();
	}

	static void printTensorToMemRefOp(OpAsmPrinter &p, TensorToMemRefOp &op) {
	p << "iree.tensor_to_memref(";
	p.printOperand(op.getOperand());
	p << " : ";
	p.printType(op.getOperand()->getType());
	p << ") : ";
	p.printType(op.getType());
	}

	OpFoldResult TensorToMemRefOp::fold(ArrayRef<Attribute> operands) {
	if (auto memrefToTensorOp = dyn_cast_or_null<IREE::MemRefToTensorOp>(
	getOperand()->getDefiningOp())) {
	return memrefToTensorOp.getOperand();
	}

	return {};
	}

	void TensorToMemRefOp::build(Builder *builder, OperationState &state,
	Value *arg) {
	build(builder, state, convertTypeToMemRef(arg->getType()), arg);
	}

	//===----------------------------------------------------------------------===//
	// iree.memref_to_tensor
	//===----------------------------------------------------------------------===//

	static ParseResult parseMemRefToTensorOp(OpAsmParser &parser,
	OperationState &state) {
	OpAsmParser::OperandType operand;
	Type operandType;
	Type resultType;
	if (failed(parser.parseLParen()) \|\| failed(parser.parseOperand(operand)) \|\|
	failed(parser.parseColonType(operandType)) \|\|
	failed(parser.resolveOperand(operand, operandType, state.operands)) \|\|
	failed(parser.parseRParen()) \|\|
	failed(parser.parseColonType(resultType)) \|\|
	failed(parser.addTypeToList(resultType, state.types))) {
	return failure();
	}
	return success();
	}

	static void printMemRefToTensorOp(OpAsmPrinter &p, MemRefToTensorOp &op) {
	p << "iree.memref_to_tensor(";
	p.printOperand(op.getOperand());
	p << " : ";
	p.printType(op.getOperand()->getType());
	p << ") : ";
	p.printType(op.getType());
	}

	OpFoldResult MemRefToTensorOp::fold(ArrayRef<Attribute> operands) {
	if (auto tensorToMemRefOp = dyn_cast_or_null<IREE::TensorToMemRefOp>(
	getOperand()->getDefiningOp())) {
	return tensorToMemRefOp.getOperand();
	}

	return {};
	}

	void MemRefToTensorOp::build(Builder *builder, OperationState &state,
	Value *arg) {
	// TODO(gcmn) Use getTensorType from MemRefUtils when circular dependency can
	// be avoided.
	auto memRefType = arg->getType().cast<MemRefType>();
	auto tensorType =
	RankedTensorType::get(memRefType.getShape(), memRefType.getElementType());
	build(builder, state, tensorType, arg);
	}

	//===----------------------------------------------------------------------===//
	// iree.scalar_to_memref
	//===----------------------------------------------------------------------===//

	static ParseResult parseScalarToMemRefOp(OpAsmParser &parser,
	OperationState &state) {
	OpAsmParser::OperandType operand;
	Type operandType;
	Type resultType;
	if (failed(parser.parseLParen()) \|\| failed(parser.parseOperand(operand)) \|\|
	failed(parser.parseColonType(operandType)) \|\|
	failed(parser.resolveOperand(operand, operandType, state.operands)) \|\|
	failed(parser.parseRParen()) \|\|
	failed(parser.parseColonType(resultType)) \|\|
	failed(parser.addTypeToList(resultType, state.types))) {
	return failure();
	}
	return success();
	}

	static void printScalarToMemRefOp(OpAsmPrinter &p, ScalarToMemRefOp &op) {
	p << "iree.scalar_to_memref(";
	p.printOperand(op.getOperand());
	p << " : ";
	p.printType(op.getOperand()->getType());
	p << ") : ";
	p.printType(op.getType());
	}

	OpFoldResult ScalarToMemRefOp::fold(ArrayRef<Attribute> operands) {
	if (auto memrefToScalarOp = dyn_cast_or_null<IREE::MemRefToScalarOp>(
	getOperand()->getDefiningOp())) {
	return memrefToScalarOp.getOperand();
	}

	return {};
	}

	void ScalarToMemRefOp::build(Builder *builder, OperationState &state,
	Value *arg) {
	build(builder, state, convertTypeToMemRef(arg->getType()), arg);
	}

	//===----------------------------------------------------------------------===//
	// iree.memref_to_scalar
	//===----------------------------------------------------------------------===//

	static ParseResult parseMemRefToScalarOp(OpAsmParser &parser,
	OperationState &state) {
	OpAsmParser::OperandType operand;
	Type operandType;
	Type resultType;
	if (failed(parser.parseLParen()) \|\| failed(parser.parseOperand(operand)) \|\|
	failed(parser.parseColonType(operandType)) \|\|
	failed(parser.resolveOperand(operand, operandType, state.operands)) \|\|
	failed(parser.parseRParen()) \|\|
	failed(parser.parseColonType(resultType)) \|\|
	failed(parser.addTypeToList(resultType, state.types))) {
	return failure();
	}
	return success();
	}

	static void printMemRefToScalarOp(OpAsmPrinter &p, MemRefToScalarOp &op) {
	p << "iree.memref_to_scalar(";
	p.printOperand(op.getOperand());
	p << " : ";
	p.printType(op.getOperand()->getType());
	p << ") : ";
	p.printType(op.getType());
	}

	OpFoldResult MemRefToScalarOp::fold(ArrayRef<Attribute> operands) {
	if (auto scalarToMemRefOp = dyn_cast_or_null<IREE::ScalarToMemRefOp>(
	getOperand()->getDefiningOp())) {
	return scalarToMemRefOp.getOperand();
	}

	return {};
	}

	void MemRefToScalarOp::build(Builder *builder, OperationState &state,
	Value *arg) {
	build(builder, state, getElementTypeOrSelf(arg), arg);
	}

	//===----------------------------------------------------------------------===//
	// iree.dispatch_region
	//===----------------------------------------------------------------------===//

	void DispatchRegionOp::build(Builder *builder, OperationState &state,
	ArrayRef<Type> resultTypes, Value *workload,
	ArrayRef<Value *> operands,
	ArrayRef<NamedAttribute> attributes) {
	state.addTypes(resultTypes);
	state.addOperands({workload});
	state.addOperands(operands);
	state.addAttributes(attributes);
	state.addRegion();
	state.setOperandListToResizable();
	}

	ParseResult parseDispatchRegionOp(OpAsmParser &parser, OperationState &state) {
	// Parse required workload.
	OpAsmParser::OperandType workloadArg;
	Type workloadArgType;
	if (failed(parser.parseLSquare()) \|\|
	failed(parser.parseOperand(workloadArg)) \|\|
	failed(parser.parseColonType(workloadArgType)) \|\|
	failed(parser.parseRSquare()) \|\|
	failed(parser.resolveOperand(workloadArg, workloadArgType,
	state.operands))) {
	return failure();
	}

	// Parse (optional) args.
	SmallVector<OpAsmParser::OperandType, 16> regionArgs;
	SmallVector<Type, 16> regionArgTypes;
	if (failed(parser.parseLParen())) {
	return failure();
	}
	if (failed(parser.parseOptionalRParen())) {
	SmallVector<OpAsmParser::OperandType, 16> regionOperands;
	auto argsLoc = parser.getCurrentLocation();
	do {
	// Reserve entries in the lists.
	regionArgs.emplace_back();
	regionOperands.emplace_back();
	regionArgTypes.emplace_back();
	if (failed(parser.parseRegionArgument(regionArgs.back())) \|\|
	failed(parser.parseEqual()) \|\|
	failed(parser.parseOperand(regionOperands.back())) \|\|
	failed(parser.parseColonType(regionArgTypes.back()))) {
	return failure();
	}
	} while (succeeded(parser.parseOptionalComma()));
	if (failed(parser.parseRParen()) \|\|
	failed(parser.resolveOperands(regionOperands, regionArgTypes, argsLoc,
	state.operands))) {
	return failure();
	}
	}
	state.setOperandListToResizable();

	// Parse (optional) results.
	if (failed(parser.parseOptionalColonTypeList(state.types))) {
	return failure();
	}

	// Parse region body.
	Region *body = state.addRegion();
	if (failed(parser.parseRegion(*body, regionArgs, regionArgTypes)) \|\|
	failed(parser.parseOptionalAttributeDict(state.attributes))) {
	return failure();
	}
	return success();
	}

	void printDispatchRegionOp(OpAsmPrinter &p, DispatchRegionOp op) {
	p << "iree.dispatch_region";

	// Print the workload argument.
	p << "[";
	p.printOperand(op.getWorkload());
	p << " : ";
	p.printType(op.getWorkload()->getType());
	p << "]";

	// Print the data argument remapping.
	p << "(";
	interleaveComma(
	llvm::zip(op.getBody().front().getArguments(), op.getArgOperands()), p,
	[&](std::tuple<BlockArgument , Value > it) {
	p << std::get<0>(it) << " = " << std::get<1>(it);
	p << " : ";
	p << std::get<1>(it)->getType();
	});
	p << ")";

	// Print the result types, if any.
	if (op.getNumResults() > 0) {
	p << " : ";
	interleaveComma(op.getResultTypes(), p);
	}

	p.printRegion(op.getBody(), /printEntryBlockArgs=/false);
	p.printOptionalAttrDict(op.getAttrs(),
	/elidedAttrs=/{});
	}

	//===----------------------------------------------------------------------===//
	// iree.reduction_region
	//===----------------------------------------------------------------------===//

	void ReductionRegionOp::build(Builder *builder, OperationState &state,
	ArrayRef<Type> resultTypes, Value *workload,
	ArrayRef<Value *> operands,
	ArrayRef<Value *> initialValues,
	ArrayRef<int64_t> dimensions,
	ArrayRef<NamedAttribute> attributes) {
	state.addTypes(resultTypes);
	state.addOperands({workload});
	state.addOperands(operands);
	state.addOperands(initialValues);
	state.addAttribute(
	"dimensions",
	DenseIntElementsAttr::get(
	RankedTensorType::get({static_cast<int64_t>(dimensions.size())},
	builder->getIntegerType(64)),
	dimensions));
	state.addAttributes(attributes);
	state.addRegion();
	state.setOperandListToResizable();
	}

	void ReductionRegionOp::build(
	Builder *builder, OperationState &state, ArrayRef<Type> resultTypes,
	Value workload, ArrayRef<Value > operands,
	ArrayRef<Value *> initialValues, ArrayRef<int64_t> windowDimensions,
	ArrayRef<int64_t> windowStrides, ArrayRef<int64_t> baseDilations,
	ArrayRef<int64_t> windowDilations, PaddingMode paddingMode,
	ArrayRef<NamedAttribute> attributes) {
	state.addTypes(resultTypes);
	state.addOperands({workload});
	state.addOperands(operands);
	state.addOperands(initialValues);
	state.addAttribute(
	"window_dimensions",
	DenseIntElementsAttr::get(
	RankedTensorType::get({static_cast<int64_t>(windowDimensions.size())},
	builder->getIntegerType(64)),
	windowDimensions));
	state.addAttribute(
	"window_strides",
	DenseIntElementsAttr::get(
	RankedTensorType::get({static_cast<int64_t>(windowStrides.size())},
	builder->getIntegerType(64)),
	windowStrides));
	state.addAttribute(
	"base_dilations",
	DenseIntElementsAttr::get(
	RankedTensorType::get({static_cast<int64_t>(baseDilations.size())},
	builder->getIntegerType(64)),
	baseDilations));
	state.addAttribute(
	"window_dilations",
	DenseIntElementsAttr::get(
	RankedTensorType::get({static_cast<int64_t>(windowDilations.size())},
	builder->getIntegerType(64)),
	windowDilations));
	state.addAttribute("padding_mode", builder->getI32IntegerAttr(
	static_cast<int32_t>(paddingMode)));
	state.addAttributes(attributes);
	state.addRegion();
	state.setOperandListToResizable();
	}

	ParseResult parseReductionRegionOp(OpAsmParser &parser, OperationState &state) {
	OpAsmParser::OperandType workloadArg;
	Type workloadArgType;
	if (failed(parser.parseLSquare()) \|\|
	failed(parser.parseOperand(workloadArg)) \|\|
	failed(parser.parseColonType(workloadArgType)) \|\|
	failed(parser.parseRSquare()) \|\|
	failed(parser.resolveOperand(workloadArg, workloadArgType,
	state.operands))) {
	return failure();
	}

	SmallVector<OpAsmParser::OperandType, 8> reductionOperands;
	Type reductionType;
	auto operandsLoc = parser.getCurrentLocation();
	if (failed(parser.parseLParen()) \|\|
	failed(parser.parseOperandList(reductionOperands)) \|\|
	failed(parser.parseRParen()) \|\|
	failed(parser.parseColonType(reductionType)) \|\|
	failed(parser.resolveOperands(
	reductionOperands, reductionType.cast<FunctionType>().getInputs(),
	operandsLoc, state.operands))) {
	return failure();
	}
	for (auto type : reductionType.cast<FunctionType>().getResults()) {
	state.types.push_back(type);
	}
	state.setOperandListToResizable();

	SmallVector<OpAsmParser::OperandType, 8> regionArgs;
	SmallVector<Type, 8> regionArgTypes;
	if (failed(parser.parseKeyword("invocation")) \|\|
	failed(parser.parseLParen())) {
	return failure();
	}
	do {
	Type argType;
	SmallVector<OpAsmParser::OperandType, 2> reductionRegionArgs;
	OpAsmParser::OperandType initialValue;
	if (failed(parser.parseLParen()) \|\|
	failed(parser.parseOperandList(reductionRegionArgs, 2)) \|\|
	failed(parser.parseRParen()) \|\| failed(parser.parseEqual()) \|\|
	failed(parser.parseOperand(initialValue)) \|\|
	failed(parser.parseColonType(argType)) \|\|
	failed(parser.resolveOperand(initialValue, argType, state.operands))) {
	return failure();
	}
	regionArgs.push_back(reductionRegionArgs[0]);
	regionArgTypes.push_back(argType);
	regionArgs.push_back(reductionRegionArgs[1]);
	regionArgTypes.push_back(argType);
	} while (succeeded(parser.parseOptionalComma()));
	if (failed(parser.parseRParen())) {
	return failure();
	}

	// Parse region body.
	Region *body = state.addRegion();
	if (failed(parser.parseRegion(*body, regionArgs, regionArgTypes)) \|\|
	failed(parser.parseOptionalAttributeDict(state.attributes))) {
	return failure();
	}

	return success();
	}

	void printReductionRegionOp(OpAsmPrinter &p, ReductionRegionOp op) {
	p << "iree.reduction_region";

	// Print the workload argument.
	p << "[";
	p.printOperand(op.getWorkload());
	p << " : ";
	p.printType(op.getWorkload()->getType());
	p << "]";

	p << "(";
	p.printOperands(op.getODSOperands(1));
	p << ")";
	if (op.getNumResults() > 0) {
	p << " : (";
	interleaveComma(op.getODSOperands(1), p,
	[&](Value *operand) { p.printType(operand->getType()); });
	p << ")";
	p << " -> (";
	interleaveComma(op.getResultTypes(), p);
	p << ")";
	}
	p << "\n";

	p << " invocation(";
	auto &entryBlock = op.getBody().getBlocks().front();
	int regionArgIndex = 0;
	interleaveComma(op.getODSOperands(2), p, [&](Value *operand) {
	p << "(";
	p.printOperand(entryBlock.getArgument(regionArgIndex++));
	p << ", ";
	p.printOperand(entryBlock.getArgument(regionArgIndex++));
	p << ") = ";
	p.printOperand(operand);
	p << " : ";
	p.printType(operand->getType());
	});
	p << ") ";

	p.printRegion(op.getBody(), /printEntryBlockArgs=/false);
	p.printOptionalAttrDict(op.getAttrs(),
	/elidedAttrs=/{});
	}

	//===----------------------------------------------------------------------===//
	// iree.return
	//===----------------------------------------------------------------------===//

	static ParseResult parseReturnOp(OpAsmParser &parser, OperationState &state) {
	SmallVector<OpAsmParser::OperandType, 2> opInfo;
	SmallVector<Type, 2> types;
	llvm::SMLoc loc = parser.getCurrentLocation();
	return failure(parser.parseOperandList(opInfo) \|\|
	(!opInfo.empty() && parser.parseColonTypeList(types)) \|\|
	parser.resolveOperands(opInfo, types, loc, state.operands));
	}

	static void printReturnOp(OpAsmPrinter &p, ReturnOp op) {
	p << "iree.return";
	if (op.getNumOperands() > 0) {
	p << ' ';
	p.printOperands(op.operand_begin(), op.operand_end());
	p << " : ";
	interleaveComma(op.getOperandTypes(), p);
	}
	}

	//===----------------------------------------------------------------------===//
	// iree.load_input
	//===----------------------------------------------------------------------===//

	ParseResult parseLoadInputOp(OpAsmParser &parser, OperationState &state) {
	OpAsmParser::OperandType operand;
	Type argType;
	if (parser.parseLParen() \|\| parser.parseOperand(operand) \|\|
	parser.parseColonType(argType) \|\| parser.parseRParen() \|\|
	parser.resolveOperand(operand, argType, state.operands)) {
	return failure();
	}
	Type outputType;
	if (parser.parseColonType(outputType) \|\|
	parser.addTypeToList(outputType, state.types)) {
	return failure();
	}
	return success();
	}

	void printLoadInputOp(OpAsmPrinter &printer, Operation *op) {
	auto *inputValue = op->getOperand(0);
	auto *outputValue = op->getResult(0);
	printer << op->getName() << '(';
	printer.printOperand(inputValue);
	printer << " : ";
	printer.printType(inputValue->getType());
	printer << ") : ";
	printer.printType(outputValue->getType());
	}

	//===----------------------------------------------------------------------===//
	// iree.store_output
	//===----------------------------------------------------------------------===//

	ParseResult parseStoreOutputOp(OpAsmParser &parser, OperationState &state) {
	OpAsmParser::OperandType op0, op1;
	Type argType0, argType1;
	if (parser.parseLParen() \|\| parser.parseOperand(op0) \|\|
	parser.parseColonType(argType0) \|\| parser.parseComma() \|\|
	parser.resolveOperand(op0, argType0, state.operands) \|\|
	parser.parseOperand(op1) \|\| parser.parseColonType(argType1) \|\|
	parser.parseRParen() \|\|
	parser.resolveOperand(op1, argType1, state.operands)) {
	return failure();
	}
	return success();
	}

	void printStoreOutputOp(OpAsmPrinter &printer, Operation *op) {
	auto *inputValue = op->getOperand(0);
	auto *outputValue = op->getOperand(1);
	printer << op->getName() << '(';
	printer.printOperand(inputValue);
	printer << " : ";
	printer.printType(inputValue->getType());
	printer << ", ";
	printer.printOperand(outputValue);
	printer << " : ";
	printer.printType(outputValue->getType());
	printer << ")";
	}

	#define GET_OP_CLASSES
	#include "compiler/IR/Ops.cpp.inc"

	} // namespace IREE
	} // namespace iree_compiler
	} // namespace mlir