[llvm-branch-commits] [mlir] [mlir][transform] Drop redundant padding_dimensions spec from pad_tiling_interface (PR #145257)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Jun 22 21:08:05 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Nicolas Vasilache (nicolasvasilache)
<details>
<summary>Changes</summary>
This revision aligns padding specification in pad_tiling_interface to that of tiling specification.
Dimensions that should be skipped are specified by "padding by 0".
Trailing dimensions that are ignored are automatically completed to "pad to 0".
---
Full diff: https://github.com/llvm/llvm-project/pull/145257.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+17-8)
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+1-20)
- (modified) mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp (+21-24)
- (modified) mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-multiple-of.mlir (+11-17)
- (modified) mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface.mlir (+4-8)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index cf3f2b70580da..c5650470fdc8d 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -1195,17 +1195,29 @@ def PadTilingInterfaceOp : Op<Transform_Dialect, "structured.pad_tiling_interfac
TransformOpInterface,
ReportTrackingListenerFailuresOpTrait]> {
let description = [{
- Pads the operations pointed to by the target handle using the options
- provided as operation attributes. The operation returns a handle to the
- padded operation and to the padding operation ("tensor.pad").
+ Pads the **iteration domain** of the operations pointed to by the target
+ handle using the options provided as operation attributes. Padding the
+ iteration domain induces a padding of the operands that is consistent
+ across the op semantics and, unlike for simple elementwise ops, may not be
+ trivially deducible or specifiable on operands only (e.g. convolutions).
+
+ The specification of `padding_sizes` follows that of `tile_sizes` during
+ tiling: the value "0" on a particular iterator encode "no padding". Like in
+ the case of tiling, an automatic completion by 0 to the operation rank
+ occurs.
+
+ This transformation returns a handle to the padded operation and to the
+ padding operation ("tensor.pad").
TODO: in the future this should be moved out of a specific Linalg
implementation file and into a more general "Structured" file.
#### Return modes
- This operation ignores non-Linalg ops and drops them in the return.
- In the future, this operation will support all TilingInterfaceOps.
+ This operation ignores non-IndexingMapOpInterface ops and drops them in the
+ return. In the future, this operation will support all TilingInterfaceOps
+ for which the contract between iteration domain and operands can be
+ reified.
This operation may produce a definite failure if the padding fails for any
reason.
@@ -1219,7 +1231,6 @@ def PadTilingInterfaceOp : Op<Transform_Dialect, "structured.pad_tiling_interfac
let arguments =
(ins TransformHandleTypeInterface:$target,
DefaultValuedAttr<ArrayAttr, "{}">:$padding_values,
- DefaultValuedAttr<I64ArrayAttr, "{}">:$padding_dimensions,
Variadic<TransformAnyParamTypeOrAnyHandle>:$padding_sizes,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
$static_padding_sizes,
@@ -1245,11 +1256,9 @@ def PadTilingInterfaceOp : Op<Transform_Dialect, "structured.pad_tiling_interfac
// add/mul ring at the moment.
// TODO: support other operations (e.g. min, max etc).
OpBuilder<(ins "Value":$target,
- "ArrayRef<int64_t>":$paddingDimensions,
CArg<"ArrayRef<int64_t>", "{}">:$staticPaddingSizes,
CArg<"bool", "false">:$padToMultipleOf)>,
OpBuilder<(ins "Value":$target,
- "ArrayRef<int64_t>":$paddingDimensions,
"ArrayRef<OpFoldResult>":$mixedPadPaddingSizes,
CArg<"bool", "false">:$usePrescribedTensorShapes)>
];
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 5d55adbf46f36..d9a0ba02f4fe4 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2163,7 +2163,6 @@ LogicalResult transform::PadOp::verify() {
void transform::PadTilingInterfaceOp::build(OpBuilder &b,
OperationState &result,
Value target,
- ArrayRef<int64_t> paddingDimensions,
ArrayRef<int64_t> paddingSizes,
bool padToMultipleOf) {
auto resultType = transform::AnyOpType::get(b.getContext());
@@ -2172,7 +2171,6 @@ void transform::PadTilingInterfaceOp::build(OpBuilder &b,
/*types=*/TypeRange{resultType, resultType},
/*target=*/target,
/*paddingValues=*/ArrayAttr(), // let inference handle this
- /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
/*paddingSizes=*/ValueRange{},
/*paddingSizes=*/
(paddingSizes.empty() ? DenseI64ArrayAttr()
@@ -2183,7 +2181,6 @@ void transform::PadTilingInterfaceOp::build(OpBuilder &b,
void transform::PadTilingInterfaceOp::build(
OpBuilder &b, OperationState &result, Value target,
- ArrayRef<int64_t> paddingDimensions,
ArrayRef<OpFoldResult> mixedPaddingSizes, bool padToMultipleOf) {
auto resultType = transform::AnyOpType::get(b.getContext());
SmallVector<int64_t> staticPaddingSizes;
@@ -2195,7 +2192,6 @@ void transform::PadTilingInterfaceOp::build(
/*types=*/TypeRange{resultType, resultType},
/*target=*/target,
/*paddingValues=*/ArrayAttr(), // let inference handle this
- /*paddingDimensions=*/b.getI64ArrayAttr(paddingDimensions),
/*paddingSizes=*/dynamicPaddingSizes,
/*paddingSizes=*/staticPaddingSizes,
/*usePrescribedTensorShapes=*/padToMultipleOf);
@@ -2277,8 +2273,6 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
TilingInterface paddedOp;
PadTilingInterfaceOptions options;
options.setPaddingValues(paddingValues)
- .setPaddingDimensions(
- extractFromIntegerArrayAttr<int64_t>(getPaddingDimensions()))
.setPaddingSizes(getMixedPaddingSizes())
.setPadToMultipleOf(getPadToMultipleOf());
@@ -2303,20 +2297,7 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
return DiagnosedSilenceableFailure::success();
}
-LogicalResult transform::PadTilingInterfaceOp::verify() {
- SmallVector<int64_t> paddingDimensions =
- extractFromIntegerArrayAttr<int64_t>(getPaddingDimensions());
- if (any_of(paddingDimensions,
- [](int64_t paddingDimension) { return paddingDimension < 0; })) {
- return emitOpError() << "expects padding_dimensions to contain positive "
- "integers, found "
- << getPaddingDimensions();
- }
- if (getMixedPaddingSizes().size() != paddingDimensions.size()) {
- return emitOpError() << "expects as many multiples as padding_dimensions";
- }
- return success();
-}
+LogicalResult transform::PadTilingInterfaceOp::verify() { return success(); }
//===---------------------------------------------------------------------===//
// HoistPadOp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 42dac0776bace..eda3373b4d639 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -32,29 +32,27 @@ using namespace mlir::tensor;
#define DBGSNL() (llvm::dbgs() << "\n")
/// Form a "full-rank" padding specification so that the application is easy.
-static llvm::SmallDenseMap<int64_t, OpFoldResult>
-getDimsToSize(Builder &b, ArrayRef<OpFoldResult> indexingSizes,
- const PadTilingInterfaceOptions &options) {
- llvm::SmallDenseMap<int64_t, OpFoldResult> dimsToSize;
- for (const auto &[paddingDim, paddingSize] :
- llvm::zip_equal(options.paddingDimensions, options.paddingSizes)) {
- dimsToSize[paddingDim] = paddingSize;
- }
+static SmallVector<OpFoldResult>
+getFullRankPaddingSizes(Builder &b, ArrayRef<OpFoldResult> indexingSizes,
+ const PadTilingInterfaceOptions &options) {
+ SmallVector<OpFoldResult> paddingSizes;
// Complete the padding specification to specify all dimensions.
- for (int64_t idx = 0, e = indexingSizes.size(); idx != e; ++idx) {
- if (dimsToSize.find(idx) != dimsToSize.end())
- continue;
- // If a dimension is not specified, either complete with:
+ for (size_t idx = 0, e = indexingSizes.size(); idx != e; ++idx) {
+ // Complete to zero if needed.
+ paddingSizes.push_back(options.paddingSizes.size() > idx
+ ? options.paddingSizes[idx]
+ : b.getIndexAttr(0));
+ // If a dimension is zero (either specified or completed), replace by:
// - 1 if we are padding to the next multiple of.
// - indexingSizes[idx] otherwise
- dimsToSize[idx] =
- options.padToMultipleOf ? b.getIndexAttr(1) : indexingSizes[idx];
- }
- for (int64_t idx = 0, e = indexingSizes.size(); idx != e; ++idx) {
- LLVM_DEBUG(DBGS() << "----idx: " << idx << " : " << dimsToSize[idx]
+ if (isZeroInteger(paddingSizes[idx])) {
+ paddingSizes[idx] =
+ options.padToMultipleOf ? b.getIndexAttr(1) : indexingSizes[idx];
+ }
+ LLVM_DEBUG(DBGS() << "----idx: " << idx << " : " << paddingSizes[idx]
<< "\n");
}
- return dimsToSize;
+ return paddingSizes;
}
/// Compute the padded shape of the given value `v` of `RankedTensorType` given
@@ -80,8 +78,8 @@ SmallVector<OpFoldResult> linalg::computePaddedShape(
"rank");
// "Full-rank" padding specification.
- llvm::SmallDenseMap<int64_t, OpFoldResult> dimsToSize =
- getDimsToSize(rewriter, indexingSizes, options);
+ SmallVector<OpFoldResult> paddingSizes =
+ getFullRankPaddingSizes(rewriter, indexingSizes, options);
// For each dimension in the operand's shape, iterate over indexingSizes and
// add the various term contributions.
@@ -97,7 +95,9 @@ SmallVector<OpFoldResult> linalg::computePaddedShape(
// Find all padding dimensions that contribute to this operand dimension
// and compute the padded term contribution to the final padded shape.
SmallVector<OpFoldResult> terms;
- for (const auto &[paddingDim, paddingSize] : dimsToSize) {
+ for (size_t paddingDim = 0, e = paddingSizes.size(); paddingDim != e;
+ ++paddingDim) {
+ OpFoldResult paddingSize = paddingSizes[paddingDim];
LLVM_DEBUG(DBGS() << "------try apply padding of dim: " << paddingDim
<< " to: " << paddingSize << "\n");
if (!enResults.value().isFunctionOfDim(paddingDim))
@@ -224,9 +224,6 @@ linalg::rewriteAsPaddedOp(RewriterBase &rewriter, TilingInterface opToPad,
SmallVector<tensor::PadOp> &padOps,
PadSizeComputationFunction computePaddingSizeFun) {
LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << opToPad << "\n");
- assert(constOptions.paddingSizes.size() ==
- constOptions.paddingDimensions.size() &&
- "invalid number of elements in padToMultipleOf");
Location loc = opToPad.getLoc();
PadTilingInterfaceOptions options(constOptions);
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-multiple-of.mlir b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-multiple-of.mlir
index 4fcbcbb2a18e3..2bba309953570 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-multiple-of.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-multiple-of.mlir
@@ -36,8 +36,7 @@ module attributes {transform.with_named_sequence} {
// Tile to 5 then pad to 8 (supposedly to better hit vector ops).
%matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%matmul_padded, %_ = transform.structured.pad_tiling_interface %matmul_l1 to padding_sizes [8] pad_to_multiple_of {
- padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
- padding_dimensions=[0]
+ padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
@@ -73,9 +72,8 @@ module {
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 = transform.structured.pad_tiling_interface %0 to padding_sizes [3, 5] pad_to_multiple_of {
- padding_dimensions = [0, 2],
- padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]
+ %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [3, 0, 5] pad_to_multiple_of {
+ padding_values = [0.0 : f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@@ -128,9 +126,8 @@ module {
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 = transform.structured.pad_tiling_interface %0 to padding_sizes [3, 5] pad_to_multiple_of {
- padding_dimensions = [0, 2],
- padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]
+ %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [3, 0, 5] pad_to_multiple_of {
+ padding_values = [0.0 : f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@@ -174,9 +171,8 @@ module attributes {transform.with_named_sequence} {
: (!transform.any_op) -> !transform.any_op
// Pad then tile should produce static shapes.
- %matmul_padded, %_ = transform.structured.pad_tiling_interface %matmul to padding_sizes [8, 16] pad_to_multiple_of {
- padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
- padding_dimensions=[0, 2]
+ %matmul_padded, %_ = transform.structured.pad_tiling_interface %matmul to padding_sizes [8, 0, 16] pad_to_multiple_of {
+ padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%m, %l0, %l1 = transform.structured.tile_using_for %matmul_padded tile_sizes [8, 0, 16]
@@ -223,9 +219,8 @@ module attributes {transform.with_named_sequence} {
%m, %l0, %l1 = transform.structured.tile_using_for %matmul tile_sizes [8, 0, 16]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
- %matmul_padded, %_ = transform.structured.pad_tiling_interface %m to padding_sizes [8, 16] pad_to_multiple_of {
- padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
- padding_dimensions=[0, 2]
+ %matmul_padded, %_ = transform.structured.pad_tiling_interface %m to padding_sizes [8, 0, 16] pad_to_multiple_of {
+ padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
@@ -258,9 +253,8 @@ module attributes {transform.with_named_sequence} {
%m, %l0, %l1 = transform.structured.tile_using_for %matmul tile_sizes [8, 0, 16]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
- %matmul_padded, %_ = transform.structured.pad_tiling_interface %m to padding_sizes [8, 16] pad_to_multiple_of {
- padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
- padding_dimensions=[0, 2]
+ %matmul_padded, %_ = transform.structured.pad_tiling_interface %m to padding_sizes [8, 0, 16] pad_to_multiple_of {
+ padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface.mlir b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface.mlir
index f0a410fa4015f..26c03ed309c05 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface.mlir
@@ -18,8 +18,7 @@ module attributes {transform.with_named_sequence} {
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%fill_padded, %_ = transform.structured.pad_tiling_interface %fill_l1 to padding_sizes [8] {
- padding_values=[0.0 : f32, 0.0 : f32],
- padding_dimensions=[0]
+ padding_values=[0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
@@ -55,8 +54,7 @@ module attributes {transform.with_named_sequence} {
// Tile to 5 then pad to 8 (supposedly to better hit vector ops).
%matmul_l1, %loops_l1 = transform.structured.tile_using_for %matmul tile_sizes [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%matmul_padded, %_ = transform.structured.pad_tiling_interface %matmul_l1 to padding_sizes [8] {
- padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
- padding_dimensions=[0]
+ padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
@@ -91,8 +89,7 @@ module {
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 = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 14] {
- padding_dimensions = [0, 2],
+ %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 0, 14] {
padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
@@ -147,8 +144,7 @@ module {
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 = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 14] {
- padding_dimensions = [0, 2],
+ %padded, %pad = transform.structured.pad_tiling_interface %0 to padding_sizes [8, 0, 14] {
padding_values = [0.000000e+00 : f32, 0.000000e+00 : f32, 0.000000e+00 : f32]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
``````````
</details>
https://github.com/llvm/llvm-project/pull/145257
More information about the llvm-branch-commits
mailing list