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 (isa<Attribute>(ofr)) 25 return ofr; 26 Value value = cast<Value>(ofr); 27 AffineMap boundMap; 28 ValueDimList mapOperands; 29 if (failed(ValueBoundsConstraintSet::computeIndependentBound( 30 boundMap, mapOperands, presburger::BoundType::UB, 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 (isa<Attribute>(prevLow)) { 84 offsets.push_back(b.getIndexAttr(0)); 85 } else { 86 offsets.push_back( 87 b.create<affine::AffineApplyOp>( 88 loc, b.getAffineDimExpr(0) - b.getAffineDimExpr(1), 89 std::initializer_list<Value>{cast<Value>(newMixedLow[i]), 90 cast<Value>(prevLow)}) 91 .getResult()); 92 } 93 // size = reified result size 94 if (!padOp.getResultType().isDynamicDim(i)) { 95 sizes.push_back(b.getIndexAttr(padOp.getResultType().getDimSize(i))); 96 } else { 97 sizes.push_back(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