1 | //===- IndependenceTransforms.cpp - Make ops independent of values --------===// |
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/Dialect/Tensor/Transforms/Transforms.h" |
10 | |
11 | #include "mlir/Dialect/Affine/IR/AffineOps.h" |
12 | #include "mlir/Dialect/Affine/Transforms/Transforms.h" |
13 | #include "mlir/Dialect/Tensor/IR/Tensor.h" |
14 | #include "mlir/Dialect/Utils/StaticValueUtils.h" |
15 | #include "mlir/Interfaces/ValueBoundsOpInterface.h" |
16 | |
17 | using namespace mlir; |
18 | using namespace mlir::tensor; |
19 | |
20 | /// Make the given OpFoldResult independent of all independencies. |
21 | static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc, |
22 | OpFoldResult ofr, |
23 | ValueRange independencies) { |
24 | if (ofr.is<Attribute>()) |
25 | return ofr; |
26 | Value value = ofr.get<Value>(); |
27 | AffineMap boundMap; |
28 | ValueDimList mapOperands; |
29 | if (failed(result: ValueBoundsConstraintSet::computeIndependentBound( |
30 | resultMap&: boundMap, mapOperands, type: presburger::BoundType::UB, var: value, |
31 | independencies, |
32 | /*closedUB=*/true))) |
33 | return failure(); |
34 | return mlir::affine::materializeComputedBound(b, loc, boundMap, mapOperands); |
35 | } |
36 | |
37 | FailureOr<Value> tensor::buildIndependentOp(OpBuilder &b, tensor::PadOp padOp, |
38 | ValueRange independencies) { |
39 | OpBuilder::InsertionGuard g(b); |
40 | b.setInsertionPoint(padOp); |
41 | Location loc = padOp.getLoc(); |
42 | |
43 | // Non-constant padding not supported. |
44 | Value constantPadding = padOp.getConstantPaddingValue(); |
45 | if (!constantPadding) |
46 | return failure(); |
47 | |
48 | SmallVector<OpFoldResult> newMixedLow, newMixedHigh; |
49 | for (OpFoldResult ofr : padOp.getMixedLowPad()) { |
50 | auto ub = makeIndependent(b, loc, ofr, independencies); |
51 | if (failed(ub)) |
52 | return failure(); |
53 | newMixedLow.push_back(*ub); |
54 | } |
55 | for (OpFoldResult ofr : padOp.getMixedHighPad()) { |
56 | auto ub = makeIndependent(b, loc, ofr, independencies); |
57 | if (failed(ub)) |
58 | return failure(); |
59 | newMixedHigh.push_back(*ub); |
60 | } |
61 | |
62 | // Return existing tensor::PadOp if nothing has changed. |
63 | if (llvm::equal(padOp.getMixedLowPad(), newMixedLow) && |
64 | llvm::equal(padOp.getMixedHighPad(), newMixedHigh)) |
65 | return padOp.getResult(); |
66 | |
67 | // Create a new tensor::PadOp. |
68 | auto newPadOp = b.create<PadOp>( |
69 | loc, padOp.getResultType(), padOp.getSource(), newMixedLow, newMixedHigh, |
70 | constantPadding, padOp.getNofold(), /*attrs=*/ArrayRef<NamedAttribute>{}); |
71 | |
72 | // Create a tensor::ExtractSliceOp. |
73 | // Reify the result sizes of the old tensor::PadOp. |
74 | ReifiedRankedShapedTypeDims reifiedSizes; |
75 | ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface = |
76 | dyn_cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation()); |
77 | if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedSizes))) |
78 | return failure(); |
79 | SmallVector<OpFoldResult> offsets, sizes, strides; |
80 | for (int64_t i = 0, e = padOp.getResultType().getRank(); i < e; ++i) { |
81 | // offset = ub(low_padding) - low_padding |
82 | OpFoldResult prevLow = padOp.getMixedLowPad()[i]; |
83 | if (prevLow.is<Attribute>()) { |
84 | offsets.push_back(b.getIndexAttr(0)); |
85 | } else { |
86 | offsets.push_back( |
87 | Elt: b.create<affine::AffineApplyOp>( |
88 | loc, b.getAffineDimExpr(position: 0) - b.getAffineDimExpr(position: 1), |
89 | std::initializer_list<Value>{newMixedLow[i].get<Value>(), |
90 | prevLow.get<Value>()}) |
91 | .getResult()); |
92 | } |
93 | // size = reified result size |
94 | if (!padOp.getResultType().isDynamicDim(i)) { |
95 | sizes.push_back(Elt: b.getIndexAttr(value: padOp.getResultType().getDimSize(i))); |
96 | } else { |
97 | sizes.push_back(Elt: reifiedSizes[0][i]); |
98 | } |
99 | // stride = 1 |
100 | strides.push_back(b.getIndexAttr(1)); |
101 | } |
102 | |
103 | return b.create<ExtractSliceOp>(loc, newPadOp, offsets, sizes, strides) |
104 | .getResult(); |
105 | } |
106 | |
107 | FailureOr<Value> tensor::buildIndependentOp(OpBuilder &b, |
108 | tensor::EmptyOp emptyOp, |
109 | ValueRange independencies) { |
110 | OpBuilder::InsertionGuard g(b); |
111 | b.setInsertionPoint(emptyOp); |
112 | Location loc = emptyOp.getLoc(); |
113 | |
114 | SmallVector<OpFoldResult> newSizes; |
115 | for (OpFoldResult ofr : emptyOp.getMixedSizes()) { |
116 | auto ub = makeIndependent(b, loc, ofr, independencies); |
117 | if (failed(ub)) |
118 | return failure(); |
119 | newSizes.push_back(*ub); |
120 | } |
121 | |
122 | // Return existing tensor::EmptyOp if nothing has changed. |
123 | if (llvm::equal(emptyOp.getMixedSizes(), newSizes)) |
124 | return emptyOp.getResult(); |
125 | |
126 | // Create a new tensor::EmptyOp. |
127 | Value newEmptyOp = |
128 | b.create<EmptyOp>(loc, newSizes, emptyOp.getType().getElementType()); |
129 | |
130 | // Create a tensor::ExtractSliceOp. |
131 | SmallVector<OpFoldResult> offsets(newSizes.size(), b.getIndexAttr(0)); |
132 | SmallVector<OpFoldResult> strides(newSizes.size(), b.getIndexAttr(1)); |
133 | return b |
134 | .create<ExtractSliceOp>(loc, newEmptyOp, offsets, emptyOp.getMixedSizes(), |
135 | strides) |
136 | .getResult(); |
137 | } |
138 | |