| // 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/ArrayRef.h" |
| #include "third_party/llvm/llvm/include/llvm/ADT/DenseMap.h" |
| #include "third_party/llvm/llvm/include/llvm/ADT/SmallVector.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/Builders.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/OperationSupport.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/Support/LogicalResult.h" |
| #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/DialectConversion.h" |
| #include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/Utils.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/Dialect.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/Ops.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/Sequencer/HLDialect.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/Sequencer/HLOps.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/StructureOps.h" |
| #include "third_party/mlir_edge/iree/compiler/Utils/MemRefUtils.h" |
| |
| namespace mlir { |
| namespace iree_compiler { |
| |
| namespace { |
| |
| class SequencerConversionPattern : public ConversionPattern { |
| public: |
| SequencerConversionPattern(StringRef operationName, int benefit, |
| MLIRContext *context, |
| MemRefTypeConverter &typeConverter) |
| : ConversionPattern(operationName, benefit, context), |
| typeConverter_(typeConverter) {} |
| |
| protected: |
| MemRefTypeConverter &typeConverter_; |
| }; |
| |
| struct ConstantOpLowering : public SequencerConversionPattern { |
| ConstantOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(ConstantOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| const auto &valueAttr = cast<ConstantOp>(op).getValue(); |
| auto midOp = rewriter.create<IREE::ConstantOp>(op->getLoc(), valueAttr); |
| |
| auto result = wrapAsTensor(midOp.getResult(), op, rewriter); |
| rewriter.replaceOp( |
| op, {loadResultValue(op->getLoc(), op->getResult(0)->getType(), result, |
| rewriter)}); |
| return matchSuccess(); |
| } |
| }; |
| |
| class CallOpLowering : public SequencerConversionPattern { |
| public: |
| CallOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(CallOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| auto callOp = cast<CallOp>(op); |
| |
| SmallVector<Type, 4> convertedResults; |
| auto result = typeConverter_.convertTypes( |
| callOp.getCalleeType().getResults(), convertedResults); |
| (void)result; |
| assert(succeeded(result) && "expected valid callee type conversion"); |
| rewriter.replaceOpWithNewOp<IREESeq::HL::CallOp>( |
| op, callOp.getCallee(), convertedResults, operands); |
| |
| return matchSuccess(); |
| } |
| }; |
| |
| class CallIndirectOpLowering : public SequencerConversionPattern { |
| public: |
| CallIndirectOpLowering(MLIRContext *context, |
| MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(CallIndirectOp::getOperationName(), 1, |
| context, typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| auto callOp = cast<CallIndirectOp>(op); |
| rewriter.replaceOpWithNewOp<IREESeq::HL::CallIndirectOp>( |
| op, callOp.getCallee(), operands); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct ReturnOpLowering : public SequencerConversionPattern { |
| ReturnOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(ReturnOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| SmallVector<Value *, 4> newOperands; |
| newOperands.reserve(operands.size()); |
| for (auto *operand : operands) { |
| newOperands.push_back(wrapAsMemRef(operand, op, rewriter)); |
| } |
| rewriter.replaceOpWithNewOp<IREESeq::HL::ReturnOp>(op, newOperands); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct BranchOpLowering : public SequencerConversionPattern { |
| BranchOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(BranchOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> properOperands, |
| ArrayRef<Block *> destinations, ArrayRef<ArrayRef<Value *>> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| rewriter.replaceOpWithNewOp<IREESeq::HL::BranchOp>(op, destinations[0], |
| operands[0]); |
| return this->matchSuccess(); |
| } |
| }; |
| |
| struct CondBranchOpLowering : public SequencerConversionPattern { |
| CondBranchOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(CondBranchOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| 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<IREESeq::HL::CondBranchOp>( |
| op, condValue, destinations[IREESeq::HL::CondBranchOp::trueIndex], |
| operands[IREESeq::HL::CondBranchOp::trueIndex], |
| destinations[IREESeq::HL::CondBranchOp::falseIndex], |
| operands[IREESeq::HL::CondBranchOp::falseIndex]); |
| return this->matchSuccess(); |
| } |
| }; |
| |
| class AllocOpLowering : public SequencerConversionPattern { |
| public: |
| AllocOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(AllocOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| // TODO(benvanik): replace with length computation. |
| rewriter.replaceOpWithNewOp<IREESeq::HL::AllocHeapOp>( |
| op, *op->getResultTypes().begin(), operands); |
| return matchSuccess(); |
| } |
| }; |
| |
| class DeallocOpLowering : public SequencerConversionPattern { |
| public: |
| DeallocOpLowering(MLIRContext *context, MemRefTypeConverter &typeConverter) |
| : SequencerConversionPattern(DeallocOp::getOperationName(), 1, context, |
| typeConverter) {} |
| |
| PatternMatchResult matchAndRewrite( |
| Operation *op, ArrayRef<Value *> operands, |
| ConversionPatternRewriter &rewriter) const override { |
| rewriter.replaceOpWithNewOp<IREESeq::HL::DiscardOp>(op, operands[0]); |
| return matchSuccess(); |
| } |
| }; |
| |
| void populateStdToSequencerConversionPatterns( |
| MLIRContext *context, MemRefTypeConverter &converter, |
| OwningRewritePatternList &patterns) { |
| patterns.insert<ConstantOpLowering, |
| // Control flow. |
| CallOpLowering, CallIndirectOpLowering, ReturnOpLowering, |
| BranchOpLowering, CondBranchOpLowering, |
| // Memory management. |
| AllocOpLowering, DeallocOpLowering>(context, converter); |
| } |
| |
| } // namespace |
| |
| // Lowers functions using std.* ops to the IREE HL sequencer dialect and buffer |
| // view types. |
| // FuncOp signatures will be updated to use the buffer view type and |
| // dispatch regions will get iree.bind_input where needed. |
| // |
| // Beyond bindings there will be no other changes within dispatchable regions. |
| // It is up to the downstream dialects to properly use the bindings to map their |
| // I/O to expected values. |
| // |
| // Note that output buffer allocation is required following this pass to either |
| // elide dispatch results entirely and provide output params or provide both |
| // while ensuring that the returned value is always sliced from an input. This |
| // should happen prior to outlining. |
| class LowerStdToSequencerDialectPass |
| : public ModulePass<LowerStdToSequencerDialectPass> { |
| public: |
| void runOnModule() override { |
| auto module = getModule(); |
| |
| // Only convert top-level functions, not ones nested in executables. |
| std::vector<Operation *> toConvert; |
| for (auto funcOp : module.getOps<FuncOp>()) { |
| toConvert.push_back(funcOp); |
| } |
| |
| // Convert the signature and body of all sequencer functions. |
| MemRefTypeConverter converter(&getContext()); |
| ConversionTarget target(getContext()); |
| target.addLegalDialect<IREEHLSequencerDialect, IREEDialect>(); |
| target.addLegalOp<LoadOp, StoreOp>(); |
| target.addDynamicallyLegalOp<FuncOp>( |
| [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); |
| |
| OwningRewritePatternList patterns; |
| populateStdToSequencerConversionPatterns(&getContext(), converter, |
| patterns); |
| if (failed( |
| applyPartialConversion(toConvert, target, patterns, &converter))) { |
| return signalPassFailure(); |
| } |
| } |
| }; |
| |
| std::unique_ptr<OpPassBase<ModuleOp>> createLowerStdToSequencerDialectPass() { |
| return std::make_unique<LowerStdToSequencerDialectPass>(); |
| } |
| |
| static PassRegistration<LowerStdToSequencerDialectPass> pass( |
| "iree-lower-std-to-sequencer-dialect", |
| "Lowers std ops to the IREE HL sequencer dialect."); |
| |
| } // namespace iree_compiler |
| } // namespace mlir |
