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/Linalg/Transforms/Transforms.h"
15 #include "mlir/Dialect/MemRef/IR/MemRef.h"
16 #include "mlir/Dialect/SCF/SCF.h"
17 #include "mlir/Dialect/StandardOps/IR/Ops.h"
18 
19 using namespace mlir;
20 using namespace mlir::linalg;
21 
22 /// Helper function to extract the operand types that are passed to the
23 /// generated CallOp. MemRefTypes have their layout canonicalized since the
24 /// information is not used in signature generation.
25 /// Note that static size information is not modified.
26 static SmallVector<Type, 4> extractOperandTypes(Operation *op) {
27   SmallVector<Type, 4> result;
28   result.reserve(op->getNumOperands());
29   if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(op)) {
30     auto *ctx = op->getContext();
31     auto numLoops = indexedGenericOp.getNumLoops();
32     result.reserve(op->getNumOperands() + numLoops);
33     result.assign(numLoops, IndexType::get(ctx));
34   }
35   for (auto type : op->getOperandTypes()) {
36     // The underlying descriptor type (e.g. LLVM) does not have layout
37     // information. Canonicalizing the type at the level of std when going into
38     // a library call avoids needing to introduce DialectCastOp.
39     if (auto memrefType = type.dyn_cast<MemRefType>())
40       result.push_back(eraseStridedLayout(memrefType));
41     else
42       result.push_back(type);
43   }
44   return result;
45 }
46 
47 // Get a SymbolRefAttr containing the library function name for the LinalgOp.
48 // If the library function does not exist, insert a declaration.
49 static FlatSymbolRefAttr getLibraryCallSymbolRef(Operation *op,
50                                                  PatternRewriter &rewriter) {
51   auto linalgOp = cast<LinalgOp>(op);
52   auto fnName = linalgOp.getLibraryCallName();
53   if (fnName.empty()) {
54     op->emitWarning("No library call defined for: ") << *op;
55     return {};
56   }
57 
58   // fnName is a dynamic std::string, unique it via a SymbolRefAttr.
59   FlatSymbolRefAttr fnNameAttr = rewriter.getSymbolRefAttr(fnName);
60   auto module = op->getParentOfType<ModuleOp>();
61   if (module.lookupSymbol(fnName)) {
62     return fnNameAttr;
63   }
64 
65   SmallVector<Type, 4> inputTypes(extractOperandTypes(op));
66   assert(op->getNumResults() == 0 &&
67          "Library call for linalg operation can be generated only for ops that "
68          "have void return types");
69   auto libFnType = rewriter.getFunctionType(inputTypes, {});
70 
71   OpBuilder::InsertionGuard guard(rewriter);
72   // Insert before module terminator.
73   rewriter.setInsertionPoint(module.getBody(),
74                              std::prev(module.getBody()->end()));
75   FuncOp funcOp =
76       rewriter.create<FuncOp>(op->getLoc(), fnNameAttr.getValue(), libFnType);
77   // Insert a function attribute that will trigger the emission of the
78   // corresponding `_mlir_ciface_xxx` interface so that external libraries see
79   // a normalized ABI. This interface is added during std to llvm conversion.
80   funcOp->setAttr("llvm.emit_c_interface", UnitAttr::get(op->getContext()));
81   funcOp.setPrivate();
82   return fnNameAttr;
83 }
84 
85 static SmallVector<Value, 4>
86 createTypeCanonicalizedMemRefOperands(OpBuilder &b, Location loc,
87                                       ValueRange operands) {
88   SmallVector<Value, 4> res;
89   res.reserve(operands.size());
90   for (auto op : operands) {
91     auto memrefType = op.getType().dyn_cast<MemRefType>();
92     if (!memrefType) {
93       res.push_back(op);
94       continue;
95     }
96     Value cast =
97         b.create<memref::CastOp>(loc, eraseStridedLayout(memrefType), op);
98     res.push_back(cast);
99   }
100   return res;
101 }
102 
103 LogicalResult mlir::linalg::LinalgOpToLibraryCallRewrite::matchAndRewrite(
104     LinalgOp op, PatternRewriter &rewriter) const {
105   // Only LinalgOp for which there is no specialized pattern go through this.
106   if (isa<CopyOp>(op) || isa<IndexedGenericOp>(op))
107     return failure();
108 
109   // TODO: remove once index ops are supported.
110   if (op.hasIndexSemantics())
111     return failure();
112 
113   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
114   if (!libraryCallName)
115     return failure();
116 
117   rewriter.replaceOpWithNewOp<mlir::CallOp>(
118       op, libraryCallName.getValue(), TypeRange(),
119       createTypeCanonicalizedMemRefOperands(rewriter, op->getLoc(),
120                                             op->getOperands()));
121   return success();
122 }
123 
124 LogicalResult mlir::linalg::CopyOpToLibraryCallRewrite::matchAndRewrite(
125     CopyOp op, PatternRewriter &rewriter) const {
126   auto inputPerm = op.inputPermutation();
127   if (inputPerm.hasValue() && !inputPerm->isIdentity())
128     return failure();
129   auto outputPerm = op.outputPermutation();
130   if (outputPerm.hasValue() && !outputPerm->isIdentity())
131     return failure();
132 
133   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
134   if (!libraryCallName)
135     return failure();
136 
137   rewriter.replaceOpWithNewOp<mlir::CallOp>(
138       op, libraryCallName.getValue(), TypeRange(),
139       createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(),
140                                             op.getOperands()));
141   return success();
142 }
143 
144 LogicalResult mlir::linalg::CopyTransposeRewrite::matchAndRewrite(
145     CopyOp op, PatternRewriter &rewriter) const {
146   Value in = op.input(), out = op.output();
147 
148   // If either inputPerm or outputPerm are non-identities, insert transposes.
149   auto inputPerm = op.inputPermutation();
150   if (inputPerm.hasValue() && !inputPerm->isIdentity())
151     in = rewriter.create<memref::TransposeOp>(op.getLoc(), in,
152                                               AffineMapAttr::get(*inputPerm));
153   auto outputPerm = op.outputPermutation();
154   if (outputPerm.hasValue() && !outputPerm->isIdentity())
155     out = rewriter.create<memref::TransposeOp>(op.getLoc(), out,
156                                                AffineMapAttr::get(*outputPerm));
157 
158   // If nothing was transposed, fail and let the conversion kick in.
159   if (in == op.input() && out == op.output())
160     return failure();
161 
162   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
163   if (!libraryCallName)
164     return failure();
165 
166   rewriter.replaceOpWithNewOp<mlir::CallOp>(
167       op, libraryCallName.getValue(), TypeRange(),
168       createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(), {in, out}));
169   return success();
170 }
171 
172 LogicalResult
173 mlir::linalg::IndexedGenericOpToLibraryCallRewrite::matchAndRewrite(
174     IndexedGenericOp op, PatternRewriter &rewriter) const {
175   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
176   if (!libraryCallName)
177     return failure();
178 
179   // TODO: Use induction variables values instead of zeros, when
180   // IndexedGenericOp is tiled.
181   auto zero = rewriter.create<mlir::ConstantOp>(
182       op.getLoc(), rewriter.getIntegerAttr(rewriter.getIndexType(), 0));
183   auto indexedGenericOp = cast<IndexedGenericOp>(op);
184   auto numLoops = indexedGenericOp.getNumLoops();
185   SmallVector<Value, 4> operands;
186   operands.reserve(numLoops + op.getNumOperands());
187   for (unsigned i = 0; i < numLoops; ++i)
188     operands.push_back(zero);
189   for (auto operand : op.getOperands())
190     operands.push_back(operand);
191   rewriter.replaceOpWithNewOp<mlir::CallOp>(
192       op, libraryCallName.getValue(), TypeRange(),
193       createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(), operands));
194   return success();
195 }
196 
197 /// Populate the given list with patterns that convert from Linalg to Standard.
198 void mlir::linalg::populateLinalgToStandardConversionPatterns(
199     RewritePatternSet &patterns) {
200   // TODO: ConvOp conversion needs to export a descriptor with relevant
201   // attribute values such as kernel striding and dilation.
202   // clang-format off
203   patterns.add<
204       CopyOpToLibraryCallRewrite,
205       CopyTransposeRewrite,
206       IndexedGenericOpToLibraryCallRewrite,
207       LinalgOpToLibraryCallRewrite>(patterns.getContext());
208   // clang-format on
209 }
210 
211 namespace {
212 struct ConvertLinalgToStandardPass
213     : public ConvertLinalgToStandardBase<ConvertLinalgToStandardPass> {
214   void runOnOperation() override;
215 };
216 } // namespace
217 
218 void ConvertLinalgToStandardPass::runOnOperation() {
219   auto module = getOperation();
220   ConversionTarget target(getContext());
221   target.addLegalDialect<AffineDialect, memref::MemRefDialect, scf::SCFDialect,
222                          StandardOpsDialect>();
223   target.addLegalOp<ModuleOp, FuncOp, ReturnOp>();
224   target.addLegalOp<linalg::ReshapeOp, linalg::RangeOp>();
225   RewritePatternSet patterns(&getContext());
226   populateLinalgToStandardConversionPatterns(patterns);
227   if (failed(applyFullConversion(module, target, std::move(patterns))))
228     signalPassFailure();
229 }
230 
231 std::unique_ptr<OperationPass<ModuleOp>>
232 mlir::createConvertLinalgToStandardPass() {
233   return std::make_unique<ConvertLinalgToStandardPass>();
234 }
235