372 lines · cpp
1//===- PaddingTilingInterface.cpp - Padding of TilingInterface ops --------===//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/Linalg/Transforms/Transforms.h"10 11#include "mlir/Dialect/Affine/IR/AffineOps.h"12#include "mlir/Dialect/Complex/IR/Complex.h"13#include "mlir/Dialect/Tensor/IR/Tensor.h"14#include "mlir/Dialect/UB/IR/UBOps.h"15#include "mlir/Dialect/Utils/StaticValueUtils.h"16#include "mlir/IR/AffineExpr.h"17#include "mlir/IR/BuiltinAttributes.h"18#include "mlir/IR/BuiltinTypeInterfaces.h"19#include "mlir/IR/BuiltinTypes.h"20#include "mlir/IR/OpDefinition.h"21#include "mlir/IR/Value.h"22#include "mlir/Interfaces/TilingInterface.h"23#include "llvm/ADT/STLExtras.h"24#include "llvm/Support/Casting.h"25 26#define DEBUG_TYPE "pad-tiling-interface"27 28using namespace mlir;29using namespace mlir::linalg;30using namespace mlir::tensor;31 32#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")33#define DBGSNL() (llvm::dbgs() << "\n")34 35/// Form a "full-rank" padding specification so that the application is easy.36static SmallVector<OpFoldResult>37getFullRankPaddingSizes(Builder &b, ArrayRef<OpFoldResult> indexingSizes,38 const PadTilingInterfaceOptions &options) {39 SmallVector<OpFoldResult> paddingSizes;40 // Complete the padding specification to specify all dimensions.41 for (size_t idx = 0, e = indexingSizes.size(); idx != e; ++idx) {42 // Complete to zero if needed.43 paddingSizes.push_back(options.paddingSizes.size() > idx44 ? options.paddingSizes[idx]45 : b.getIndexAttr(0));46 // If a dimension is zero (either specified or completed), replace by:47 // - 1 if we are padding to the next multiple of.48 // - indexingSizes[idx] otherwise49 if (isZeroInteger(paddingSizes[idx])) {50 paddingSizes[idx] =51 options.padToMultipleOf ? b.getIndexAttr(1) : indexingSizes[idx];52 }53 LLVM_DEBUG(DBGS() << "----idx: " << idx << " : " << paddingSizes[idx]54 << "\n");55 }56 return paddingSizes;57}58 59/// Extracts the constant multiplier from an affine expression of the form60/// `d * c` or `c * d`, where `d` is an AffineDimExpr and `c` is an61/// AffineConstantExpr. Returns 1 if the expression is not a simple62/// multiplication of a dimension and a constant.63static int64_t extractConstantMultiplier(AffineExpr expr) {64 if (auto binOp = dyn_cast<AffineBinaryOpExpr>(expr)) {65 if (binOp.getKind() == AffineExprKind::Mul) {66 auto lhsD = dyn_cast<AffineDimExpr>(binOp.getLHS());67 auto rhsC = dyn_cast<AffineConstantExpr>(binOp.getRHS());68 if (lhsD && rhsC) {69 return rhsC.getValue();70 }71 auto lhsC = dyn_cast<AffineConstantExpr>(binOp.getLHS());72 auto rhsD = dyn_cast<AffineDimExpr>(binOp.getRHS());73 if (lhsC && rhsD) {74 return lhsC.getValue();75 }76 }77 }78 return 1;79}80 81/// Compute the padded shape of the given value `v` of `RankedTensorType` given82/// - `indexingSizes` a list of OpFoldResult.83/// - an `indexingMap` that encodes how the shape of varies with increases84/// in `indexingSizes`.85/// The `indexingMap` encodes how the shape of varies with `indexingSizes`.86/// The `indexingMap` + `indexingSizes` encoding suits StructuredOps.87/// The implementaiton below iteratively combines increases from contributing88/// dimensions using affine.apply operations.89/// The padded shape is computed by evaluating the maximum accessed index per90/// dimension, which may involve multiplying by constant factors derived from91/// the affine indexing expressions. Currently, only a limited set of projected92/// permutation indexing maps are supported, such as93/// - affine_map<(d0, d1, d2) -> (d0, d1)>94/// - affine_map<(d0, d1, d2) -> (d0, d1 + d2)>95/// - affine_map<(d0, d1) -> (d0 * 3 + d1)>96/// In the future, more general interfaces can be devised to encode similar97/// shape evolutions and map between an op and its operands.98SmallVector<OpFoldResult>99linalg::computePaddedShape(OpBuilder &builder, TypedValue<RankedTensorType> v,100 AffineMap indexingMap,101 ArrayRef<OpFoldResult> indexingSizes,102 const PadTilingInterfaceOptions &options) {103 Location loc = v.getLoc();104 SmallVector<OpFoldResult> paddedShape;105 auto tensorType = cast<RankedTensorType>(v.getType());106 paddedShape.resize_for_overwrite(tensorType.getRank());107 assert(tensorType.getRank() == indexingMap.getNumResults() &&108 "expect the number of results of the affine map to match the tensor "109 "rank");110 111 // "Full-rank" padding specification.112 SmallVector<OpFoldResult> paddingSizes =113 getFullRankPaddingSizes(builder, indexingSizes, options);114 115 // For each dimension in the operand's shape, iterate over indexingSizes and116 // add the various term contributions.117 for (const auto &enResults : enumerate(indexingMap.getResults())) {118 int64_t resultIndex = enResults.index();119 AffineMap partialIndexingMap = indexingMap.getSubMap(120 ArrayRef<unsigned>{static_cast<unsigned>(resultIndex)});121 122 LLVM_DEBUG(DBGS() << "----resultIndex: " << resultIndex123 << " with partialIndexingMap: " << partialIndexingMap124 << "\n");125 126 // Find all padding dimensions that contribute to this operand dimension127 // and compute the padded term contribution to the final padded shape.128 SmallVector<OpFoldResult> terms;129 for (size_t paddingDim = 0, e = paddingSizes.size(); paddingDim != e;130 ++paddingDim) {131 OpFoldResult paddingSize = paddingSizes[paddingDim];132 LLVM_DEBUG(DBGS() << "------try apply padding of dim: " << paddingDim133 << " to: " << paddingSize << "\n");134 if (!enResults.value().isFunctionOfDim(paddingDim))135 continue;136 137 LLVM_DEBUG(DBGS() << "------apply padding of dim: " << paddingDim138 << " to: " << paddingSize << "\n");139 140 // Project non-'paddingDim' dimensions and compress the result.141 llvm::SmallBitVector projectedDims(partialIndexingMap.getNumDims(), true);142 projectedDims.flip(paddingDim);143 AffineMap projectedMap =144 mlir::projectDims(partialIndexingMap, projectedDims,145 /*compressDimsFlag=*/true);146 147 // If we are padding to the next multiple of, compose with ceil(sz) * sz.148 OpFoldResult paddingDimOfr;149 if (options.padToMultipleOf) {150 AffineExpr d0, s0;151 bindDims(builder.getContext(), d0);152 bindSymbols(builder.getContext(), s0);153 AffineMap ceilMap = AffineMap::get(1, 1, d0.ceilDiv(s0) * s0);154 AffineMap composedMap = projectedMap.compose(ceilMap);155 paddingDimOfr = affine::makeComposedFoldedAffineApply(156 builder, loc, composedMap, {indexingSizes[paddingDim], paddingSize},157 /*composeAffineMin=*/true);158 } else {159 // Otherwise just set to paddingSize.160 paddingDimOfr = affine::makeComposedFoldedAffineApply(161 builder, loc, projectedMap, paddingSize);162 }163 164 // Adjust for the maximum accessed index, which is (paddingSize - 1) *165 // multiplier.166 AffineExpr d0;167 bindDims(builder.getContext(), d0);168 int64_t multiplier = extractConstantMultiplier(projectedMap.getResult(0));169 AffineMap subtractMap = AffineMap::get(1, 0, d0 - multiplier);170 OpFoldResult maxAccessIdx = affine::makeComposedFoldedAffineApply(171 builder, loc, subtractMap, {paddingDimOfr});172 terms.push_back(maxAccessIdx);173 174 LLVM_DEBUG(DBGS() << "------new term: " << terms.back() << "\n");175 }176 177 // If there are no terms, just return the dim.178 if (terms.empty()) {179 paddedShape[resultIndex] =180 createFoldedDimOp(builder, loc, v, resultIndex);181 continue;182 }183 184 // Sum individual terms' contributions.185 SmallVector<AffineExpr> dims(terms.size());186 bindDimsList(builder.getContext(), MutableArrayRef{dims});187 AffineExpr sumExpr = dims.front();188 for (unsigned i = 1; i < dims.size(); ++i)189 sumExpr = sumExpr + dims[i];190 // Add 1 to the maximum accessed index and get the final padded size.191 OpFoldResult paddedDimOfr =192 affine::makeComposedFoldedAffineApply(builder, loc, sumExpr + 1, terms);193 paddedShape[resultIndex] = paddedDimOfr;194 }195 196 return paddedShape;197}198 199FailureOr<SmallVector<OpFoldResult>>200linalg::computeIndexingMapOpInterfacePaddedShape(201 OpBuilder &builder, OpOperand &operandToPad,202 ArrayRef<Range> iterationDomain, const PadTilingInterfaceOptions &options) {203 auto transferOp =204 llvm::dyn_cast<IndexingMapOpInterface>(operandToPad.getOwner());205 if (!transferOp)206 return failure();207 208 // clang-format off209 assert(llvm::all_of(iterationDomain, [&builder](Range r) {210 return r.offset == OpFoldResult(builder.getIndexAttr(0)) &&211 r.stride == OpFoldResult(builder.getIndexAttr(1));212 }) && "expected 0-offset 1-stride loop ranges");213 // clang-format on214 SmallVector<OpFoldResult> loopUpperBounds;215 loopUpperBounds.reserve(iterationDomain.size());216 for (const Range &range : iterationDomain)217 loopUpperBounds.push_back(range.size);218 219 AffineMap indexingMap = transferOp.getMatchingIndexingMap(&operandToPad);220 return computePaddedShape(221 builder, cast<TypedValue<RankedTensorType>>(operandToPad.get()),222 indexingMap, loopUpperBounds, options);223}224 225/// Pad a single operand to `paddedShape` using `paddingValueAttr` as padding226/// Value.227static Value padOperand(OpBuilder &builder, TilingInterface opToPad,228 TypedValue<RankedTensorType> v,229 ArrayRef<OpFoldResult> paddedShape,230 Attribute paddingValueAttr) {231 Value paddingValue;232 if (auto complexTy =233 dyn_cast<ComplexType>(getElementTypeOrSelf(v.getType()))) {234 if (auto complexAttr = dyn_cast<ArrayAttr>(paddingValueAttr)) {235 paddingValue = complex::ConstantOp::create(builder, opToPad.getLoc(),236 complexTy, complexAttr);237 }238 } else if (isa<ub::PoisonAttr>(paddingValueAttr)) {239 paddingValue = ub::PoisonOp::create(builder, opToPad.getLoc(),240 getElementTypeOrSelf(v.getType()));241 } else if (auto typedAttr = dyn_cast<TypedAttr>(paddingValueAttr)) {242 paddingValue =243 arith::ConstantOp::create(builder, opToPad.getLoc(), typedAttr);244 }245 assert(paddingValue && "failed to create value from padding attribute");246 247 // Pad the operand to the bounding box defined by `paddedShape`.248 SmallVector<int64_t> tensorShape;249 SmallVector<Value> dynDims;250 for (OpFoldResult ofr : paddedShape) {251 std::optional<int64_t> cst = getConstantIntValue(ofr);252 tensorShape.push_back(cst.has_value() ? *cst : ShapedType::kDynamic);253 if (!cst.has_value())254 dynDims.push_back(ofr.dyn_cast<Value>());255 }256 // TODO: use dispatchIndexOpFoldResults(paddedShape, dynDims, paddedShape);257 258 auto paddedTensorType =259 RankedTensorType::get(tensorShape, getElementTypeOrSelf(v));260 LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "261 << paddedTensorType);262 return makeComposedPadHighOp(builder, opToPad.getLoc(), paddedTensorType, v,263 paddingValue, /*nofold=*/false, dynDims);264}265 266FailureOr<PadTilingInterfaceResult> linalg::rewriteAsPaddedOp(267 OpBuilder &builder, TilingInterface toPad,268 PadTilingInterfaceOptions options,269 const PadSizeComputationFunction &computePaddingSizeFun) {270 LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << toPad << "\n");271 SmallVector<tensor::PadOp> padOps;272 Location loc = toPad.getLoc();273 274 // Allow inference of pad values if they are not explicitly specified.275 // TODO: be mindful about the value depending on the actual operation.276 if (options.paddingValues.empty()) {277 SmallVector<Type> types(toPad->getOperandTypes());278 llvm::append_range(types, toPad->getResultTypes());279 for (Type t : types) {280 options.paddingValues.push_back(281 builder.getZeroAttr(getElementTypeOrSelf(t)));282 }283 }284 285 if (llvm::any_of(toPad->getOperands(),286 [](Value v) { return isa<MemRefType>(v.getType()); })) {287 LLVM_DEBUG(DBGS() << "Not an operation on tensors: FAIL\n");288 return failure();289 }290 291 // 1. Get the loopUpperBounds from the TilingInterface.292 SmallVector<Range> iterationDomain = toPad.getIterationDomain(builder);293 294 // 2. For each operand.295 SmallVector<Value> newOperands;296 newOperands.reserve(toPad->getNumOperands());297 for (OpOperand &opOperand : toPad->getOpOperands()) {298 Value operand = opOperand.get();299 LLVM_DEBUG(DBGS() << "--start padding operand: " << operand << "\n");300 301 // 2.a. Skip scalar-like operands.302 Type operandType = operand.getType();303 if (!isa<RankedTensorType>(operandType)) {304 assert((!isa<ShapedType>(operandType) || isa<VectorType>(operandType)) &&305 "Unexpected non-vector ShapedType");306 newOperands.push_back(operand);307 continue;308 }309 310 // 2.a. Compute padded shape.311 FailureOr<SmallVector<OpFoldResult>> maybePaddedShape =312 computePaddingSizeFun(builder, opOperand, iterationDomain, options);313 if (failed(maybePaddedShape)) {314 LLVM_DEBUG(DBGS() << "Could not get padded shape of operand: FAIL\n");315 return failure();316 }317 318 // 2.b. Expect proper `paddingValues`.319 // TODO: we may want to allow garbage padding in the future, in which case320 // we would just not assert.321 if (opOperand.getOperandNumber() >= options.paddingValues.size()) {322 LLVM_DEBUG(DBGS() << "Too few padding values specified: FAIL\n");323 return failure();324 }325 Attribute paddingValueAttr =326 options.paddingValues[opOperand.getOperandNumber()];327 328 // 2.c. Perform actual padding.329 Value paddedOperand =330 padOperand(builder, toPad, cast<TypedValue<RankedTensorType>>(operand),331 *maybePaddedShape, paddingValueAttr);332 LLVM_DEBUG(DBGS() << "--done padding operand: " << paddedOperand << "\n");333 334 newOperands.push_back(paddedOperand);335 if (auto padOp = paddedOperand.getDefiningOp<tensor::PadOp>())336 padOps.push_back(padOp);337 }338 339 // 3. Form the resulting tensor::ExtractSliceOp.340 ReifiedRankedShapedTypeDims reifiedResultShapes;341 if (failed(reifyResultShapes(builder, toPad, reifiedResultShapes))) {342 LLVM_DEBUG(DBGS() << "Failed to reify result shapes: FAIL\n");343 return failure();344 }345 assert(reifiedResultShapes.size() == toPad->getNumResults() &&346 "expected same number of results");347 348 // Clone `toPad` to operate on the statically padded shapes.349 auto resultTensorTypes =350 ValueRange(newOperands).take_back(toPad->getNumResults()).getTypes();351 // clone **should** properly notify the builder.352 TilingInterface paddedOp =353 clone(builder, toPad, resultTensorTypes, newOperands);354 LLVM_DEBUG(DBGS() << "--cloned padded op: " << paddedOp << "\n");355 356 // Recover the slice out of the new static results.357 SmallVector<Value> paddedSubtensorResults;358 paddedSubtensorResults.reserve(toPad->getNumResults());359 for (const auto &en : llvm::enumerate(paddedOp->getResults())) {360 Value paddedResult = en.value();361 int64_t resultNumber = en.index();362 int64_t rank = cast<RankedTensorType>(paddedResult.getType()).getRank();363 SmallVector<OpFoldResult> offsets(rank, builder.getIndexAttr(0));364 SmallVector<OpFoldResult> strides(rank, builder.getIndexAttr(1));365 paddedSubtensorResults.push_back(tensor::ExtractSliceOp::create(366 builder, loc, paddedResult, offsets, reifiedResultShapes[resultNumber],367 strides));368 }369 370 return PadTilingInterfaceResult{padOps, paddedOp, paddedSubtensorResults};371}372