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
17using namespace mlir;
18using namespace mlir::tensor;
19
20/// Make the given OpFoldResult independent of all independencies.
21static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc,
22 OpFoldResult ofr,
23 ValueRange independencies) {
24 if (isa<Attribute>(Val: ofr))
25 return ofr;
26 Value value = cast<Value>(Val&: ofr);
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
37FailureOr<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 (isa<Attribute>(Val: prevLow)) {
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>{cast<Value>(Val&: newMixedLow[i]),
90 cast<Value>(Val&: prevLow)})
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
107FailureOr<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

source code of mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp