1 //===- Bufferize.cpp - Bufferization of linalg ops ------------------===//
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/Transforms/Bufferize.h"
10 #include "PassDetail.h"
11 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
12 #include "mlir/Dialect/Linalg/Passes.h"
13 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
14 #include "mlir/Dialect/Linalg/Utils/Utils.h"
15 #include "mlir/Dialect/Vector/VectorOps.h"
16 #include "mlir/IR/Function.h"
17 #include "mlir/IR/Operation.h"
18 #include "mlir/Pass/Pass.h"
19 
20 using namespace ::mlir;
21 using namespace ::mlir::linalg;
22 
23 static SmallVector<Range, 4> computeLoopRanges(Location loc, LinalgOp linalgOp,
24                                                OpBuilder &b) {
25   auto indexingMaps = llvm::to_vector<4>(
26       linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
27   auto inputIndexingMaps =
28       llvm::makeArrayRef(indexingMaps).take_front(linalgOp.getNumInputs());
29 
30   mlir::edsc::ScopedContext scope(b, loc);
31   return emitLoopRanges(scope.getBuilderRef(), loc,
32                         concatAffineMaps(inputIndexingMaps),
33                         getShape(b, linalgOp));
34 }
35 
36 static Value maybeConvertToIndex(Location loc, Value val, OpBuilder &b) {
37   if (val.getType().isIndex())
38     return val;
39   return b.create<IndexCastOp>(loc, val, b.getIndexType());
40 }
41 
42 static LogicalResult
43 allocateBuffersForResults(Location loc, LinalgOp linalgOp,
44                           linalg::GenericOpAdaptor &adaptor,
45                           SmallVectorImpl<Value> &resultBuffers, OpBuilder &b) {
46   // Lazily compute loopRanges.
47   SmallVector<Range, 4> loopRanges;
48 
49   // Allocate a buffer for every tensor result.
50   for (auto en : llvm::enumerate(linalgOp.getOperation()->getResultTypes())) {
51     size_t resultIndex = en.index();
52     Type resultType = en.value();
53 
54     auto tensorType = resultType.dyn_cast<RankedTensorType>();
55     if (tensorType == nullptr) {
56       linalgOp.emitOpError()
57           << "tensor to buffer conversion expects ranked tensor results";
58       return failure();
59     }
60     auto tensorShape = tensorType.getShape();
61     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
62 
63     // Allocate buffers for init tensors that are assumed to fold onto the first
64     // results.
65     // TODO: update this assumption because the reality is more complex
66     // under linalg on tensor based transformations.
67     bool foldedInitTensor = resultIndex < linalgOp.getNumInitTensors();
68     if (foldedInitTensor) {
69       // Dealing with an init tensor requires distinguishing between 1-use
70       // and many-use cases which would create aliasing and WAR hazards.
71       Value initTensor = linalgOp.getInitTensor(resultIndex);
72       Value initBuffer = adaptor.init_tensors()[resultIndex];
73       if (initTensor.hasOneUse()) {
74         resultBuffers.push_back(initBuffer);
75         continue;
76       }
77       SmallVector<Value, 4> dynOperands;
78       for (auto dim : llvm::enumerate(tensorShape)) {
79         if (dim.value() == TensorType::kDynamicSize) {
80           dynOperands.push_back(b.create<DimOp>(loc, initTensor, dim.index()));
81         }
82       }
83       auto alloc = b.create<AllocOp>(loc, memrefType, dynOperands);
84       b.create<linalg::CopyOp>(loc, initBuffer, alloc);
85       resultBuffers.push_back(alloc);
86       continue;
87     }
88 
89     // Allocate buffers for statically-shaped results.
90     if (memrefType.hasStaticShape()) {
91       resultBuffers.push_back(b.create<AllocOp>(loc, memrefType));
92       continue;
93     }
94 
95     // Perform a naive shape inference for the dynamically-shaped results.
96     // Extract the required element out of the vector.
97     SmallVector<Value, 4> dynOperands;
98     auto resultIndexingMap = linalgOp.getOutputIndexingMap(resultIndex);
99     for (auto shapeElement : llvm::enumerate(tensorType.getShape())) {
100       if (loopRanges.empty())
101         loopRanges = computeLoopRanges(loc, linalgOp, b);
102 
103       if (shapeElement.value() != ShapedType::kDynamicSize)
104         continue;
105 
106       AffineExpr expr = resultIndexingMap.getResult(shapeElement.index());
107       switch (expr.getKind()) {
108       case AffineExprKind::DimId: {
109         int64_t loopIndex = expr.cast<AffineDimExpr>().getPosition();
110         Value size = maybeConvertToIndex(loc, loopRanges[loopIndex].size, b);
111         dynOperands.push_back(size);
112         break;
113       }
114       default:
115         return failure();
116       }
117     }
118     resultBuffers.push_back(b.create<AllocOp>(loc, memrefType, dynOperands));
119   }
120   return success();
121 }
122 
123 // Specialization for `linalg::GenericOp`.
124 /// A pattern to convert Generic Linalg operations which work on tensors to
125 /// use buffers. A buffer is allocated using BufferAssignmentPlacer for
126 /// each operation result. BufferPlacement pass should be later used to move
127 /// Alloc operations to the correct positions and insert the missing Dealloc
128 /// operations in the correct places.
129 static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
130                                      linalg::GenericOp genericOp,
131                                      ValueRange inputs, ValueRange outputs) {
132   // Generate a new linalg operation that works on buffers.
133   auto newGenericOp = rewriter.create<linalg::GenericOp>(
134       genericOp.getLoc(),
135       /*resultTensorTypes=*/llvm::None,
136       /*inputs=*/inputs,
137       /*outputBuffers=*/outputs,
138       /*initTensors=*/llvm::None, genericOp.indexing_maps(),
139       genericOp.iterator_types(), genericOp.docAttr(),
140       genericOp.library_callAttr(), genericOp.symbol_sourceAttr());
141 
142   // Create a new block in the region of the new Generic Op.
143   Block *oldBlock = genericOp.getBody();
144   Region &newRegion = newGenericOp.region();
145   Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
146                                          oldBlock->getArgumentTypes());
147 
148   // Add the result arguments to the new block.
149   for (Value v : ValueRange(outputs).drop_front(genericOp.getNumInitTensors()))
150     newBlock->addArgument(v.getType().cast<MemRefType>().getElementType());
151 
152   // Clone the body of the old block to the new block.
153   BlockAndValueMapping mapping;
154   mapping.map(oldBlock->getArguments(), newBlock->getArguments());
155 
156   OpBuilder::InsertionGuard guard(rewriter);
157   rewriter.setInsertionPointToEnd(newBlock);
158   for (auto &op : oldBlock->getOperations()) {
159     Operation *clonedOp = rewriter.clone(op, mapping);
160     mapping.map(op.getResults(), clonedOp->getResults());
161   }
162 
163   // Replace the results of the old op with the new output buffers.
164   rewriter.replaceOp(genericOp, outputs);
165 }
166 
167 // TODO: Specialization for `linalg::IndexedGenericOp`.
168 
169 // Specialization for all other `linalg::LinalgOp`.
170 static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
171                                      linalg::LinalgOp linalgOp,
172                                      ValueRange inputs, ValueRange outputs) {
173   assert(!isa<linalg::GenericOp>(linalgOp.getOperation()));
174   assert(!isa<linalg::IndexedGenericOp>(linalgOp.getOperation()));
175   SmallVector<Value, 8> newOperands = inputs;
176   newOperands.append(outputs.begin(), outputs.end());
177   auto otherOperands = linalgOp.getAssumedNonShapedOperands();
178   newOperands.append(otherOperands.begin(), otherOperands.end());
179   LinalgOp res = cast<LinalgOp>(linalgOp.clone(rewriter, linalgOp.getLoc(),
180                                                /*resultTypes=*/ArrayRef<Type>{},
181                                                newOperands));
182   // Need to mutate the operands_segment_sizes in the resulting op.
183   res.setNumOutputBuffers(outputs.size());
184   res.setNumInitTensors(0);
185   // Replace the results of the old op with the new output buffers.
186   rewriter.replaceOp(linalgOp, outputs);
187 }
188 
189 LogicalResult mlir::linalg::LinalgOpConverter::matchAndRewrite(
190     Operation *op, ArrayRef<Value> operands,
191     ConversionPatternRewriter &rewriter) const {
192   LinalgOp linalgOp = dyn_cast<linalg::LinalgOp>(op);
193   if (!linalgOp)
194     return failure();
195 
196   // We abuse the GenericOpAdaptor here.
197   // TODO: Manually create an Adaptor that captures inputs, output_buffers and
198   // init_tensors for all linalg::LinalgOp interface ops.
199   linalg::GenericOpAdaptor adaptor(operands, op->getAttrDictionary());
200 
201   // All inputs need to be turned into buffers first. Until then, bail out.
202   if (llvm::any_of(adaptor.inputs(),
203                    [](Value in) { return !in.getType().isa<MemRefType>(); }))
204     return failure();
205 
206   // All init_tensors need to be turned into buffers first. Until then, bail
207   // out.
208   if (llvm::any_of(adaptor.init_tensors(),
209                    [](Value in) { return !in.getType().isa<MemRefType>(); }))
210     return failure();
211 
212   Location loc = linalgOp.getLoc();
213   SmallVector<Value, 2> newOutputBuffers(adaptor.output_buffers().begin(),
214                                          adaptor.output_buffers().end());
215 
216   if (failed(allocateBuffersForResults(loc, linalgOp, adaptor, newOutputBuffers,
217                                        rewriter))) {
218     linalgOp.emitOpError() << "Failed to allocate buffers for tensor results.";
219     return failure();
220   }
221 
222   // Delegate to the linalg generic pattern.
223   if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
224     finalizeBufferAllocation(rewriter, genericOp, adaptor.inputs(),
225                              newOutputBuffers);
226     return success();
227   }
228 
229   finalizeBufferAllocation(rewriter, linalgOp, adaptor.inputs(),
230                            newOutputBuffers);
231   return success();
232 }
233 
234 LogicalResult mlir::linalg::TensorConstantOpConverter::matchAndRewrite(
235     ConstantOp op, ArrayRef<Value> operands,
236     ConversionPatternRewriter &rewriter) const {
237   RankedTensorType rankedTensorType = op.getType().dyn_cast<RankedTensorType>();
238   if (!rankedTensorType)
239     return failure();
240   if (llvm::any_of(rankedTensorType.getShape(), [](int64_t s) {
241         return s == 0 || ShapedType::isDynamic(s);
242       }))
243     return failure();
244 
245   int64_t nElements = 1;
246   for (int64_t s : rankedTensorType.getShape())
247     nElements *= s;
248   Type elementType = rankedTensorType.getElementType();
249   MemRefType memrefType =
250       converter.convertType(op.getType()).cast<MemRefType>();
251   VectorType flatVectorType = VectorType::get({nElements}, elementType);
252   MemRefType memrefOfFlatVectorType = MemRefType::get({}, flatVectorType);
253   MemRefType flatMemrefType = MemRefType::get({nElements}, elementType);
254 
255   Location loc = op.getLoc();
256   auto attr = op.getValue().cast<DenseElementsAttr>();
257   Value alloc =
258       rewriter.create<AllocOp>(loc, memrefOfFlatVectorType, ValueRange{});
259   Value cstVec = rewriter.create<ConstantOp>(loc, flatVectorType,
260                                              attr.reshape(flatVectorType));
261   rewriter.create<StoreOp>(loc, cstVec, alloc);
262 
263   Value memref =
264       rewriter.create<vector::TypeCastOp>(loc, flatMemrefType, alloc);
265   if (rankedTensorType.getRank() > 1) {
266     // Introduce a linalg.reshape to flatten the memref.
267     AffineMap collapseAllDims = AffineMap::getMultiDimIdentityMap(
268         /*numDims=*/rankedTensorType.getRank(), op.getContext());
269     memref = rewriter.create<linalg::ReshapeOp>(
270         loc, memrefType, memref,
271         rewriter.getAffineMapArrayAttr(collapseAllDims));
272   }
273   rewriter.replaceOp(op, memref);
274 
275   return success();
276 }
277 
278 LogicalResult mlir::linalg::TensorCastOpConverter::matchAndRewrite(
279     TensorCastOp op, ArrayRef<Value> operands,
280     ConversionPatternRewriter &rewriter) const {
281   if (op.getType().hasRank())
282     return failure();
283   Type t = UnrankedMemRefType::get(op.getType().getElementType(),
284                                    /*memorySpace=*/0);
285   rewriter.replaceOpWithNewOp<MemRefCastOp>(op, t, operands.front());
286   return success();
287 }
288 
289 namespace {
290 
291 /// Converts Linalg operations that work on tensor-type operands or results to
292 /// work on buffers.
293 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
294   void runOnOperation() override {
295     MLIRContext &context = getContext();
296     ConversionTarget target(context);
297     BufferAssignmentTypeConverter converter;
298 
299     // Mark all Standard operations legal.
300     target.addLegalDialect<StandardOpsDialect, vector::VectorDialect>();
301     target.addLegalOp<ModuleOp>();
302     target.addLegalOp<ModuleTerminatorOp>();
303 
304     // Mark all Linalg operations illegal as long as they work on tensors.
305     auto isLegalOperation = [&](Operation *op) {
306       return converter.isLegal(op);
307     };
308     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
309         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
310             isLegalOperation));
311 
312     // Mark operations that consume or return tensors illegal.
313     auto isLegal = [&](Operation *op) {
314       if (llvm::any_of(op->getOperandTypes(),
315                        [&](Type t) { return !converter.isLegal(t); }))
316         return false;
317       if (llvm::any_of(op->getResultTypes(),
318                        [&](Type t) { return !converter.isLegal(t); }))
319         return false;
320       return true;
321     };
322     target.addDynamicallyLegalOp<
323         // clang-format off
324         CallOp,
325         ConstantOp,
326         ConstantIntOp,
327         ConstantIndexOp,
328         ConstantFloatOp,
329         ReturnOp,
330         TensorCastOp
331         // clang-format on
332         >(isLegal);
333 
334     // Mark the function operation illegal as long as an argument is tensor.
335     // TODO: if the FuncOp is a FuncOp that only has a declaration (e.g. to an
336     // externally defined symbol like an external library calls), only convert
337     // if some special attribute is set. This will allow more control of interop
338     // across ABI boundaries.
339     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
340       return converter.isSignatureLegal(funcOp.getType()) &&
341              llvm::none_of(funcOp.getType().getResults(),
342                            [&](Type type) { return type.isa<MemRefType>(); }) &&
343              converter.isLegal(&funcOp.getBody());
344     });
345 
346     converter.setResultConversionKind<RankedTensorType, MemRefType>(
347         BufferAssignmentTypeConverter::AppendToArgumentsList);
348 
349     OwningRewritePatternList patterns;
350     populateLinalgBufferizePatterns(&context, converter, patterns);
351     populateWithBufferAssignmentOpConversionPatterns<
352         mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(&context, converter,
353                                                         patterns);
354     if (failed(applyFullConversion(this->getOperation(), target, patterns)))
355       this->signalPassFailure();
356   }
357 };
358 } // end anonymous namespace
359 
360 std::unique_ptr<OperationPass<ModuleOp>> mlir::createLinalgBufferizePass() {
361   return std::make_unique<LinalgBufferizePass>();
362 }
363 void mlir::linalg::populateLinalgBufferizePatterns(
364     MLIRContext *context, BufferAssignmentTypeConverter &converter,
365     OwningRewritePatternList &patterns) {
366   patterns.insert<
367       // clang-format off
368       LinalgOpConverter,
369       TensorCastOpConverter,
370       TensorConstantOpConverter
371       // clang-format on
372       >(context, converter);
373 }
374