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 class RankOpConverter : public OpConversionPattern<shape::RankOp> {
94 public:
95   using OpConversionPattern<shape::RankOp>::OpConversionPattern;
96 
97   LogicalResult
98   matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
99                   ConversionPatternRewriter &rewriter) const override {
100     shape::RankOp::Adaptor transformed(operands);
101     rewriter.replaceOpWithNewOp<DimOp>(op.getOperation(), transformed.shape(),
102                                        0);
103     return success();
104   }
105 };
106 
107 /// Type conversions.
108 class ShapeTypeConverter : public TypeConverter {
109 public:
110   using TypeConverter::convertType;
111 
112   ShapeTypeConverter(MLIRContext *ctx) {
113     // Add default pass-through conversion.
114     addConversion([&](Type type) { return type; });
115 
116     addConversion([ctx](SizeType type) { return IndexType::get(ctx); });
117     addConversion([ctx](ShapeType type) {
118       return RankedTensorType::get({ShapedType::kDynamicSize},
119                                    IndexType::get(ctx));
120     });
121   }
122 };
123 
124 /// Conversion pass.
125 class ConvertShapeToStandardPass
126     : public ConvertShapeToStandardBase<ConvertShapeToStandardPass> {
127 
128   void runOnOperation() override {
129     // Setup type conversion.
130     MLIRContext &ctx = getContext();
131     ShapeTypeConverter typeConverter(&ctx);
132 
133     // Setup target legality.
134     ConversionTarget target(ctx);
135     target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
136     target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();
137     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
138       return typeConverter.isSignatureLegal(op.getType()) &&
139              typeConverter.isLegal(&op.getBody());
140     });
141 
142     // Setup conversion patterns.
143     OwningRewritePatternList patterns;
144     populateShapeToStandardConversionPatterns(patterns, &ctx);
145     populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
146 
147     // Apply conversion.
148     auto module = getOperation();
149     if (failed(applyFullConversion(module, target, patterns)))
150       signalPassFailure();
151   }
152 };
153 
154 } // namespace
155 
156 void mlir::populateShapeToStandardConversionPatterns(
157     OwningRewritePatternList &patterns, MLIRContext *ctx) {
158   populateWithGenerated(ctx, &patterns);
159   // clang-format off
160   patterns.insert<
161       BinaryOpConversion<AddOp, AddIOp>,
162       BinaryOpConversion<MulOp, MulIOp>,
163       ConstSizeOpConverter,
164       RankOpConverter,
165       ShapeOfOpConversion>(ctx);
166   // clang-format on
167 }
168 
169 std::unique_ptr<OperationPass<ModuleOp>>
170 mlir::createConvertShapeToStandardPass() {
171   return std::make_unique<ConvertShapeToStandardPass>();
172 }
173