1 //===- LinalgToStandard.cpp - conversion from Linalg 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/LinalgToStandard/LinalgToStandard.h"
10 
11 #include "../PassDetail.h"
12 #include "mlir/Dialect/Affine/IR/AffineOps.h"
13 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
14 #include "mlir/Dialect/SCF/SCF.h"
15 #include "mlir/Dialect/StandardOps/IR/Ops.h"
16 
17 using namespace mlir;
18 using namespace mlir::linalg;
19 
20 /// Helper function to extract the operand types that are passed to the
21 /// generated CallOp. MemRefTypes have their layout canonicalized since the
22 /// information is not used in signature generation.
23 /// Note that static size information is not modified.
24 template <typename LinalgOp>
25 static SmallVector<Type, 4> extractOperandTypes(Operation *op) {
26   SmallVector<Type, 4> result;
27   result.reserve(op->getNumOperands());
28   for (auto type : op->getOperandTypes()) {
29     // The underlying descriptor type (e.g. LLVM) does not have layout
30     // information. Canonicalizing the type at the level of std when going into
31     // a library call avoids needing to introduce DialectCastOp.
32     if (auto memrefType = type.dyn_cast<MemRefType>())
33       result.push_back(eraseStridedLayout(memrefType));
34     else
35       result.push_back(type);
36   }
37   return result;
38 }
39 
40 template <>
41 SmallVector<Type, 4> extractOperandTypes<IndexedGenericOp>(Operation *op) {
42   auto *ctx = op->getContext();
43   auto indexedGenericOp = cast<IndexedGenericOp>(op);
44   auto numLoops = indexedGenericOp.getNumLoops();
45 
46   SmallVector<Type, 4> result(numLoops, IndexType::get(ctx));
47   auto canonicalizedOperands = extractOperandTypes<LinalgOp>(op);
48   result.append(canonicalizedOperands.begin(), canonicalizedOperands.end());
49   return result;
50 }
51 
52 // Get a SymbolRefAttr containing the library function name for the LinalgOp.
53 // If the library function does not exist, insert a declaration.
54 template <typename LinalgOp>
55 static FlatSymbolRefAttr getLibraryCallSymbolRef(Operation *op,
56                                                  PatternRewriter &rewriter) {
57   auto linalgOp = cast<LinalgOp>(op);
58   auto fnName = linalgOp.getLibraryCallName();
59   if (fnName.empty()) {
60     op->emitWarning("No library call defined for: ") << *op;
61     return {};
62   }
63 
64   // fnName is a dynamic std::string, unique it via a SymbolRefAttr.
65   FlatSymbolRefAttr fnNameAttr = rewriter.getSymbolRefAttr(fnName);
66   auto module = op->getParentOfType<ModuleOp>();
67   if (module.lookupSymbol(fnName)) {
68     return fnNameAttr;
69   }
70 
71   SmallVector<Type, 4> inputTypes(extractOperandTypes<LinalgOp>(op));
72   assert(op->getNumResults() == 0 &&
73          "Library call for linalg operation can be generated only for ops that "
74          "have void return types");
75   auto libFnType = FunctionType::get(inputTypes, {}, rewriter.getContext());
76 
77   OpBuilder::InsertionGuard guard(rewriter);
78   // Insert before module terminator.
79   rewriter.setInsertionPoint(module.getBody(),
80                              std::prev(module.getBody()->end()));
81   FuncOp funcOp =
82       rewriter.create<FuncOp>(op->getLoc(), fnNameAttr.getValue(), libFnType,
83                               ArrayRef<NamedAttribute>{});
84   // Insert a function attribute that will trigger the emission of the
85   // corresponding `_mlir_ciface_xxx` interface so that external libraries see
86   // a normalized ABI. This interface is added during std to llvm conversion.
87   funcOp.setAttr("llvm.emit_c_interface", UnitAttr::get(op->getContext()));
88   return fnNameAttr;
89 }
90 
91 namespace {
92 
93 SmallVector<Value, 4>
94 createTypeCanonicalizedMemRefOperands(OpBuilder &b, Location loc,
95                                       ValueRange operands) {
96   SmallVector<Value, 4> res;
97   res.reserve(operands.size());
98   for (auto op : operands) {
99     auto memrefType = op.getType().dyn_cast<MemRefType>();
100     if (!memrefType) {
101       res.push_back(op);
102       continue;
103     }
104     Value cast =
105         b.create<MemRefCastOp>(loc, eraseStridedLayout(memrefType), op);
106     res.push_back(cast);
107   }
108   return res;
109 }
110 
111 // LinalgOpConversion<LinalgOp> creates a new call to the type-canonicalized
112 // `LinalgOp::getLibraryCallName()` function.
113 // The implementation of the function can be either in the same module or in an
114 // externally linked library.
115 template <typename LinalgOp>
116 class LinalgOpConversion : public OpRewritePattern<LinalgOp> {
117 public:
118   using OpRewritePattern<LinalgOp>::OpRewritePattern;
119 
120   LogicalResult matchAndRewrite(LinalgOp op,
121                                 PatternRewriter &rewriter) const override {
122     auto libraryCallName = getLibraryCallSymbolRef<LinalgOp>(op, rewriter);
123     if (!libraryCallName)
124       return failure();
125 
126     rewriter.replaceOpWithNewOp<mlir::CallOp>(
127         op, libraryCallName.getValue(), ArrayRef<Type>{},
128         createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(),
129                                               op.getOperands()));
130     return success();
131   }
132 };
133 
134 /// Conversion pattern specialization for CopyOp. This kicks in when both input
135 /// and output permutations are left unspecified or are the identity.
136 template <>
137 class LinalgOpConversion<CopyOp> : public OpRewritePattern<CopyOp> {
138 public:
139   using OpRewritePattern<CopyOp>::OpRewritePattern;
140 
141   LogicalResult matchAndRewrite(CopyOp op,
142                                 PatternRewriter &rewriter) const override {
143     auto inputPerm = op.inputPermutation();
144     if (inputPerm.hasValue() && !inputPerm->isIdentity())
145       return failure();
146     auto outputPerm = op.outputPermutation();
147     if (outputPerm.hasValue() && !outputPerm->isIdentity())
148       return failure();
149 
150     auto libraryCallName = getLibraryCallSymbolRef<CopyOp>(op, rewriter);
151     if (!libraryCallName)
152       return failure();
153 
154     rewriter.replaceOpWithNewOp<mlir::CallOp>(
155         op, libraryCallName.getValue(), ArrayRef<Type>{},
156         createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(),
157                                               op.getOperands()));
158     return success();
159   }
160 };
161 
162 /// Conversion pattern specialization for IndexedGenericOp.
163 template <>
164 class LinalgOpConversion<IndexedGenericOp>
165     : public OpRewritePattern<IndexedGenericOp> {
166 public:
167   using OpRewritePattern<IndexedGenericOp>::OpRewritePattern;
168 
169   LogicalResult matchAndRewrite(IndexedGenericOp op,
170                                 PatternRewriter &rewriter) const override {
171     auto libraryCallName =
172         getLibraryCallSymbolRef<IndexedGenericOp>(op, rewriter);
173     if (!libraryCallName)
174       return failure();
175 
176     // TODO(pifon, ntv): Use induction variables values instead of zeros, when
177     // IndexedGenericOp is tiled.
178     auto zero = rewriter.create<mlir::ConstantOp>(
179         op.getLoc(), rewriter.getIntegerAttr(rewriter.getIndexType(), 0));
180     auto indexedGenericOp = cast<IndexedGenericOp>(op);
181     auto numLoops = indexedGenericOp.getNumLoops();
182     SmallVector<Value, 4> operands;
183     operands.reserve(numLoops + op.getNumOperands());
184     for (unsigned i = 0; i < numLoops; ++i)
185       operands.push_back(zero);
186     for (auto operand : op.getOperands())
187       operands.push_back(operand);
188     rewriter.replaceOpWithNewOp<mlir::CallOp>(
189         op, libraryCallName.getValue(), ArrayRef<Type>{},
190         createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(), operands));
191     return success();
192   }
193 };
194 
195 /// A non-conversion rewrite pattern kicks in to convert CopyOp with
196 /// permutations into a sequence of TransposeOp and permutation-free CopyOp.
197 /// This interplays together with TransposeOpConversion and
198 /// LinalgConversion<CopyOp> to create a path to the LLVM dialect.
199 class CopyTransposeConversion : public OpRewritePattern<CopyOp> {
200 public:
201   using OpRewritePattern<CopyOp>::OpRewritePattern;
202 
203   LogicalResult matchAndRewrite(CopyOp op,
204                                 PatternRewriter &rewriter) const override {
205     Value in = op.input(), out = op.output();
206 
207     // If either inputPerm or outputPerm are non-identities, insert transposes.
208     auto inputPerm = op.inputPermutation();
209     if (inputPerm.hasValue() && !inputPerm->isIdentity())
210       in = rewriter.create<linalg::TransposeOp>(op.getLoc(), in,
211                                                 AffineMapAttr::get(*inputPerm));
212     auto outputPerm = op.outputPermutation();
213     if (outputPerm.hasValue() && !outputPerm->isIdentity())
214       out = rewriter.create<linalg::TransposeOp>(
215           op.getLoc(), out, AffineMapAttr::get(*outputPerm));
216 
217     // If nothing was transposed, fail and let the conversion kick in.
218     if (in == op.input() && out == op.output())
219       return failure();
220 
221     rewriter.replaceOpWithNewOp<CopyOp>(op, in, out);
222     return success();
223   }
224 };
225 } // namespace
226 
227 /// Populate the given list with patterns that convert from Linalg to Standard.
228 void mlir::populateLinalgToStandardConversionPatterns(
229     OwningRewritePatternList &patterns, MLIRContext *ctx) {
230   // TODO(ntv) ConvOp conversion needs to export a descriptor with relevant
231   // attribute values such as kernel striding and dilation.
232   // clang-format off
233   patterns.insert<
234       CopyTransposeConversion,
235       LinalgOpConversion<ConvOp>,
236       LinalgOpConversion<PoolingMaxOp>,
237       LinalgOpConversion<PoolingMinOp>,
238       LinalgOpConversion<PoolingSumOp>,
239       LinalgOpConversion<CopyOp>,
240       LinalgOpConversion<DotOp>,
241       LinalgOpConversion<FillOp>,
242       LinalgOpConversion<GenericOp>,
243       LinalgOpConversion<IndexedGenericOp>,
244       LinalgOpConversion<MatmulOp>,
245       LinalgOpConversion<MatvecOp>>(ctx);
246   // clang-format on
247 }
248 
249 namespace {
250 struct ConvertLinalgToStandardPass
251     : public ConvertLinalgToStandardBase<ConvertLinalgToStandardPass> {
252   void runOnOperation() override;
253 };
254 } // namespace
255 
256 void ConvertLinalgToStandardPass::runOnOperation() {
257   auto module = getOperation();
258   ConversionTarget target(getContext());
259   target.addLegalDialect<AffineDialect, scf::SCFDialect, StandardOpsDialect>();
260   target.addLegalOp<ModuleOp, FuncOp, ModuleTerminatorOp, ReturnOp>();
261   target.addLegalOp<linalg::TransposeOp, linalg::ReshapeOp, linalg::RangeOp>();
262   OwningRewritePatternList patterns;
263   populateLinalgToStandardConversionPatterns(patterns, &getContext());
264   if (failed(applyFullConversion(module, target, patterns)))
265     signalPassFailure();
266 }
267 
268 std::unique_ptr<OperationPass<ModuleOp>>
269 mlir::createConvertLinalgToStandardPass() {
270   return std::make_unique<ConvertLinalgToStandardPass>();
271 }
272