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

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