1 //===- ShapeToStandard.cpp - conversion from Shape to Standard dialect ----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
10 
11 #include "../PassDetail.h"
12 #include "mlir/Dialect/SCF/SCF.h"
13 #include "mlir/Dialect/Shape/IR/Shape.h"
14 #include "mlir/Dialect/StandardOps/IR/Ops.h"
15 #include "mlir/Transforms/DialectConversion.h"
16 
17 using namespace mlir;
18 using namespace mlir::shape;
19 
20 namespace {
21 
22 /// Generated conversion patterns.
23 #include "ShapeToStandardPatterns.inc"
24 
25 /// Conversion patterns.
26 template <typename SrcOpTy, typename DstOpTy>
27 class BinaryOpConversion : public OpConversionPattern<SrcOpTy> {
28 public:
29   using OpConversionPattern<SrcOpTy>::OpConversionPattern;
30 
31   LogicalResult
32   matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
33                   ConversionPatternRewriter &rewriter) const override {
34     typename SrcOpTy::Adaptor adaptor(operands);
35     rewriter.replaceOpWithNewOp<DstOpTy>(op.getOperation(), adaptor.lhs(),
36                                          adaptor.rhs());
37     return success();
38   }
39 };
40 
41 class ShapeOfOpConversion : public OpConversionPattern<ShapeOfOp> {
42 public:
43   using OpConversionPattern<ShapeOfOp>::OpConversionPattern;
44 
45   LogicalResult
46   matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
47                   ConversionPatternRewriter &rewriter) const override {
48     ShapeOfOp::Adaptor transformed(operands);
49     auto loc = op.getLoc();
50     auto tensorVal = transformed.arg();
51     auto tensorTy = tensorVal.getType();
52 
53     // For unranked tensors `shape_of` lowers to `scf` and the pattern can be
54     // found in the corresponding pass.
55     if (tensorTy.isa<UnrankedTensorType>())
56       return failure();
57 
58     // Build values for individual dimensions.
59     SmallVector<Value, 8> dimValues;
60     auto rankedTensorTy = tensorTy.cast<RankedTensorType>();
61     int64_t rank = rankedTensorTy.getRank();
62     for (int64_t i = 0; i < rank; i++) {
63       if (rankedTensorTy.isDynamicDim(i)) {
64         auto dimVal = rewriter.create<DimOp>(loc, tensorVal, i);
65         dimValues.push_back(dimVal);
66       } else {
67         int64_t dim = rankedTensorTy.getDimSize(i);
68         auto dimVal = rewriter.create<ConstantIndexOp>(loc, dim);
69         dimValues.push_back(dimVal);
70       }
71     }
72 
73     // Materialize shape as ranked tensor.
74     rewriter.replaceOpWithNewOp<TensorFromElementsOp>(op.getOperation(),
75                                                       dimValues);
76     return success();
77   }
78 };
79 
80 class ConstSizeOpConverter : public OpConversionPattern<ConstSizeOp> {
81 public:
82   using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
83 
84   LogicalResult
85   matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
86                   ConversionPatternRewriter &rewriter) const override {
87     rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(),
88                                                  op.value().getSExtValue());
89     return success();
90   }
91 };
92 
93 /// Type conversions.
94 class ShapeTypeConverter : public TypeConverter {
95 public:
96   using TypeConverter::convertType;
97 
98   ShapeTypeConverter(MLIRContext *ctx) {
99     // Add default pass-through conversion.
100     addConversion([&](Type type) { return type; });
101 
102     addConversion([ctx](SizeType type) { return IndexType::get(ctx); });
103     addConversion([ctx](ShapeType type) {
104       return RankedTensorType::get({ShapedType::kDynamicSize},
105                                    IndexType::get(ctx));
106     });
107   }
108 };
109 
110 /// Conversion pass.
111 class ConvertShapeToStandardPass
112     : public ConvertShapeToStandardBase<ConvertShapeToStandardPass> {
113 
114   void runOnOperation() override {
115     // Setup type conversion.
116     MLIRContext &ctx = getContext();
117     ShapeTypeConverter typeConverter(&ctx);
118 
119     // Setup target legality.
120     ConversionTarget target(ctx);
121     target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
122     target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();
123     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
124       return typeConverter.isSignatureLegal(op.getType()) &&
125              typeConverter.isLegal(&op.getBody());
126     });
127 
128     // Setup conversion patterns.
129     OwningRewritePatternList patterns;
130     populateShapeToStandardConversionPatterns(patterns, &ctx);
131     populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
132 
133     // Apply conversion.
134     auto module = getOperation();
135     if (failed(applyFullConversion(module, target, patterns)))
136       signalPassFailure();
137   }
138 };
139 
140 } // namespace
141 
142 void mlir::populateShapeToStandardConversionPatterns(
143     OwningRewritePatternList &patterns, MLIRContext *ctx) {
144   populateWithGenerated(ctx, &patterns);
145   // clang-format off
146   patterns.insert<
147       BinaryOpConversion<AddOp, AddIOp>,
148       BinaryOpConversion<MulOp, MulIOp>,
149       ConstSizeOpConverter,
150       ShapeOfOpConversion>(ctx);
151   // clang-format on
152 }
153 
154 std::unique_ptr<OperationPass<ModuleOp>>
155 mlir::createConvertShapeToStandardPass() {
156   return std::make_unique<ConvertShapeToStandardPass>();
157 }
158