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/StandardOps/Transforms/Passes.h"
16 #include "mlir/Dialect/Vector/VectorOps.h"
17 #include "mlir/IR/Function.h"
18 #include "mlir/IR/Operation.h"
19 #include "mlir/Pass/Pass.h"
20 
21 using namespace ::mlir;
22 using namespace ::mlir::linalg;
23 
24 static SmallVector<Range, 4> computeLoopRanges(Location loc, LinalgOp linalgOp,
25                                                OpBuilder &b) {
26   auto indexingMaps = llvm::to_vector<4>(
27       linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
28   auto inputIndexingMaps =
29       llvm::makeArrayRef(indexingMaps).take_front(linalgOp.getNumInputs());
30 
31   mlir::edsc::ScopedContext scope(b, loc);
32   return emitLoopRanges(scope.getBuilderRef(), loc,
33                         concatAffineMaps(inputIndexingMaps),
34                         getShape(b, linalgOp));
35 }
36 
37 static Value maybeConvertToIndex(Location loc, Value val, OpBuilder &b) {
38   if (val.getType().isIndex())
39     return val;
40   return b.create<IndexCastOp>(loc, val, b.getIndexType());
41 }
42 
43 static LogicalResult
44 allocateBuffersForResults(Location loc, LinalgOp linalgOp,
45                           linalg::GenericOpAdaptor &adaptor,
46                           SmallVectorImpl<Value> &resultBuffers, OpBuilder &b) {
47   // Lazily compute loopRanges.
48   SmallVector<Range, 4> loopRanges;
49 
50   // Allocate a buffer for every tensor result.
51   for (auto en : llvm::enumerate(linalgOp.getOperation()->getResultTypes())) {
52     size_t resultIndex = en.index();
53     Type resultType = en.value();
54 
55     auto tensorType = resultType.dyn_cast<RankedTensorType>();
56     if (tensorType == nullptr) {
57       linalgOp.emitOpError()
58           << "tensor to buffer conversion expects ranked tensor results";
59       return failure();
60     }
61     auto tensorShape = tensorType.getShape();
62     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
63 
64     // Allocate buffers for init tensors that are assumed to fold onto the first
65     // results.
66     // TODO: update this assumption because the reality is more complex
67     // under linalg on tensor based transformations.
68     bool foldedInitTensor = resultIndex < linalgOp.getNumInitTensors();
69     if (foldedInitTensor) {
70       Value initTensor = linalgOp.getInitTensor(resultIndex);
71       Value initBuffer = adaptor.init_tensors()[resultIndex];
72       SmallVector<Value, 4> dynOperands;
73       for (auto dim : llvm::enumerate(tensorShape)) {
74         if (dim.value() == TensorType::kDynamicSize) {
75           dynOperands.push_back(b.create<DimOp>(loc, initTensor, dim.index()));
76         }
77       }
78       auto alloc = b.create<AllocOp>(loc, memrefType, dynOperands);
79       b.create<linalg::CopyOp>(loc, initBuffer, alloc);
80       resultBuffers.push_back(alloc);
81       continue;
82     }
83 
84     // Allocate buffers for statically-shaped results.
85     if (memrefType.hasStaticShape()) {
86       resultBuffers.push_back(b.create<AllocOp>(loc, memrefType));
87       continue;
88     }
89 
90     // Perform a naive shape inference for the dynamically-shaped results.
91     // Extract the required element out of the vector.
92     SmallVector<Value, 4> dynOperands;
93     auto resultIndexingMap = linalgOp.getOutputIndexingMap(resultIndex);
94     for (auto shapeElement : llvm::enumerate(tensorType.getShape())) {
95       if (loopRanges.empty())
96         loopRanges = computeLoopRanges(loc, linalgOp, b);
97 
98       if (shapeElement.value() != ShapedType::kDynamicSize)
99         continue;
100 
101       AffineExpr expr = resultIndexingMap.getResult(shapeElement.index());
102       switch (expr.getKind()) {
103       case AffineExprKind::DimId: {
104         int64_t loopIndex = expr.cast<AffineDimExpr>().getPosition();
105         Value size = maybeConvertToIndex(loc, loopRanges[loopIndex].size, b);
106         dynOperands.push_back(size);
107         break;
108       }
109       default:
110         return failure();
111       }
112     }
113     resultBuffers.push_back(b.create<AllocOp>(loc, memrefType, dynOperands));
114   }
115   return success();
116 }
117 
118 // Specialization for `linalg::GenericOp`.
119 /// A pattern to convert Generic Linalg operations which work on tensors to
120 /// use buffers. BufferPlacement pass should be later used to move
121 /// Alloc operations to the correct positions and insert the missing Dealloc
122 /// operations in the correct places.
123 static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
124                                      linalg::GenericOp genericOp,
125                                      ValueRange inputs, ValueRange outputs) {
126   // Generate a new linalg operation that works on buffers.
127   auto newGenericOp = rewriter.create<linalg::GenericOp>(
128       genericOp.getLoc(),
129       /*resultTensorTypes=*/llvm::None,
130       /*inputs=*/inputs,
131       /*outputBuffers=*/outputs,
132       /*initTensors=*/llvm::None, genericOp.indexing_maps(),
133       genericOp.iterator_types(), genericOp.docAttr(),
134       genericOp.library_callAttr(), genericOp.sparseAttr(),
135       genericOp.symbol_sourceAttr());
136 
137   // Create a new block in the region of the new Generic Op.
138   Block *oldBlock = genericOp.getBody();
139   Region &newRegion = newGenericOp.region();
140   Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
141                                          oldBlock->getArgumentTypes());
142 
143   // Add the result arguments to the new block.
144   for (Value v : ValueRange(outputs).drop_front(genericOp.getNumInitTensors()))
145     newBlock->addArgument(v.getType().cast<MemRefType>().getElementType());
146 
147   // Clone the body of the old block to the new block.
148   BlockAndValueMapping mapping;
149   mapping.map(oldBlock->getArguments(), newBlock->getArguments());
150 
151   OpBuilder::InsertionGuard guard(rewriter);
152   rewriter.setInsertionPointToEnd(newBlock);
153   for (auto &op : oldBlock->getOperations()) {
154     Operation *clonedOp = rewriter.clone(op, mapping);
155     mapping.map(op.getResults(), clonedOp->getResults());
156   }
157 
158   // Replace the results of the old op with the new output buffers.
159   rewriter.replaceOp(genericOp, outputs);
160 }
161 
162 // TODO: Specialization for `linalg::IndexedGenericOp`.
163 
164 // Specialization for all other `linalg::LinalgOp`.
165 static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
166                                      linalg::LinalgOp linalgOp,
167                                      ValueRange inputs, ValueRange outputs) {
168   assert(!isa<linalg::GenericOp>(linalgOp.getOperation()));
169   assert(!isa<linalg::IndexedGenericOp>(linalgOp.getOperation()));
170   SmallVector<Value, 8> newOperands = inputs;
171   newOperands.append(outputs.begin(), outputs.end());
172   auto otherOperands = linalgOp.getAssumedNonShapedOperands();
173   newOperands.append(otherOperands.begin(), otherOperands.end());
174   LinalgOp res = cast<LinalgOp>(linalgOp.clone(rewriter, linalgOp.getLoc(),
175                                                /*resultTypes=*/ArrayRef<Type>{},
176                                                newOperands));
177   // Need to mutate the operands_segment_sizes in the resulting op.
178   res.setNumOutputBuffers(outputs.size());
179   res.setNumInitTensors(0);
180   // Replace the results of the old op with the new output buffers.
181   rewriter.replaceOp(linalgOp, outputs);
182 }
183 
184 //===----------------------------------------------------------------------===//
185 // Bufferization patterns.
186 //===----------------------------------------------------------------------===//
187 
188 namespace {
189 /// Generic conversion pattern that matches any LinalgOp. This avoids template
190 /// instantiating one pattern for each LinalgOp.
191 class BufferizeAnyLinalgOp : public ConversionPattern {
192 public:
193   BufferizeAnyLinalgOp(TypeConverter &typeConverter)
194       : ConversionPattern(/*benefit=*/1, typeConverter, MatchAnyOpTypeTag()) {}
195 
196   LogicalResult
197   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
198                   ConversionPatternRewriter &rewriter) const final {
199 
200     LinalgOp linalgOp = dyn_cast<linalg::LinalgOp>(op);
201     if (!linalgOp)
202       return failure();
203 
204     // We abuse the GenericOpAdaptor here.
205     // TODO: Manually create an Adaptor that captures inputs, output_buffers and
206     // init_tensors for all linalg::LinalgOp interface ops.
207     linalg::GenericOpAdaptor adaptor(operands, op->getAttrDictionary());
208 
209     Location loc = linalgOp.getLoc();
210     SmallVector<Value, 2> newOutputBuffers(adaptor.output_buffers().begin(),
211                                            adaptor.output_buffers().end());
212 
213     if (failed(allocateBuffersForResults(loc, linalgOp, adaptor,
214                                          newOutputBuffers, rewriter))) {
215       linalgOp.emitOpError()
216           << "Failed to allocate buffers for tensor results.";
217       return failure();
218     }
219 
220     // Delegate to the linalg generic pattern.
221     if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
222       finalizeBufferAllocation(rewriter, genericOp, adaptor.inputs(),
223                                newOutputBuffers);
224       return success();
225     }
226 
227     finalizeBufferAllocation(rewriter, linalgOp, adaptor.inputs(),
228                              newOutputBuffers);
229     return success();
230   }
231 };
232 
233 // Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
234 static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
235   return llvm::to_vector<4>(
236       llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
237         return a.cast<IntegerAttr>().getInt();
238       }));
239 }
240 
241 /// Convert `subtensor %t [offsets][sizes][strides] -> %st` to an alloc + copy
242 /// pattern.
243 /// ```
244 ///   %a = alloc(sizes)
245 ///   %sv = subview %source [offsets][sizes][strides]
246 ///   linalg_copy(%sv, %a)
247 /// ```
248 ///
249 /// This pattern is arguable a std pattern once linalg::CopyOp becomes
250 /// std::CopyOp.
251 class SubTensorOpConverter : public OpConversionPattern<SubTensorOp> {
252 public:
253   using OpConversionPattern<SubTensorOp>::OpConversionPattern;
254 
255   LogicalResult
256   matchAndRewrite(SubTensorOp op, ArrayRef<Value> operands,
257                   ConversionPatternRewriter &rewriter) const final {
258     SubTensorOpAdaptor adaptor(operands,
259                                op.getOperation()->getAttrDictionary());
260     Value sourceMemref = adaptor.source();
261     assert(sourceMemref.getType().isa<MemRefType>());
262 
263     MemRefType subviewMemRefType =
264         getTypeConverter()->convertType(op.getType()).cast<MemRefType>();
265     // op.sizes() capture exactly the dynamic alloc operands matching the
266     // subviewMemRefType thanks to subview/subtensor canonicalization and
267     // verification.
268     Value alloc =
269         rewriter.create<AllocOp>(op.getLoc(), subviewMemRefType, op.sizes());
270     Value subView = rewriter.create<SubViewOp>(
271         op.getLoc(), sourceMemref, extractFromI64ArrayAttr(op.static_offsets()),
272         extractFromI64ArrayAttr(op.static_sizes()),
273         extractFromI64ArrayAttr(op.static_strides()), op.offsets(), op.sizes(),
274         op.strides());
275     rewriter.create<linalg::CopyOp>(op.getLoc(), subView, alloc);
276     rewriter.replaceOp(op, alloc);
277     return success();
278   }
279 };
280 
281 /// Convert `subtensor_insert %source into %dest [offsets][sizes][strides] ->
282 /// %t` to an tensor_to_memref + subview + copy + tensor_load pattern.
283 /// tensor_to_memref and tensor_load are inserted automatically by the
284 /// conversion infra:
285 /// ```
286 ///   %sv = subview %dest [offsets][sizes][strides]
287 ///   linalg_copy(%source, %sv)
288 ///   // replace with %dest
289 /// ```
290 ///
291 /// This pattern is arguable a std pattern once linalg::CopyOp becomes
292 /// std::CopyOp.
293 class SubTensorInsertOpConverter
294     : public OpConversionPattern<SubTensorInsertOp> {
295 public:
296   using OpConversionPattern<SubTensorInsertOp>::OpConversionPattern;
297 
298   LogicalResult
299   matchAndRewrite(SubTensorInsertOp op, ArrayRef<Value> operands,
300                   ConversionPatternRewriter &rewriter) const final {
301     SubTensorInsertOpAdaptor adaptor(operands,
302                                      op.getOperation()->getAttrDictionary());
303     Value sourceMemRef = adaptor.source();
304     assert(sourceMemRef.getType().isa<MemRefType>());
305 
306     Value destMemRef = adaptor.dest();
307     assert(destMemRef.getType().isa<MemRefType>());
308 
309     // Take a subview to copy the small memref.
310     Value subview = rewriter.create<SubViewOp>(
311         op.getLoc(), destMemRef, extractFromI64ArrayAttr(op.static_offsets()),
312         extractFromI64ArrayAttr(op.static_sizes()),
313         extractFromI64ArrayAttr(op.static_strides()), adaptor.offsets(),
314         adaptor.sizes(), adaptor.strides());
315     // Copy the small memref.
316     rewriter.create<linalg::CopyOp>(op.getLoc(), sourceMemRef, subview);
317     rewriter.replaceOp(op, destMemRef);
318     return success();
319   }
320 };
321 
322 /// TensorConstantOp conversion inserts a linearized 1-D vector constant that is
323 /// stored in memory. A linalg.reshape is introduced to convert to the desired
324 /// n-D buffer form.
325 class TensorConstantOpConverter : public OpConversionPattern<ConstantOp> {
326 public:
327   using OpConversionPattern::OpConversionPattern;
328 
329   LogicalResult
330   matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
331                   ConversionPatternRewriter &rewriter) const final {
332 
333     RankedTensorType rankedTensorType =
334         op.getType().dyn_cast<RankedTensorType>();
335     if (!rankedTensorType)
336       return failure();
337     if (llvm::any_of(rankedTensorType.getShape(), [](int64_t s) {
338           return s == 0 || ShapedType::isDynamic(s);
339         }))
340       return failure();
341 
342     int64_t nElements = 1;
343     for (int64_t s : rankedTensorType.getShape())
344       nElements *= s;
345     Type elementType = rankedTensorType.getElementType();
346     MemRefType memrefType =
347         getTypeConverter()->convertType(op.getType()).cast<MemRefType>();
348     VectorType flatVectorType = VectorType::get({nElements}, elementType);
349     MemRefType memrefOfFlatVectorType = MemRefType::get({}, flatVectorType);
350     MemRefType flatMemrefType = MemRefType::get({nElements}, elementType);
351 
352     Location loc = op.getLoc();
353     auto attr = op.getValue().cast<DenseElementsAttr>();
354     Value alloc =
355         rewriter.create<AllocOp>(loc, memrefOfFlatVectorType, ValueRange{});
356     Value cstVec = rewriter.create<ConstantOp>(loc, flatVectorType,
357                                                attr.reshape(flatVectorType));
358     rewriter.create<StoreOp>(loc, cstVec, alloc);
359 
360     Value memref =
361         rewriter.create<vector::TypeCastOp>(loc, flatMemrefType, alloc);
362     if (rankedTensorType.getRank() > 1) {
363       // Introduce a linalg.reshape to flatten the memref.
364       AffineMap collapseAllDims = AffineMap::getMultiDimIdentityMap(
365           /*numDims=*/rankedTensorType.getRank(), op.getContext());
366       memref = rewriter.create<linalg::ReshapeOp>(
367           loc, memrefType, memref,
368           rewriter.getAffineMapArrayAttr(collapseAllDims));
369     }
370     rewriter.replaceOp(op, memref);
371 
372     return success();
373   }
374 };
375 } // namespace
376 
377 namespace {
378 /// Converts Linalg operations that work on tensor-type operands or results to
379 /// work on buffers.
380 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
381   void runOnOperation() override {
382     MLIRContext &context = getContext();
383     ConversionTarget target(context);
384     BufferizeTypeConverter typeConverter;
385 
386     // Mark all Standard operations legal.
387     target.addLegalDialect<StandardOpsDialect, vector::VectorDialect>();
388     target.addIllegalOp<SubTensorOp, SubTensorInsertOp>();
389 
390     // Mark all Linalg operations illegal as long as they work on tensors.
391     auto isLegalOperation = [&](Operation *op) {
392       return typeConverter.isLegal(op);
393     };
394     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
395     target.addDynamicallyLegalOp<ConstantOp>(isLegalOperation);
396 
397     OwningRewritePatternList patterns;
398     populateLinalgBufferizePatterns(&context, typeConverter, patterns);
399     if (failed(applyPartialConversion(getOperation(), target,
400                                       std::move(patterns))))
401       signalPassFailure();
402   }
403 };
404 } // end anonymous namespace
405 
406 std::unique_ptr<OperationPass<ModuleOp>> mlir::createLinalgBufferizePass() {
407   return std::make_unique<LinalgBufferizePass>();
408 }
409 
410 void mlir::linalg::populateLinalgBufferizePatterns(
411     MLIRContext *context, BufferizeTypeConverter &typeConverter,
412     OwningRewritePatternList &patterns) {
413   patterns.insert<BufferizeAnyLinalgOp>(typeConverter);
414   // TODO: Drop this once tensor constants work in standard.
415   patterns.insert<
416       // clang-format off
417       SubTensorOpConverter,
418       SubTensorInsertOpConverter,
419       TensorConstantOpConverter
420       // clang-format on
421       >(typeConverter, context);
422 }
423