[Mlir-commits] [mlir] [mlir][linalg] Fix padding transform and update transform padding op. (PR #144354)
Fabian Mora
llvmlistbot at llvm.org
Mon Jun 16 09:41:44 PDT 2025
https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/144354
>From e59c3ffebc86b46bd2c226683b266a4463e48f45 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fabian.mora-cordero at amd.com>
Date: Mon, 16 Jun 2025 13:43:10 +0000
Subject: [PATCH 1/2] [mlir][linalg] Fix padding transform and update transform
padding op.
---
.../Linalg/TransformOps/LinalgTransformOps.td | 4 +-
.../Dialect/Linalg/Transforms/Transforms.h | 16 ++
.../include/mlir/Dialect/Linalg/Utils/Utils.h | 17 +-
.../include/mlir/Dialect/Tensor/Utils/Utils.h | 2 +-
.../TransformOps/LinalgTransformOps.cpp | 9 +
.../lib/Dialect/Linalg/Transforms/Padding.cpp | 250 ++++++++++++++----
mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 24 +-
mlir/lib/Dialect/Tensor/Utils/Utils.cpp | 3 +-
.../test/Dialect/Linalg/transform-op-pad.mlir | 2 +-
9 files changed, 250 insertions(+), 77 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 15ea5e7bf7159..bf56b633d2872 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1134,7 +1134,8 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
DefaultValuedAttr<
TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
"{}">:$transpose_paddings,
- DefaultValuedAttr<StrAttr, "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copy_back_op);
+ DefaultValuedAttr<StrAttr, "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copy_back_op,
+ UnitProp:$use_default_tensor_shapes);
let results = (outs TransformHandleTypeInterface:$padded,
TransformHandleTypeInterface:$pad,
TransformHandleTypeInterface:$copy);
@@ -1142,6 +1143,7 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
let assemblyFormat = [{
$target
(`pad_to_multiple_of` custom<DynamicIndexList>($pad_to_multiple_of, $static_pad_to_multiple_of)^)?
+ (`use_default_tensor_shapes` $use_default_tensor_shapes^)?
attr-dict
`:` functional-type(operands, results)
}];
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 2eef0a06d0eb4..7004d98c22475 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -295,6 +295,22 @@ struct LinalgPaddingOptions {
padToMultipleOf.emplace(m.begin(), m.end());
return *this;
}
+ /// A mapping between an operand and shape dim, and a size for a padding
+ /// dimension. Each size is expected to be greater or equal than the
+ /// corresponding shape dim. If no size is provided for a padding dim, then
+ /// the corresponding tensor size will be used when padding.
+ DenseMap<std::pair<unsigned, unsigned>, OpFoldResult> padToSizeOf;
+ LinalgPaddingOptions &setPadToSizeOf(unsigned operandIndex, unsigned dimIndex,
+ OpFoldResult size) {
+ assert(size && "expected non-null size");
+ padToSizeOf[{operandIndex, dimIndex}] = size;
+ return *this;
+ }
+ /// Populates the `padToSizeOf` map with constant upper bounds for each padded
+ /// dim and operand of `opToPad`. Returns failure if any of the sizes cannot
+ /// be computed.
+ LogicalResult computeConstantUpperShapeBounds(linalg::LinalgOp opToPad);
+
/// A flag for every operand to mark the PadOp as nofold which enables
/// packing for statically shaped operands.
SmallVector<bool> nofoldFlags;
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 80aa034d2199d..fc151d02ceef6 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -71,12 +71,14 @@ bool isParallelIterator(utils::IteratorType iteratorType);
/// Check if iterator type has "reduction" semantics.
bool isReductionIterator(utils::IteratorType iteratorType);
-/// Create a tensor::PadOp that pads `source` to the size of the statically
-/// sized `type` whose static sizes are assumed to be greater than the dynamic
-/// `source` size. The padding introduces trailing `pad` values until the
-/// target size is met. If `source` is defined by one or more LinalgOps that
-/// have been padded with the same value and sizes, return their padded result
-/// instead of creating a tensor::PadOp.
+/// Create a tensor::PadOp that pads `source` to the shape of `type` whose sizes
+/// are assumed to be greater than the dynamic `source` size. If `typeDynDims`
+/// is specified, then it must contain the sizes of all the dynamic dimensions
+/// in order of appearance in `type`, otherwise the function will pad those
+/// values to `0`. The padding introduces trailing `pad` values until the target
+/// size is met. If `source` is defined by one or more LinalgOps that have been
+/// padded with the same value and sizes, return their padded result instead of
+/// creating a tensor::PadOp.
///
/// Example:
/// ```
@@ -91,7 +93,8 @@ bool isReductionIterator(utils::IteratorType iteratorType);
/// %4 = tensor.pad %3 low[0, 0] high[...] { tensor.yield %other_cst }
/// ```
Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
- Value source, Value pad, bool nofold);
+ Value source, Value padding, bool nofold,
+ ValueRange typeDynDims = std::nullopt);
/// Returns GenericOp that copies an n-D memref. Unlike the current
/// implementation of memref::CopyOp, this op can further tile, lower to loops
diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index 1a4733df3f187..a1ce4e252c2f4 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -30,7 +30,7 @@ namespace tensor {
// for _static_ dimensions.
PadOp createPadHighOp(RankedTensorType resType, Value source, Value pad,
bool nofold, Location loc, OpBuilder &builder,
- SmallVector<Value> dynOutDims = {});
+ ValueRange dynOutDims = std::nullopt);
// Creates dim ops for each dynamic dimension of the ranked tensor argument and
// returns these as values.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index b2c28f5eed33c..45d83b9a8d3c8 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2051,6 +2051,15 @@ transform::PadOp::apply(transform::TransformRewriter &rewriter,
} else {
llvm_unreachable("unsupported copy_back op");
}
+ if (!getUseDefaultTensorShapes()) {
+ if (failed(options.computeConstantUpperShapeBounds(linalgTarget))) {
+ auto diag =
+ emitSilenceableError()
+ << "could not compute upper constant bounds for padded dims";
+ diag.attachNote(target->getLoc()) << "target op";
+ return diag;
+ }
+ }
SmallVector<Value> replacements;
SmallVector<tensor::PadOp> newPadOps;
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp b/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
index 9a685f6dc96ac..475f338ffbc90 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
@@ -8,10 +8,12 @@
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#define DEBUG_TYPE "linalg-padding"
@@ -22,69 +24,147 @@ using namespace mlir::linalg;
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
#define DBGSNL() (llvm::dbgs() << "\n")
-/// Compute the padded shape of the given operand. The operand is padded to a
-/// static bounding box according to the specified padding options.
-static LogicalResult computePaddedShape(linalg::LinalgOp opToPad,
- OpOperand *opOperand,
- const LinalgPaddingOptions &options,
- SmallVector<int64_t> &paddedShape,
- bool &alreadyHasRequestedShape) {
+namespace {
+/// Helper class for storing padding information.
+struct PaddingInfo {
+ PaddingInfo(int64_t padToMultipleOf = 1, OpFoldResult size = {})
+ : padToMultipleOf(padToMultipleOf), size(size) {}
+ bool isTrivial() const { return padToMultipleOf == 1 && size.isNull(); }
+ /// Pad the tensor to a multiple of.
+ int64_t padToMultipleOf = 1;
+ /// The size used for padding.
+ OpFoldResult size = {};
+};
+
+/// Helper class for storing and computing the padded shape.
+struct PaddedShape {
+ /// Initializes the shape information and returns whether the shape of the
+ /// operand will change.
+ bool initialize(linalg::LinalgOp opToPad, OpOperand *opOperand,
+ const LinalgPaddingOptions &options);
+
+ /// Computs the padded shape.
+ void computePadding(OpBuilder &builder, Value operand);
+
+ /// Returns the new tensor type.
+ RankedTensorType getType(Type elemTy) {
+ return RankedTensorType::get(shape, elemTy);
+ }
+
+ /// Return the dynamic dimensions of the shape.
+ ValueRange getDynamicDims() { return dynDims; }
+
+private:
+ SmallVector<int64_t> shape;
+ SmallVector<Value> dynDims;
+ DenseMap<int64_t, PaddingInfo> dimToInfo;
+};
+} // namespace
+
+bool PaddedShape::initialize(linalg::LinalgOp opToPad, OpOperand *opOperand,
+ const LinalgPaddingOptions &options) {
AffineMap indexingMap = opToPad.getMatchingIndexingMap(opOperand);
- ArrayRef<int64_t> shape = opToPad.getShape(opOperand);
+
+ // Initialize the padded shape.
+ llvm::append_range(shape, opToPad.getShape(opOperand));
+
+ // Return early if there's no padding involved.
+ if (!options.padToMultipleOf && options.padToSizeOf.empty())
+ return true;
// Collect the shape dimensions that are a function of "paddingDimensions",
// along with the multiple that they should be padded to ("1" if none).
- alreadyHasRequestedShape = true;
- DenseMap<int64_t, int64_t> shapeDimToMultiple;
- for (const auto &dimEn : enumerate(options.paddingDimensions)) {
- for (const auto &en : enumerate(indexingMap.getResults())) {
- if (en.value().isFunctionOfDim(dimEn.value())) {
- int64_t dimSize = shape[en.index()];
+ bool alreadyHasRequestedShape = true;
+ for (const auto [shapeIndex, shapeExpr] :
+ enumerate(indexingMap.getResults())) {
+ PaddingInfo paddingInfo;
+
+ // Flag indicating whether the current dim in the operand has to be padded.
+ bool isPaddedDim = false;
+
+ // Construct the padding info according to the options.
+ for (const auto [dimIndex, dim] : enumerate(options.paddingDimensions)) {
+ if (shapeExpr.isFunctionOfDim(dim)) {
+ isPaddedDim = true;
if (options.padToMultipleOf.has_value()) {
- shapeDimToMultiple[en.index()] =
- (*options.padToMultipleOf)[dimEn.index()];
- } else {
- shapeDimToMultiple[en.index()] = 1;
- }
- if (ShapedType::isDynamic(dimSize)) {
- alreadyHasRequestedShape = false;
- } else if (dimSize % shapeDimToMultiple[en.index()] != 0) {
- alreadyHasRequestedShape = false;
+ // We use the least common multiple as multiple dim iterators can
+ // appear in a shape dim, for example:
+ // `affine_map<(i, j, k) -> (i, k, i + j)>` could impose more than one
+ // multiple of constraint on the last dim of the shape.
+ paddingInfo.padToMultipleOf =
+ std::lcm((*options.padToMultipleOf)[dimIndex],
+ paddingInfo.padToMultipleOf);
}
}
}
+
+ // If the dimension is not being padded, continue.
+ if (!isPaddedDim)
+ continue;
+
+ // Check if the dim is being padded to a specified size.
+ if (auto it = options.padToSizeOf.find(
+ {opOperand->getOperandNumber(), shapeIndex});
+ it != options.padToSizeOf.end()) {
+ paddingInfo.size = it->second;
+ assert(paddingInfo.size && "expected non-null `OpFoldResult`");
+ }
+
+ int64_t dimSize = shape[shapeIndex];
+
+ // Skip if the padding information is trivial. Note that dynamic
+ // dimensions never have trivial padding information.
+ if (paddingInfo.isTrivial() && !ShapedType::isDynamic(dimSize))
+ continue;
+
+ // Set the padding info.
+ dimToInfo[shapeIndex] = paddingInfo;
+ if (ShapedType::isDynamic(dimSize) ||
+ dimSize % paddingInfo.padToMultipleOf != 0) {
+ alreadyHasRequestedShape = false;
+ }
}
- // Helper function to round a number up to a given multiple.
- auto ceil = [](int64_t val, int64_t multiple) {
- return ((val + multiple - 1) / multiple) * multiple;
- };
+ return alreadyHasRequestedShape;
+}
+
+void PaddedShape::computePadding(OpBuilder &builder, Value operand) {
+ Location loc = operand.getLoc();
+ AffineExpr sizeSym = builder.getAffineSymbolExpr(0);
- // Upper bound the sizes to obtain a static bounding box.
- paddedShape.assign(shape.begin(), shape.end());
- for (int64_t i = 0, e = shape.size(); i < e; ++i) {
+ // Compute the padding for each dimension.
+ for (auto &&[i, dim] : llvm::enumerate(shape)) {
LLVM_DEBUG(DBGS() << "--compute padded size for dim " << i << "\n");
+
// Skip dimensions that do not require padding.
- if (!shapeDimToMultiple.contains(i)) {
+ if (!dimToInfo.contains(i)) {
LLVM_DEBUG(DBGS() << "----dim does not require padding, SKIP\n");
+ if (ShapedType::isDynamic(dim)) {
+ dynDims.push_back(
+ cast<Value>(tensor::getMixedSize(builder, loc, operand, i)));
+ }
continue;
}
- // Otherwise, try to compute a constant upper bound for the size value.
- FailureOr<int64_t> upperBound =
- ValueBoundsConstraintSet::computeConstantBound(
- presburger::BoundType::UB,
- {opOperand->get(),
- /*dim=*/i},
- /*stopCondition=*/nullptr, /*closedUB=*/true);
- if (failed(upperBound)) {
- LLVM_DEBUG(DBGS() << "----could not compute a bounding box for padding");
- return failure();
+ PaddingInfo paddingInfo = dimToInfo[i];
+ OpFoldResult size = paddingInfo.size;
+ // Get the tensor dim size if none was provided.
+ if (size.isNull())
+ size = tensor::getMixedSize(builder, loc, operand, i);
+
+ // Compute the padded size to be a multiple of `padToMultipleOf`.
+ AffineExpr szExpr = (sizeSym).ceilDiv(paddingInfo.padToMultipleOf) *
+ paddingInfo.padToMultipleOf;
+ OpFoldResult paddedSize =
+ affine::makeComposedFoldedAffineApply(builder, loc, szExpr, size);
+ assert(paddedSize && "invalid arguments to affine apply");
+ if (auto cstSzAttr = dyn_cast<Attribute>(paddedSize)) {
+ dim = cast<IntegerAttr>(cstSzAttr).getValue().getZExtValue();
+ } else {
+ dim = ShapedType::kDynamic;
+ dynDims.push_back(cast<Value>(paddedSize));
}
- paddedShape[i] = ceil(*upperBound, shapeDimToMultiple[i]);
- LLVM_DEBUG(DBGS() << "----new dim size: " << paddedShape[i] << "\n");
+ LLVM_DEBUG(DBGS() << "----new dim size: " << paddedSize << "\n");
}
-
- return success();
}
/// Pad the `opOperand` in the "paddingDimensions" using the padding value and
@@ -107,13 +187,9 @@ static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
options.padToMultipleOf->size() == options.paddingDimensions.size()) &&
"invalid number of elements in padToMultipleOf");
- // Compute padded shape.
- SmallVector<int64_t> paddedShape;
- bool alreadyHasRequestedShape = false;
- if (failed(computePaddedShape(opToPad, opOperand, options, paddedShape,
- alreadyHasRequestedShape)))
- return rewriter.notifyMatchFailure(opToPad,
- "--failed to compute padded shape");
+ // Initialize the padded shape.
+ PaddedShape shape;
+ bool alreadyHasRequestedShape = shape.initialize(opToPad, opOperand, options);
// Return the unpadded operand if padding to a static shape is not needed and
// if the nofold flag is not set.
@@ -140,13 +216,73 @@ static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
}
- // Pad the operand to the bounding box defined by `paddedShape`.
- auto paddedTensorType = RankedTensorType::get(
- paddedShape, getElementTypeOrSelf(opOperand->get()));
+ // If needed, the padding for each dimension.
+ if (!alreadyHasRequestedShape)
+ shape.computePadding(rewriter, opOperand->get());
+
+ // Compute the new tensor type.
+ RankedTensorType paddedTensorType =
+ shape.getType(getElementTypeOrSelf(opOperand->get()));
LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
- << paddedTensorType);
+ << paddedTensorType << "\n");
+
+ // Pad the operand to the bounding box defined by `shape`.
return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
- opOperand->get(), paddingValue, nofold);
+ opOperand->get(), paddingValue, nofold,
+ shape.getDynamicDims());
+}
+
+LogicalResult LinalgPaddingOptions::computeConstantUpperShapeBounds(
+ linalg::LinalgOp opToPad) {
+ LLVM_DEBUG(DBGS() << "-Computing constant upper bounds for " << opToPad
+ << "\n");
+ if (!opToPad.hasPureTensorSemantics()) {
+ LLVM_DEBUG(DBGS() << "--FAILURE: op does not have pure tensor semantics"
+ << "\n");
+ return failure();
+ }
+ Builder builder(opToPad.getContext());
+ // For each operand compute the sizes.
+ for (OpOperand &operand : opToPad->getOpOperands()) {
+ AffineMap indexingMap = opToPad.getMatchingIndexingMap(&operand);
+ LLVM_DEBUG(DBGS() << "--Computing constant upper bounds for operand "
+ << operand.getOperandNumber() << "\n");
+
+ // Get the size for each dimension.
+ for (const auto [shapeIndex, shapeExpr] :
+ enumerate(indexingMap.getResults())) {
+
+ // Get whether the iterator dimension is being padded.
+ // TODO[c++20]: Remove the `expr = shapeExpr` copy, this was added as
+ // `captured structured bindings are a C++20 extension`.
+ AffineExpr expr = shapeExpr;
+ bool isPaddedDim = llvm::any_of(paddingDimensions, [expr](unsigned dim) {
+ return expr.isFunctionOfDim(dim);
+ });
+ if (!isPaddedDim)
+ continue;
+
+ // Compute the constant upper bound.
+ LLVM_DEBUG(DBGS() << "---compute upper bound size for shape dim "
+ << shapeIndex << "\n");
+ FailureOr<int64_t> upperBound =
+ ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::UB,
+ {operand.get(),
+ /*dim=*/static_cast<int64_t>(shapeIndex)},
+ /*stopCondition=*/nullptr, /*closedUB=*/true);
+ if (failed(upperBound)) {
+ LLVM_DEBUG(DBGS() << "---could not compute a bounding box for padding"
+ << "\n");
+ return failure();
+ }
+
+ // Set the upper bound.
+ setPadToSizeOf(operand.getOperandNumber(), shapeIndex,
+ builder.getIndexAttr(*upperBound));
+ }
+ }
+ return success();
}
LogicalResult
diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 2527d90cfa2e6..209309ddb413a 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -244,11 +244,13 @@ bool isReductionIterator(utils::IteratorType iteratorType) {
}
Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
- Value source, Value pad, bool nofold) {
+ Value source, Value pad, bool nofold,
+ ValueRange typeDynDims) {
// Exit if `source` is not defined by an ExtractSliceOp.
auto sliceOp = source.getDefiningOp<tensor::ExtractSliceOp>();
if (!sliceOp)
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Search the `source` use-def chain for padded LinalgOps.
Value current = sliceOp.getSource();
@@ -264,24 +266,28 @@ Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
// Exit if the search fails to match a tensor::PadOp at the end of the matched
// LinalgOp sequence.
if (!padOp)
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Exit if the padded result type does not match.
if (sliceOp.getSource().getType() != type)
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Exit if the LinalgOps are not high padded.
if (llvm::any_of(padOp.getMixedLowPad(), [](OpFoldResult ofr) {
return getConstantIntValue(ofr) != static_cast<int64_t>(0);
}))
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Exit if `padOpSliceOp`, which defines the slice used by
// `padOp`, is rank-reducing.
auto padOpSliceOp = padOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
if (!padOpSliceOp ||
sliceOp.getMixedSizes().size() != padOpSliceOp.getMixedSizes().size())
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Exit if the sizes of the dynamic sizes of `sliceOp` do not match the size
// of the slice padded by `padOp`.
@@ -290,14 +296,16 @@ Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
[](std::tuple<OpFoldResult, OpFoldResult> it) {
return !isEqualConstantIntOrValue(std::get<0>(it), std::get<1>(it));
}))
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Exit if the padding values do not match.
Attribute padOpPadAttr, padAttr;
Value padOpPad = padOp.getConstantPaddingValue();
if (!padOpPad || !matchPattern(padOpPad, m_Constant(&padOpPadAttr)) ||
!matchPattern(pad, m_Constant(&padAttr)) || padOpPadAttr != padAttr)
- return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
+ return tensor::createPadHighOp(type, source, pad, nofold, loc, b,
+ typeDynDims);
// Return the padded result if the padding values and sizes match.
return sliceOp.getSource();
diff --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index 11ae0108594dd..289296a07d9d3 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -24,8 +24,7 @@ using namespace mlir::tensor;
PadOp mlir::tensor::createPadHighOp(RankedTensorType resType, Value source,
Value pad, bool nofold, Location loc,
- OpBuilder &b,
- SmallVector<Value> dynOutDims) {
+ OpBuilder &b, ValueRange dynOutDims) {
// This assumption simplifies the following logic without limiting what's
// required _today_. If needed, we can relax it in the future.
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad.mlir b/mlir/test/Dialect/Linalg/transform-op-pad.mlir
index ab2711545405e..926b63311acf6 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad.mlir
@@ -300,7 +300,7 @@ func.func @negative_no_ub_estimate(%arg0: tensor<?x12xf32>,
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
- // expected-error @below {{ailed to pad op}}
+ // expected-error @below {{could not compute upper constant bounds for padded dims}}
%padded, %pad, %copy_back = transform.structured.pad %0 {
padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
// Note - attempting to pad non-static dim
>From dc095f41b317faec8f304e7ba38288276f904d52 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fabian.mora-cordero at amd.com>
Date: Mon, 16 Jun 2025 16:41:16 +0000
Subject: [PATCH 2/2] add tests and address reviewer comments
---
.../Linalg/TransformOps/LinalgTransformOps.td | 15 ++--
.../TransformOps/LinalgTransformOps.cpp | 62 ++++++++--------
.../test/Dialect/Linalg/transform-op-pad.mlir | 72 +++++++++++++++++++
3 files changed, 117 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index bf56b633d2872..bd809232d1086 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1110,6 +1110,11 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
* "linalg.copy"
* "none" (no copy back)
+ If `use_prescribed_tensor_shapes` is set to false, the transform will
+ compute the constant upper bounds for each of the padded dimensions
+ for each of the input operands. Otherwise, it will use `tensor.dim`
+ to compute the size of the padded dimension.
+
#### Return modes
This operation ignores non-Linalg ops and drops them in the return.
@@ -1135,7 +1140,7 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
"{}">:$transpose_paddings,
DefaultValuedAttr<StrAttr, "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copy_back_op,
- UnitProp:$use_default_tensor_shapes);
+ DefaultValuedAttr<UnitAttr, "false">:$use_prescribed_tensor_shapes);
let results = (outs TransformHandleTypeInterface:$padded,
TransformHandleTypeInterface:$pad,
TransformHandleTypeInterface:$copy);
@@ -1143,7 +1148,7 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
let assemblyFormat = [{
$target
(`pad_to_multiple_of` custom<DynamicIndexList>($pad_to_multiple_of, $static_pad_to_multiple_of)^)?
- (`use_default_tensor_shapes` $use_default_tensor_shapes^)?
+ (`use_prescribed_tensor_shapes` $use_prescribed_tensor_shapes^)?
attr-dict
`:` functional-type(operands, results)
}];
@@ -1161,13 +1166,15 @@ def PadOp : Op<Transform_Dialect, "structured.pad",
CArg<"ArrayRef<int64_t>", "{}">:$staticPadToMultipleOf,
CArg<"ArrayRef<int64_t>", "{}">:$nofoldFlags,
CArg<"ArrayRef<Attribute>", "{}">:$transposePaddings,
- CArg<"StringRef", "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copyBackOp)>,
+ CArg<"StringRef", "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copyBackOp,
+ CArg<"bool", "false">:$usePrescribedTensorShapes)>,
OpBuilder<(ins "Value":$target,
"ArrayRef<int64_t>":$paddingDimensions,
"ArrayRef<OpFoldResult>":$mixedPadToMultipleOf,
CArg<"ArrayRef<int64_t>", "{}">:$nofoldFlags,
CArg<"ArrayRef<Attribute>", "{}">:$transposePaddings,
- CArg<"StringRef", "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copyBackOp)>
+ CArg<"StringRef", "::mlir::bufferization::MaterializeInDestinationOp::getOperationName()">:$copyBackOp,
+ CArg<"bool", "false">:$usePrescribedTensorShapes)>
];
let extraClassDeclaration = [{
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 45d83b9a8d3c8..f015b84c4f631 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1907,22 +1907,25 @@ void transform::PadOp::build(OpBuilder &b, OperationState &result, Value target,
ArrayRef<int64_t> padToMultipleOf,
ArrayRef<int64_t> nofoldFlags,
ArrayRef<Attribute> transposePaddings,
- StringRef copyBackOp) {
+ StringRef copyBackOp,
+ bool usePrescribedTensorShapes) {
auto resultType = transform::AnyOpType::get(b.getContext());
- return build(/*builder=*/b,
- /*result=*/result,
- /*types=*/TypeRange{resultType, resultType},
- /*target=*/target,
- /*paddingValues=*/ArrayAttr(), // let inference handle this
- /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
- /*padToMultipleOf=*/ValueRange{},
- /*padToMultipleOf=*/
- (padToMultipleOf.empty()
- ? DenseI64ArrayAttr()
- : b.getDenseI64ArrayAttr(padToMultipleOf)),
- /*nofoldFlags=*/b.getI64ArrayAttr(nofoldFlags),
- /*transposePaddings=*/b.getArrayAttr(transposePaddings),
- /*copyBackOp=*/b.getStringAttr(copyBackOp));
+ return build(
+ /*builder=*/b,
+ /*result=*/result,
+ /*types=*/TypeRange{resultType, resultType},
+ /*target=*/target,
+ /*paddingValues=*/ArrayAttr(), // let inference handle this
+ /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
+ /*padToMultipleOf=*/ValueRange{},
+ /*padToMultipleOf=*/
+ (padToMultipleOf.empty() ? DenseI64ArrayAttr()
+ : b.getDenseI64ArrayAttr(padToMultipleOf)),
+ /*nofoldFlags=*/b.getI64ArrayAttr(nofoldFlags),
+ /*transposePaddings=*/b.getArrayAttr(transposePaddings),
+ /*copyBackOp=*/b.getStringAttr(copyBackOp),
+ /*usePrescribedTensorShapes=*/
+ usePrescribedTensorShapes ? b.getUnitAttr() : nullptr);
}
void transform::PadOp::build(OpBuilder &b, OperationState &result, Value target,
@@ -1930,23 +1933,26 @@ void transform::PadOp::build(OpBuilder &b, OperationState &result, Value target,
ArrayRef<OpFoldResult> mixedPadToMultipleOf,
ArrayRef<int64_t> nofoldFlags,
ArrayRef<Attribute> transposePaddings,
- StringRef copyBackOp) {
+ StringRef copyBackOp,
+ bool usePrescribedTensorShapes) {
auto resultType = transform::AnyOpType::get(b.getContext());
SmallVector<int64_t> staticPadToMultipleOf;
SmallVector<Value> dynamicPadToMultipleOf;
dispatchIndexOpFoldResults(mixedPadToMultipleOf, dynamicPadToMultipleOf,
staticPadToMultipleOf);
- return build(/*builder=*/b,
- /*result=*/result,
- /*types=*/TypeRange{resultType, resultType},
- /*target=*/target,
- /*paddingValues=*/ArrayAttr(), // let inference handle this
- /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
- /*padToMultipleOf=*/dynamicPadToMultipleOf,
- /*padToMultipleOf=*/staticPadToMultipleOf,
- /*nofoldFlags=*/b.getI64ArrayAttr(nofoldFlags),
- /*transposePaddings=*/b.getArrayAttr(transposePaddings),
- /*copyBackOp=*/b.getStringAttr(copyBackOp));
+ return build(
+ /*builder=*/b,
+ /*result=*/result,
+ /*types=*/TypeRange{resultType, resultType},
+ /*target=*/target,
+ /*paddingValues=*/ArrayAttr(), // let inference handle this
+ /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
+ /*padToMultipleOf=*/dynamicPadToMultipleOf,
+ /*padToMultipleOf=*/staticPadToMultipleOf,
+ /*nofoldFlags=*/b.getI64ArrayAttr(nofoldFlags),
+ /*transposePaddings=*/b.getArrayAttr(transposePaddings),
+ /*copyBackOp=*/copyBackOp,
+ /*usePrescribedTensorShapes=*/usePrescribedTensorShapes);
}
void PadOp::getEffects(
@@ -2051,7 +2057,7 @@ transform::PadOp::apply(transform::TransformRewriter &rewriter,
} else {
llvm_unreachable("unsupported copy_back op");
}
- if (!getUseDefaultTensorShapes()) {
+ if (!getUsePrescribedTensorShapes()) {
if (failed(options.computeConstantUpperShapeBounds(linalgTarget))) {
auto diag =
emitSilenceableError()
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad.mlir b/mlir/test/Dialect/Linalg/transform-op-pad.mlir
index 926b63311acf6..374d06845be71 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad.mlir
@@ -313,6 +313,39 @@ module attributes {transform.with_named_sequence} {
// -----
+// Test dynamic padding using `use_prescribed_tensor_shapes`
+
+// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (-s1 + (s0 ceildiv 7) * 7)>
+// CHECK: @use_prescribed_tensor_shapes
+// CHECK: (%[[ARG0:.*]]: tensor<?x12xf32>, %[[ARG1:.*]]: tensor<12x?xf32>
+func.func @use_prescribed_tensor_shapes(%arg0: tensor<?x12xf32>,
+ %arg1: tensor<12x?xf32>,
+ %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ // CHECK: %[[C1_0:.*]] = arith.constant 1 : index
+ // CHECK: %[[DIM_0:.*]] = tensor.dim %[[ARG1]], %[[C1_0]] : tensor<12x?xf32>
+ // CHECK: %[[C1_1:.*]] = arith.constant 1 : index
+ // CHECK: %[[DIM_1:.*]] = tensor.dim %[[ARG1]], %[[C1_1]] : tensor<12x?xf32>
+ // CHECK: %[[PADDING:.*]] = affine.apply #[[MAP]]()[%[[DIM_0]], %[[DIM_1]]]
+ // CHECK: %[[PADDED:.*]] = tensor.pad %[[ARG1]] low[0, 0] high[0, %[[PADDING]]] {
+ // CHECK: linalg.matmul ins(%[[ARG0]], %[[PADDED]] : tensor<?x12xf32>, tensor<12x?xf32>)
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x12xf32>, tensor<12x?xf32>) outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ func.return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %padded, %pad, %copy_back = transform.structured.pad %0
+ pad_to_multiple_of [7] use_prescribed_tensor_shapes {
+ padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
+ padding_dimensions=[1]
+ } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
// Check that the padding can be applied even when the output argument of the
// linalg op is not produced by an empty op or an extract_slice op.
@@ -419,3 +452,42 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Test padding when multiple padded dimensions appear in the indexing map
+// result expression, for example:
+// `affine_map<(d0, d1, d2) -> (d0, d2, d0 + d1)>`
+// In this test we pad the first and second dimensions to be multiples of 13 and
+// 17 elements, which results in a padding of `13 * 17 = 221` elements.
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2, d0 + d1)>
+// CHECK-LABEL: @test_lcm_padding
+func.func @test_lcm_padding(%arg0: tensor<2x3xf32>, %arg1: tensor<2x7x11xf32>) -> tensor<2x7x11xf32> {
+ // CHECK: tensor.pad %{{.*}} low[0, 0] high[11, 14] {
+ // CHECK: } : tensor<2x3xf32> to tensor<13x17xf32>
+ // CHECK: tensor.pad %{{.*}} low[0, 0, 0] high[11, 12, 210] {
+ // CHECK: } : tensor<2x7x11xf32> to tensor<13x19x221xf32>
+ // CHECK: linalg.generic {indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"]} ins(%{{.*}} : tensor<13x17xf32>) outs(%{{.*}} : tensor<13x19x221xf32>)
+ %0 = linalg.generic {
+ indexing_maps = [#map, #map1],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ } ins(%arg0 : tensor<2x3xf32>) outs(%arg1 : tensor<2x7x11xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<2x7x11xf32>
+ return %0 : tensor<2x7x11xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %padded, %pad, %copy = transform.structured.pad %0
+ pad_to_multiple_of [13, 17, 19] {
+ copy_back_op = "none",
+ padding_dimensions = [0, 1, 2],
+ padding_values = [0. : f32, 0. : f32, 0. : f32]
+ } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list