137 lines · cpp
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-exception6//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>(ofr))24 return ofr;25 Value value = cast<Value>(ofr);26 AffineMap boundMap;27 ValueDimList mapOperands;28 if (failed(ValueBoundsConstraintSet::computeIndependentBound(29 boundMap, mapOperands, presburger::BoundType::UB, 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(ub))51 return failure();52 newMixedLow.push_back(*ub);53 }54 for (OpFoldResult ofr : padOp.getMixedHighPad()) {55 auto ub = makeIndependent(b, loc, ofr, independencies);56 if (failed(ub))57 return failure();58 newMixedHigh.push_back(*ub);59 }60 61 // Return existing tensor::PadOp if nothing has changed.62 if (llvm::equal(padOp.getMixedLowPad(), newMixedLow) &&63 llvm::equal(padOp.getMixedHighPad(), newMixedHigh))64 return padOp.getResult();65 66 // Create a new tensor::PadOp.67 auto newPadOp =68 PadOp::create(b, loc, padOp.getResultType(), padOp.getSource(),69 newMixedLow, newMixedHigh, constantPadding,70 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_padding82 OpFoldResult prevLow = padOp.getMixedLowPad()[i];83 if (isa<Attribute>(prevLow)) {84 offsets.push_back(b.getIndexAttr(0));85 } else {86 offsets.push_back(87 affine::AffineApplyOp::create(88 b, 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 size94 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 = 1100 strides.push_back(b.getIndexAttr(1));101 }102 103 return ExtractSliceOp::create(b, 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 EmptyOp::create(b, 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 ExtractSliceOp::create(b, loc, newEmptyOp, offsets,134 emptyOp.getMixedSizes(), strides)135 .getResult();136}137