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