1 //===- BubbleUpExtractSlice.cpp - bubble up tensor.extract_slice ----------===//
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 // This file implements patterns that transforms linalg.<op> +
10 // tensor.extract_slice into tensor.extract_slice + linalg.<op> to reduce
11 // the computation for the linalg op.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "PassDetail.h"
16 #include "mlir/Dialect/Affine/IR/AffineOps.h"
17 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
18 #include "mlir/Dialect/Linalg/IR/Linalg.h"
19 #include "mlir/Dialect/Linalg/Passes.h"
20 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
21 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
22 
23 using namespace mlir;
24 using namespace mlir::linalg;
25 
26 namespace {
27 /// Bubble up extract_slice above Linalg operation.
28 ///
29 /// A sequence of operations
30 ///
31 /// ```mlir
32 /// %0 = linalg.<op> ... arg0, arg1, ...
33 /// %1 = tensor.extract_slice %0 ...
34 /// ```
35 ///
36 /// can be replaced with
37 ///
38 /// ```mlir
39 /// %0 = tensor.extract_slice %arg0
40 /// %1 = tensor.extract_slice %arg1
41 /// %2 = linalg.<op> ... %0, %1, ...
42 /// ```
43 ///
44 /// This results in the reduce computation of the linalg operation.
45 ///
46 struct BubbleUpExtractSliceOpPattern
47     : OpRewritePattern<tensor::ExtractSliceOp> {
48   using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
49 
50   LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
51                                 PatternRewriter &rewriter) const final {
52     Value source = sliceOp.source();
53     auto linalgOp = source.getDefiningOp<LinalgOp>();
54     if (!linalgOp) {
55       return rewriter.notifyMatchFailure(sliceOp,
56                                          "expected source to be linalg op");
57     }
58 
59     // TODO: we might relax this if we want heuristics to detect that all uses
60     // are small portion of the output.
61     if (!linalgOp->hasOneUse()) {
62       return rewriter.notifyMatchFailure(sliceOp,
63                                          "expected single use of linalg op");
64     }
65 
66     if (linalgOp.getNumOutputs() != 1) {
67       return rewriter.notifyMatchFailure(sliceOp,
68                                          "expected single output of linalg op");
69     }
70 
71     if (!linalgOp.hasTensorSemantics()) {
72       return rewriter.notifyMatchFailure(sliceOp,
73                                          "expected tensor of linalg op");
74     }
75 
76     if (!sliceOp.hasUnitStride())
77       return rewriter.notifyMatchFailure(sliceOp, "expected unit stride");
78 
79     OpOperand *outOperand = linalgOp.getOutputOperand(0);
80     AffineMap indexingMap = linalgOp.getTiedIndexingMap(outOperand);
81     if (!indexingMap.isProjectedPermutation()) {
82       return rewriter.notifyMatchFailure(
83           sliceOp, "expected a projected permutation for output");
84     }
85 
86     auto linalgLoc = linalgOp.getLoc();
87     auto allShapeSizes =
88         linalgOp.createFlatListOfOperandDims(rewriter, linalgLoc);
89     AffineMap shapeSizesToLoopsMap = linalgOp.getShapesToLoopsMap();
90     if (!shapeSizesToLoopsMap) {
91       return rewriter.notifyMatchFailure(
92           linalgOp, "failed to get loops map from shape sizes");
93     }
94     auto sizeBounds = applyMapToValues(rewriter, linalgLoc,
95                                        shapeSizesToLoopsMap, allShapeSizes);
96 
97     auto sliceLoc = sliceOp.getLoc();
98     auto offsetVals = getValueOrCreateConstantIndexOp(
99         rewriter, sliceLoc, sliceOp.getMixedOffsets());
100     auto sizeVals = getValueOrCreateConstantIndexOp(rewriter, sliceLoc,
101                                                     sliceOp.getMixedSizes());
102 
103     // The offsets and sizes from the slice operation only give you the tile
104     // size of the output. Use that compute the tile sizes and offsets of the
105     // loops. For loops not used to access the output, set the tile sizes to
106     // loop bounds and set the offset to 0.
107     Value zero = rewriter.create<arith::ConstantIndexOp>(linalgLoc, 0);
108     SmallVector<Value, 4> tileOffsets(sizeBounds.size(), zero);
109     SmallVector<Value, 4> tileSizes = sizeBounds;
110     for (auto const &result : enumerate(indexingMap.getResults())) {
111       unsigned position = result.value().cast<AffineDimExpr>().getPosition();
112       tileOffsets[position] = offsetVals[result.index()];
113       tileSizes[position] = sizeVals[result.index()];
114     }
115 
116     SmallVector<Value> valuesToTile = linalgOp.getInputAndOutputOperands();
117 
118     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
119         rewriter, linalgLoc, linalgOp, valuesToTile, tileOffsets, tileSizes,
120         sizeBounds, /*omitPartialTileCheck=*/true);
121 
122     SmallVector<Type, 4> resultTensorTypes;
123     for (OpOperand *opOperand : linalgOp.getOutputTensorOperands())
124       resultTensorTypes.push_back(
125           tiledOperands[opOperand->getOperandNumber()].getType());
126 
127     Operation *newOp =
128         linalgOp.clone(rewriter, linalgLoc, resultTensorTypes, tiledOperands);
129     rewriter.replaceOp(sliceOp, newOp->getResults());
130     return success();
131   }
132 };
133 } // namespace
134 
135 void mlir::linalg::populateBubbleUpExtractSliceOpPatterns(
136     RewritePatternSet &patterns) {
137   auto *context = patterns.getContext();
138   patterns.add<BubbleUpExtractSliceOpPattern>(context);
139 }
140