| // 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/PatternMatch.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/Sequencer/LLDialect.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/Sequencer/LLOps.h" |
| #include "third_party/mlir_edge/iree/compiler/IR/StructureOps.h" |
| |
| namespace mlir { |
| namespace iree_compiler { |
| |
| namespace { |
| |
| // Returns an integer scalar memref containing the offset specified by |indices| |
| // within |type|. |
| Value *computeOffset(Location loc, Value *reference, Value *indices, |
| OpBuilder &builder) { |
| auto referenceType = reference->getType().cast<ShapedType>(); |
| auto *shapeMemRef = builder |
| .create<IREESeq::LL::AllocHeapOp>( |
| loc, |
| builder.getMemRefType({referenceType.getRank()}, |
| builder.getIntegerType(32)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| builder.create<IREESeq::LL::ShapeOp>(loc, reference, shapeMemRef); |
| auto *resultMemRef = |
| builder |
| .create<IREESeq::LL::AllocHeapOp>( |
| loc, builder.getMemRefType({}, builder.getIntegerType(32)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| auto elementSizeAttr = builder.getIntegerAttr( |
| builder.getIntegerType(8), referenceType.getElementTypeBitWidth() / 8); |
| builder.create<IREESeq::LL::ComputeOffsetOp>( |
| loc, shapeMemRef, elementSizeAttr, indices, resultMemRef); |
| return resultMemRef; |
| } |
| |
| // Returns a tuple of (offset, length) integer scalar memrefs with the range |
| // specified by |indices| and |lengths| within |type|. |
| std::pair<Value *, Value *> computeRange(Location loc, Value *reference, |
| Value *indices, Value *lengths, |
| OpBuilder &builder) { |
| auto referenceType = reference->getType().cast<ShapedType>(); |
| auto *shapeMemRef = builder |
| .create<IREESeq::LL::AllocHeapOp>( |
| loc, |
| builder.getMemRefType({referenceType.getRank()}, |
| builder.getIntegerType(32)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| builder.create<IREESeq::LL::ShapeOp>(loc, reference, shapeMemRef); |
| auto *offsetMemRef = |
| builder |
| .create<IREESeq::LL::AllocHeapOp>( |
| loc, builder.getMemRefType({}, builder.getIntegerType(32)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| auto *lengthMemRef = |
| builder |
| .create<IREESeq::LL::AllocHeapOp>( |
| loc, builder.getMemRefType({}, builder.getIntegerType(32)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| auto elementSizeAttr = builder.getIntegerAttr( |
| builder.getIntegerType(8), referenceType.getElementTypeBitWidth() / 8); |
| builder.create<IREESeq::LL::ComputeRangeOp>(loc, shapeMemRef, elementSizeAttr, |
| indices, lengths, offsetMemRef, |
| lengthMemRef); |
| return {offsetMemRef, lengthMemRef}; |
| } |
| |
| struct LowerSliceOpPattern : public OpRewritePattern<IREESeq::HL::SliceOp> { |
| using OpRewritePattern::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::SliceOp op, |
| PatternRewriter &rewriter) const { |
| auto range = computeRange(op.getLoc(), op.src(), op.indices(), op.lengths(), |
| rewriter); |
| rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicSliceOp>( |
| op, ArrayRef<Type>{op.getResult()->getType()}, |
| ArrayRef<Value *>{op.src(), range.first, range.second}, op.getAttrs()); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct LowerShapeOpPattern : public OpRewritePattern<IREESeq::HL::ShapeOp> { |
| using OpRewritePattern::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::ShapeOp op, |
| PatternRewriter &rewriter) const { |
| auto *shapeMemRef = |
| rewriter |
| .create<IREESeq::LL::AllocHeapOp>( |
| op.getLoc(), |
| rewriter.getMemRefType( |
| {op.getResult()->getType().cast<ShapedType>().getRank()}, |
| rewriter.getIntegerType(64)), |
| ArrayRef<Value *>{}) |
| .getResult(); |
| op.replaceAllUsesWith(shapeMemRef); |
| rewriter.replaceOpWithNewOp<IREESeq::LL::ShapeOp>(op, op.getOperand(), |
| shapeMemRef); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct LowerCopyOpPattern : public OpRewritePattern<IREESeq::HL::CopyOp> { |
| using OpRewritePattern::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::CopyOp op, |
| PatternRewriter &rewriter) const { |
| auto *srcOffsetMemRef = |
| computeOffset(op.getLoc(), op.src(), op.srcIndices(), rewriter); |
| auto dstRange = computeRange(op.getLoc(), op.dst(), op.dstIndices(), |
| op.lengths(), rewriter); |
| rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicCopyOp>( |
| op, op.src(), srcOffsetMemRef, op.dst(), dstRange.first, |
| dstRange.second); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct LowerFillOpPattern : public OpRewritePattern<IREESeq::HL::FillOp> { |
| using OpRewritePattern::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::FillOp op, |
| PatternRewriter &rewriter) const { |
| auto dstRange = computeRange(op.getLoc(), op.dst(), op.dstIndices(), |
| op.lengths(), rewriter); |
| rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicFillOp>( |
| op, op.value(), op.dst(), dstRange.first, dstRange.second); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct LowerBranchOpPattern : public OpRewritePattern<IREESeq::HL::BranchOp> { |
| using OpRewritePattern<IREESeq::HL::BranchOp>::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::BranchOp op, |
| PatternRewriter &rewriter) const { |
| SmallVector<Value *, 8> operands{op.getOperation()->getOperands()}; |
| |
| rewriter.replaceOpWithNewOp<IREESeq::LL::BranchOp>(op, op.getDest(), |
| operands); |
| return matchSuccess(); |
| } |
| }; |
| |
| struct LowerCondCondBranchOpPattern |
| : public OpRewritePattern<IREESeq::HL::CondBranchOp> { |
| using OpRewritePattern<IREESeq::HL::CondBranchOp>::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(IREESeq::HL::CondBranchOp op, |
| PatternRewriter &rewriter) const { |
| SmallVector<Value *, 8> trueOperands{op.getTrueOperands()}; |
| SmallVector<Value *, 8> falseOperands{op.getFalseOperands()}; |
| |
| rewriter.replaceOpWithNewOp<IREESeq::LL::CondBranchOp>( |
| op, op.getCondition(), op.getTrueDest(), trueOperands, |
| op.getFalseDest(), falseOperands); |
| return matchSuccess(); |
| } |
| }; |
| |
| // Rewrites an op into one with all the same operands, results, and attributes. |
| // Operands and results in the ops must have the same order and attributes must |
| // have the same name. They must also be constructed properly by the default |
| // builders. |
| template <typename SRC, typename DST> |
| struct LowerIdenticalOpPattern : public OpRewritePattern<SRC> { |
| using OpRewritePattern<SRC>::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(SRC op, PatternRewriter &rewriter) const { |
| SmallVector<Type, 8> resultTypes{op.getOperation()->getResultTypes()}; |
| SmallVector<Value *, 8> operands{op.getOperation()->getOperands()}; |
| |
| rewriter.replaceOpWithNewOp<DST>(op, resultTypes, operands, op.getAttrs()); |
| return this->matchSuccess(); |
| } |
| }; |
| |
| } // namespace |
| |
| class LowerSequencerDialectPass |
| : public FunctionPass<LowerSequencerDialectPass> { |
| public: |
| void runOnFunction() override { |
| OwningRewritePatternList patterns; |
| patterns.insert< |
| LowerIdenticalOpPattern<IREE::ConstantOp, IREESeq::LL::ConstantOp>, |
| LowerIdenticalOpPattern<IREESeq::HL::DispatchOp, |
| IREESeq::LL::DynamicDispatchOp>, |
| LowerShapeOpPattern, LowerCopyOpPattern, LowerSliceOpPattern, |
| LowerBranchOpPattern, LowerCondCondBranchOpPattern>(&getContext()); |
| #define IDENTICAL_OP_LOWERING(op_name) \ |
| LowerIdenticalOpPattern<IREESeq::HL::op_name, IREESeq::LL::op_name> |
| patterns.insert< |
| IDENTICAL_OP_LOWERING(AllocHeapOp), IDENTICAL_OP_LOWERING(CloneOp), |
| IDENTICAL_OP_LOWERING(ReshapeOp), IDENTICAL_OP_LOWERING(CallOp), |
| IDENTICAL_OP_LOWERING(ReturnOp)>(&getContext()); |
| #undef IDENTICAL_OP_LOWERING |
| |
| ConversionTarget target(getContext()); |
| target.addLegalDialect<IREELLSequencerDialect>(); |
| target.addLegalOp<FuncOp, ModuleOp>(); |
| |
| if (failed(applyFullConversion(getFunction(), target, patterns))) { |
| return signalPassFailure(); |
| } |
| } |
| }; |
| |
| std::unique_ptr<OpPassBase<FuncOp>> createLowerSequencerDialectPass() { |
| return std::make_unique<LowerSequencerDialectPass>(); |
| } |
| |
| static PassRegistration<LowerSequencerDialectPass> pass( |
| "iree-lower-sequencer-dialect", |
| "Lowers the IREE HL sequencer dialect to the LL sequencer dialect."); |
| |
| } // namespace iree_compiler |
| } // namespace mlir |
