compiler/Transforms/Interpreter/LowerStdToInterpreterDialect.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 "third_party/llvm/llvm/include/llvm/ADT/DenseSet.h"
 #include "third_party/llvm/llvm/include/llvm/Support/Allocator.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Dialect/StandardOps/Ops.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Attributes.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/BlockAndValueMapping.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Function.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Location.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/MLIRContext.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Module.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/StandardTypes.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/Pass.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/PassRegistry.h"
 #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/DialectConversion.h"
 #include "third_party/mlir_edge/iree/compiler/IR/Dialect.h"
 #include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLDialect.h"
 #include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLOps.h"
 #include "third_party/mlir_edge/iree/compiler/IR/Interpreter/LLDialect.h"
 #include "third_party/mlir_edge/iree/compiler/IR/Ops.h"
 #include "third_party/mlir_edge/iree/compiler/Transforms/ConversionUtils.h"
 #include "third_party/mlir_edge/iree/compiler/Utils/MemRefUtils.h"

 namespace mlir {
 namespace iree_compiler {

 struct ConstantOpLowering : public ConversionPattern {
   explicit ConstantOpLowering(MLIRContext *context)
       : ConversionPattern(ConstantOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto midOp = rewriter.create<IREE::ConstantOp>(
         op->getLoc(), cast<ConstantOp>(op).getValue());

     auto result = wrapAsTensor(midOp.getResult(), op, rewriter);
     rewriter.replaceOp(
         op, {loadResultValue(op->getLoc(), op->getResult(0)->getType(), result,
                              rewriter)});
     return matchSuccess();
   }
 };

 class CallOpLowering : public ConversionPattern {
  public:
   explicit CallOpLowering(MLIRContext *context)
       : ConversionPattern(CallOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto callOp = cast<CallOp>(op);
     auto calleeType = callOp.getCalleeType();
     rewriter.replaceOpWithNewOp<IREEInterp::HL::CallOp>(
         op, callOp.getCallee(), calleeType.getResults(), operands);
     return matchSuccess();
   }
 };

 class CallIndirectOpLowering : public ConversionPattern {
  public:
   explicit CallIndirectOpLowering(MLIRContext *context)
       : ConversionPattern(CallIndirectOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto callOp = cast<CallIndirectOp>(op);
     rewriter.replaceOpWithNewOp<IREEInterp::HL::CallIndirectOp>(
         op, callOp.getCallee(), operands);
     return matchSuccess();
   }
 };

 struct ReturnOpLowering : public ConversionPattern {
   explicit ReturnOpLowering(MLIRContext *context)
       : ConversionPattern(ReturnOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<IREEInterp::HL::ReturnOp>(op, operands);
     return matchSuccess();
   }
 };

 struct BranchOpLowering : public ConversionPattern {
   explicit BranchOpLowering(MLIRContext *context)
       : ConversionPattern(BranchOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> properOperands,
       ArrayRef<Block *> destinations, ArrayRef<ArrayRef<Value *>> operands,
       ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<IREEInterp::HL::BranchOp>(op, destinations[0],
                                                           operands[0]);
     return this->matchSuccess();
   }
 };

 struct CondBranchOpLowering : public ConversionPattern {
   explicit CondBranchOpLowering(MLIRContext *context)
       : ConversionPattern(CondBranchOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> properOperands,
       ArrayRef<Block *> destinations, ArrayRef<ArrayRef<Value *>> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto *condValue =
         loadAccessValue(op->getLoc(), properOperands[0], rewriter);
     rewriter.replaceOpWithNewOp<IREEInterp::HL::CondBranchOp>(
         op, condValue, destinations[IREEInterp::HL::CondBranchOp::trueIndex],
         operands[IREEInterp::HL::CondBranchOp::trueIndex],
         destinations[IREEInterp::HL::CondBranchOp::falseIndex],
         operands[IREEInterp::HL::CondBranchOp::falseIndex]);
     return this->matchSuccess();
   }
 };

 template <typename SrcOp, typename DstOp>
 struct CompareOpLowering : public ConversionPattern {
   explicit CompareOpLowering(MLIRContext *context)
       : ConversionPattern(SrcOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto lhValue = loadAccessValue(op->getLoc(), operands[0], rewriter);
     auto rhValue = loadAccessValue(op->getLoc(), operands[1], rewriter);

     lhValue = wrapAsMemRef(lhValue, op, rewriter);
     rhValue = wrapAsMemRef(rhValue, op, rewriter);

     // TODO(benvanik): map predicate to stable value.
     auto predicate = rewriter.getI32IntegerAttr(
         static_cast<int32_t>(dyn_cast<SrcOp>(op).getPredicate()));

     auto dstType = getMemRefType(op->getResult(0), rewriter);
     auto midOp = rewriter.create<DstOp>(op->getLoc(), dstType, predicate,
                                         lhValue, rhValue);

     auto result = wrapAsTensor(midOp.getResult(), op, rewriter);
     rewriter.replaceOp(
         op, {loadResultValue(op->getLoc(), op->getResult(0)->getType(), result,
                              rewriter)});
     return this->matchSuccess();
   }
 };

 struct CmpIOpLowering
     : public CompareOpLowering<CmpIOp, IREEInterp::HL::CmpIOp> {
   using CompareOpLowering::CompareOpLowering;
 };

 struct CmpFOpLowering
     : public CompareOpLowering<CmpFOp, IREEInterp::HL::CmpFOp> {
   using CompareOpLowering::CompareOpLowering;
 };

 struct AllocOpLowering : public ConversionPattern {
   explicit AllocOpLowering(MLIRContext *context)
       : ConversionPattern(AllocOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     // TODO(benvanik): replace with length computation.
     rewriter.replaceOpWithNewOp<IREEInterp::HL::AllocHeapOp>(
         op, *op->getResultTypes().begin(), operands);
     return matchSuccess();
   }
 };

 struct DeallocOpLowering : public ConversionPattern {
   explicit DeallocOpLowering(MLIRContext *context)
       : ConversionPattern(DeallocOp::getOperationName(), 1, context) {}

   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<IREEInterp::HL::DiscardOp>(op, operands[0]);
     return matchSuccess();
   }
 };

 struct ExtractElementOpLowering : public ConversionPattern {
   explicit ExtractElementOpLowering(MLIRContext *context)
       : ConversionPattern(ExtractElementOp::getOperationName(), 1, context) {}
   PatternMatchResult matchAndRewrite(
       Operation *op, ArrayRef<Value *> operands,
       ConversionPatternRewriter &rewriter) const override {
     auto extractOp = cast<ExtractElementOp>(op);
     Value *memRefInput = wrapAsMemRef(
         loadAccessValue(op->getLoc(), extractOp.getAggregate(), rewriter), op,
         rewriter);

     SmallVector<Value *, 4> indices = {extractOp.indices().begin(),
                                        extractOp.indices().end()};
     rewriter.replaceOpWithNewOp<LoadOp>(op, memRefInput, indices);
     return matchSuccess();
   }
 };

 #define UNARY_OP_LOWERING(StdOpType, IREEOpType)                               \
   struct StdOpType##Lowering : public UnaryOpLowering<StdOpType, IREEOpType> { \
     using UnaryOpLowering::UnaryOpLowering;                                    \
   };

 #define BINARY_OP_LOWERING(StdOpType, IREEOpType)        \
   struct StdOpType##Lowering                             \
       : public BinaryOpLowering<StdOpType, IREEOpType> { \
     using BinaryOpLowering::BinaryOpLowering;            \
   };

 #define TERNARY_OP_LOWERING(StdOpType, IREEOpType)        \
   struct StdOpType##Lowering                              \
       : public TernaryOpLowering<StdOpType, IREEOpType> { \
     using TernaryOpLowering::TernaryOpLowering;           \
   };

 // UNARY_OP_LOWERING(RankOp, IREEInterp::HL::RankOp);
 UNARY_OP_LOWERING(DimOp, IREEInterp::HL::DimOp);
 // UNARY_OP_LOWERING(ShapeOp, IREEInterp::HL::ShapeOp);
 // UNARY_OP_LOWERING(LengthOp, IREEInterp::HL::LengthOp);

 // UNARY_OP_LOWERING(NotOp, IREEInterp::HL::NotOp);
 BINARY_OP_LOWERING(AndOp, IREEInterp::HL::AndOp);
 BINARY_OP_LOWERING(OrOp, IREEInterp::HL::OrOp);
 // BINARY_OP_LOWERING(XorOp, IREEInterp::HL::XorOp);
 // BINARY_OP_LOWERING(ShiftLeftOp, IREEInterp::HL::ShiftLeftOp);
 // BINARY_OP_LOWERING(ShiftRightLogicalOp, IREEInterp::HL::ShiftRightLogicalOp);
 // BINARY_OP_LOWERING(ShiftRightArithmeticOp,
 // IREEInterp::HL::ShiftRightArithmeticOp);

 BINARY_OP_LOWERING(AddIOp, IREEInterp::HL::AddIOp);
 BINARY_OP_LOWERING(AddFOp, IREEInterp::HL::AddFOp);
 BINARY_OP_LOWERING(SubIOp, IREEInterp::HL::SubIOp);
 BINARY_OP_LOWERING(SubFOp, IREEInterp::HL::SubFOp);
 // UNARY_OP_LOWERING(AbsIOp, IREEInterp::HL::AbsIOp);
 // UNARY_OP_LOWERING(AbsFOp, IREEInterp::HL::AbsFOp);
 BINARY_OP_LOWERING(MulIOp, IREEInterp::HL::MulIOp);
 BINARY_OP_LOWERING(MulFOp, IREEInterp::HL::MulFOp);
 BINARY_OP_LOWERING(DivISOp, IREEInterp::HL::DivISOp);
 BINARY_OP_LOWERING(DivIUOp, IREEInterp::HL::DivIUOp);
 BINARY_OP_LOWERING(DivFOp, IREEInterp::HL::DivFOp);
 // BINARY_OP_LOWERING(MulAddIOp, IREEInterp::HL::MulAddIOp);
 // BINARY_OP_LOWERING(MulAddFOp, IREEInterp::HL::MulAddFOp);
 // UNARY_OP_LOWERING(ExpFOp, IREEInterp::HL::ExpFOp);
 // UNARY_OP_LOWERING(LogFOp, IREEInterp::HL::LogFOp);
 // UNARY_OP_LOWERING(RsqrtFOp, IREEInterp::HL::RsqrtFOp);
 // UNARY_OP_LOWERING(CosFOp, IREEInterp::HL::CosFOp);
 // UNARY_OP_LOWERING(SinFOp, IREEInterp::HL::SinFOp);
 // UNARY_OP_LOWERING(TanhFOp, IREEInterp::HL::TanhFOp);
 // UNARY_OP_LOWERING(Atan2FOp, IREEInterp::HL::Atan2FOp);

 // BINARY_OP_LOWERING(MinISOp, IREEInterp::HL::MinISOp);
 // BINARY_OP_LOWERING(MinIUOp, IREEInterp::HL::MinIUOp);
 // BINARY_OP_LOWERING(MinFOp, IREEInterp::HL::MinFOp);
 // BINARY_OP_LOWERING(MaxISOp, IREEInterp::HL::MaxISOp);
 // BINARY_OP_LOWERING(MaxIUOp, IREEInterp::HL::MaxIUOp);
 // BINARY_OP_LOWERING(MaxFOp, IREEInterp::HL::MaxFOp);
 // TERNARY_OP_LOWERING(ClampFOp, IREEInterp::HL::ClampFOp);
 // UNARY_OP_LOWERING(FloorFOp, IREEInterp::HL::FloorFOp);
 // UNARY_OP_LOWERING(CeilFOp, IREEInterp::HL::CeilFOp);

 class LowerStdToInterpreterDialectPass
     : public FunctionPass<LowerStdToInterpreterDialectPass> {
  public:
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     patterns.insert<
         ConstantOpLowering,
         // Control flow.
         CallOpLowering, CallIndirectOpLowering, ReturnOpLowering,
         BranchOpLowering, CondBranchOpLowering, CmpIOpLowering, CmpFOpLowering,
         // Memory management.
         AllocOpLowering, DeallocOpLowering, ExtractElementOpLowering,
         // Shape operations.
         DimOpLowering,
         // Logical ops.
         AndOpLowering, OrOpLowering,
         // Arithmetic ops.
         AddIOpLowering, AddFOpLowering, SubIOpLowering, SubFOpLowering,
         MulIOpLowering, MulFOpLowering, DivISOpLowering, DivIUOpLowering,
         DivFOpLowering>(&getContext());

     ConversionTarget target(getContext());
     target.addLegalDialect<IREEHLInterpreterDialect, IREELLInterpreterDialect,
                            IREEDialect>();
     target.addLegalOp<LoadOp, StoreOp, FuncOp, ModuleOp>();
     target.addDynamicallyLegalOp<ConstantOp>([](ConstantOp constOp) {
       // std.constant is legal for index integers.
       return constOp.getValue().isa<IntegerAttr>() &&
              constOp.getType().isIndex();
     });
     // This is only able to be a full conversion because we know that we do it
     // last. Other conversions (e.g. XLA->IREE) can't be full because we might
     // still have e.g. std ops.
     // TODO(b/139012931) Move fully to the conversion framework.
     if (failed(applyFullConversion(getFunction(), target, patterns))) {
       return signalPassFailure();
     }
   }
 };

 std::unique_ptr<OpPassBase<FuncOp>> createLowerStdToInterpreterDialectPass() {
   return std::make_unique<LowerStdToInterpreterDialectPass>();
 }

 static PassRegistration<LowerStdToInterpreterDialectPass> pass(
     "lower-std-to-iree-interpreter",
     "Convert all std functions to the IREE dialect");

 }  // 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 "third_party/llvm/llvm/include/llvm/ADT/DenseSet.h"
	#include "third_party/llvm/llvm/include/llvm/Support/Allocator.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Dialect/StandardOps/Ops.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Attributes.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/BlockAndValueMapping.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Function.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Location.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/MLIRContext.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Module.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/StandardTypes.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/Pass.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/PassRegistry.h"
	#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/DialectConversion.h"
	#include "third_party/mlir_edge/iree/compiler/IR/Dialect.h"
	#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLDialect.h"
	#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLOps.h"
	#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/LLDialect.h"
	#include "third_party/mlir_edge/iree/compiler/IR/Ops.h"
	#include "third_party/mlir_edge/iree/compiler/Transforms/ConversionUtils.h"
	#include "third_party/mlir_edge/iree/compiler/Utils/MemRefUtils.h"

	namespace mlir {
	namespace iree_compiler {

	struct ConstantOpLowering : public ConversionPattern {
	explicit ConstantOpLowering(MLIRContext *context)
	: ConversionPattern(ConstantOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	auto midOp = rewriter.create<IREE::ConstantOp>(
	op->getLoc(), cast<ConstantOp>(op).getValue());

	auto result = wrapAsTensor(midOp.getResult(), op, rewriter);
	rewriter.replaceOp(
	op, {loadResultValue(op->getLoc(), op->getResult(0)->getType(), result,
	rewriter)});
	return matchSuccess();
	}
	};

	class CallOpLowering : public ConversionPattern {
	public:
	explicit CallOpLowering(MLIRContext *context)
	: ConversionPattern(CallOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	auto callOp = cast<CallOp>(op);
	auto calleeType = callOp.getCalleeType();
	rewriter.replaceOpWithNewOp<IREEInterp::HL::CallOp>(
	op, callOp.getCallee(), calleeType.getResults(), operands);
	return matchSuccess();
	}
	};

	class CallIndirectOpLowering : public ConversionPattern {
	public:
	explicit CallIndirectOpLowering(MLIRContext *context)
	: ConversionPattern(CallIndirectOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	auto callOp = cast<CallIndirectOp>(op);
	rewriter.replaceOpWithNewOp<IREEInterp::HL::CallIndirectOp>(
	op, callOp.getCallee(), operands);
	return matchSuccess();
	}
	};

	struct ReturnOpLowering : public ConversionPattern {
	explicit ReturnOpLowering(MLIRContext *context)
	: ConversionPattern(ReturnOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	rewriter.replaceOpWithNewOp<IREEInterp::HL::ReturnOp>(op, operands);
	return matchSuccess();
	}
	};

	struct BranchOpLowering : public ConversionPattern {
	explicit BranchOpLowering(MLIRContext *context)
	: ConversionPattern(BranchOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > properOperands,
	ArrayRef<Block > destinations, ArrayRef<ArrayRef<Value >> operands,
	ConversionPatternRewriter &rewriter) const override {
	rewriter.replaceOpWithNewOp<IREEInterp::HL::BranchOp>(op, destinations[0],
	operands[0]);
	return this->matchSuccess();
	}
	};

	struct CondBranchOpLowering : public ConversionPattern {
	explicit CondBranchOpLowering(MLIRContext *context)
	: ConversionPattern(CondBranchOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > properOperands,
	ArrayRef<Block > destinations, ArrayRef<ArrayRef<Value >> operands,
	ConversionPatternRewriter &rewriter) const override {
	auto *condValue =
	loadAccessValue(op->getLoc(), properOperands[0], rewriter);
	rewriter.replaceOpWithNewOp<IREEInterp::HL::CondBranchOp>(
	op, condValue, destinations[IREEInterp::HL::CondBranchOp::trueIndex],
	operands[IREEInterp::HL::CondBranchOp::trueIndex],
	destinations[IREEInterp::HL::CondBranchOp::falseIndex],
	operands[IREEInterp::HL::CondBranchOp::falseIndex]);
	return this->matchSuccess();
	}
	};

	template <typename SrcOp, typename DstOp>
	struct CompareOpLowering : public ConversionPattern {
	explicit CompareOpLowering(MLIRContext *context)
	: ConversionPattern(SrcOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	auto lhValue = loadAccessValue(op->getLoc(), operands[0], rewriter);
	auto rhValue = loadAccessValue(op->getLoc(), operands[1], rewriter);

	lhValue = wrapAsMemRef(lhValue, op, rewriter);
	rhValue = wrapAsMemRef(rhValue, op, rewriter);

	// TODO(benvanik): map predicate to stable value.
	auto predicate = rewriter.getI32IntegerAttr(
	static_cast<int32_t>(dyn_cast<SrcOp>(op).getPredicate()));

	auto dstType = getMemRefType(op->getResult(0), rewriter);
	auto midOp = rewriter.create<DstOp>(op->getLoc(), dstType, predicate,
	lhValue, rhValue);

	auto result = wrapAsTensor(midOp.getResult(), op, rewriter);
	rewriter.replaceOp(
	op, {loadResultValue(op->getLoc(), op->getResult(0)->getType(), result,
	rewriter)});
	return this->matchSuccess();
	}
	};

	struct CmpIOpLowering
	: public CompareOpLowering<CmpIOp, IREEInterp::HL::CmpIOp> {
	using CompareOpLowering::CompareOpLowering;
	};

	struct CmpFOpLowering
	: public CompareOpLowering<CmpFOp, IREEInterp::HL::CmpFOp> {
	using CompareOpLowering::CompareOpLowering;
	};

	struct AllocOpLowering : public ConversionPattern {
	explicit AllocOpLowering(MLIRContext *context)
	: ConversionPattern(AllocOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	// TODO(benvanik): replace with length computation.
	rewriter.replaceOpWithNewOp<IREEInterp::HL::AllocHeapOp>(
	op, *op->getResultTypes().begin(), operands);
	return matchSuccess();
	}
	};

	struct DeallocOpLowering : public ConversionPattern {
	explicit DeallocOpLowering(MLIRContext *context)
	: ConversionPattern(DeallocOp::getOperationName(), 1, context) {}

	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	rewriter.replaceOpWithNewOp<IREEInterp::HL::DiscardOp>(op, operands[0]);
	return matchSuccess();
	}
	};

	struct ExtractElementOpLowering : public ConversionPattern {
	explicit ExtractElementOpLowering(MLIRContext *context)
	: ConversionPattern(ExtractElementOp::getOperationName(), 1, context) {}
	PatternMatchResult matchAndRewrite(
	Operation op, ArrayRef<Value > operands,
	ConversionPatternRewriter &rewriter) const override {
	auto extractOp = cast<ExtractElementOp>(op);
	Value *memRefInput = wrapAsMemRef(
	loadAccessValue(op->getLoc(), extractOp.getAggregate(), rewriter), op,
	rewriter);

	SmallVector<Value *, 4> indices = {extractOp.indices().begin(),
	extractOp.indices().end()};
	rewriter.replaceOpWithNewOp<LoadOp>(op, memRefInput, indices);
	return matchSuccess();
	}
	};

	#define UNARY_OP_LOWERING(StdOpType, IREEOpType) \
	struct StdOpType##Lowering : public UnaryOpLowering<StdOpType, IREEOpType> { \
	using UnaryOpLowering::UnaryOpLowering; \
	};

	#define BINARY_OP_LOWERING(StdOpType, IREEOpType) \
	struct StdOpType##Lowering \
	: public BinaryOpLowering<StdOpType, IREEOpType> { \
	using BinaryOpLowering::BinaryOpLowering; \
	};

	#define TERNARY_OP_LOWERING(StdOpType, IREEOpType) \
	struct StdOpType##Lowering \
	: public TernaryOpLowering<StdOpType, IREEOpType> { \
	using TernaryOpLowering::TernaryOpLowering; \
	};

	// UNARY_OP_LOWERING(RankOp, IREEInterp::HL::RankOp);
	UNARY_OP_LOWERING(DimOp, IREEInterp::HL::DimOp);
	// UNARY_OP_LOWERING(ShapeOp, IREEInterp::HL::ShapeOp);
	// UNARY_OP_LOWERING(LengthOp, IREEInterp::HL::LengthOp);

	// UNARY_OP_LOWERING(NotOp, IREEInterp::HL::NotOp);
	BINARY_OP_LOWERING(AndOp, IREEInterp::HL::AndOp);
	BINARY_OP_LOWERING(OrOp, IREEInterp::HL::OrOp);
	// BINARY_OP_LOWERING(XorOp, IREEInterp::HL::XorOp);
	// BINARY_OP_LOWERING(ShiftLeftOp, IREEInterp::HL::ShiftLeftOp);
	// BINARY_OP_LOWERING(ShiftRightLogicalOp, IREEInterp::HL::ShiftRightLogicalOp);
	// BINARY_OP_LOWERING(ShiftRightArithmeticOp,
	// IREEInterp::HL::ShiftRightArithmeticOp);

	BINARY_OP_LOWERING(AddIOp, IREEInterp::HL::AddIOp);
	BINARY_OP_LOWERING(AddFOp, IREEInterp::HL::AddFOp);
	BINARY_OP_LOWERING(SubIOp, IREEInterp::HL::SubIOp);
	BINARY_OP_LOWERING(SubFOp, IREEInterp::HL::SubFOp);
	// UNARY_OP_LOWERING(AbsIOp, IREEInterp::HL::AbsIOp);
	// UNARY_OP_LOWERING(AbsFOp, IREEInterp::HL::AbsFOp);
	BINARY_OP_LOWERING(MulIOp, IREEInterp::HL::MulIOp);
	BINARY_OP_LOWERING(MulFOp, IREEInterp::HL::MulFOp);
	BINARY_OP_LOWERING(DivISOp, IREEInterp::HL::DivISOp);
	BINARY_OP_LOWERING(DivIUOp, IREEInterp::HL::DivIUOp);
	BINARY_OP_LOWERING(DivFOp, IREEInterp::HL::DivFOp);
	// BINARY_OP_LOWERING(MulAddIOp, IREEInterp::HL::MulAddIOp);
	// BINARY_OP_LOWERING(MulAddFOp, IREEInterp::HL::MulAddFOp);
	// UNARY_OP_LOWERING(ExpFOp, IREEInterp::HL::ExpFOp);
	// UNARY_OP_LOWERING(LogFOp, IREEInterp::HL::LogFOp);
	// UNARY_OP_LOWERING(RsqrtFOp, IREEInterp::HL::RsqrtFOp);
	// UNARY_OP_LOWERING(CosFOp, IREEInterp::HL::CosFOp);
	// UNARY_OP_LOWERING(SinFOp, IREEInterp::HL::SinFOp);
	// UNARY_OP_LOWERING(TanhFOp, IREEInterp::HL::TanhFOp);
	// UNARY_OP_LOWERING(Atan2FOp, IREEInterp::HL::Atan2FOp);

	// BINARY_OP_LOWERING(MinISOp, IREEInterp::HL::MinISOp);
	// BINARY_OP_LOWERING(MinIUOp, IREEInterp::HL::MinIUOp);
	// BINARY_OP_LOWERING(MinFOp, IREEInterp::HL::MinFOp);
	// BINARY_OP_LOWERING(MaxISOp, IREEInterp::HL::MaxISOp);
	// BINARY_OP_LOWERING(MaxIUOp, IREEInterp::HL::MaxIUOp);
	// BINARY_OP_LOWERING(MaxFOp, IREEInterp::HL::MaxFOp);
	// TERNARY_OP_LOWERING(ClampFOp, IREEInterp::HL::ClampFOp);
	// UNARY_OP_LOWERING(FloorFOp, IREEInterp::HL::FloorFOp);
	// UNARY_OP_LOWERING(CeilFOp, IREEInterp::HL::CeilFOp);

	class LowerStdToInterpreterDialectPass
	: public FunctionPass<LowerStdToInterpreterDialectPass> {
	public:
	void runOnFunction() override {
	OwningRewritePatternList patterns;
	patterns.insert<
	ConstantOpLowering,
	// Control flow.
	CallOpLowering, CallIndirectOpLowering, ReturnOpLowering,
	BranchOpLowering, CondBranchOpLowering, CmpIOpLowering, CmpFOpLowering,
	// Memory management.
	AllocOpLowering, DeallocOpLowering, ExtractElementOpLowering,
	// Shape operations.
	DimOpLowering,
	// Logical ops.
	AndOpLowering, OrOpLowering,
	// Arithmetic ops.
	AddIOpLowering, AddFOpLowering, SubIOpLowering, SubFOpLowering,
	MulIOpLowering, MulFOpLowering, DivISOpLowering, DivIUOpLowering,
	DivFOpLowering>(&getContext());

	ConversionTarget target(getContext());
	target.addLegalDialect<IREEHLInterpreterDialect, IREELLInterpreterDialect,
	IREEDialect>();
	target.addLegalOp<LoadOp, StoreOp, FuncOp, ModuleOp>();
	target.addDynamicallyLegalOp<ConstantOp>([](ConstantOp constOp) {
	// std.constant is legal for index integers.
	return constOp.getValue().isa<IntegerAttr>() &&
	constOp.getType().isIndex();
	});
	// This is only able to be a full conversion because we know that we do it
	// last. Other conversions (e.g. XLA->IREE) can't be full because we might
	// still have e.g. std ops.
	// TODO(b/139012931) Move fully to the conversion framework.
	if (failed(applyFullConversion(getFunction(), target, patterns))) {
	return signalPassFailure();
	}
	}
	};

	std::unique_ptr<OpPassBase<FuncOp>> createLowerStdToInterpreterDialectPass() {
	return std::make_unique<LowerStdToInterpreterDialectPass>();
	}

	static PassRegistration<LowerStdToInterpreterDialectPass> pass(
	"lower-std-to-iree-interpreter",
	"Convert all std functions to the IREE dialect");

	} // namespace iree_compiler
	} // namespace mlir