[Mlir-commits] [mlir] [[mlir][linalg] Infer reduction-neutral padding values in rewriteAsPaddedOp (PR #216517)
Jianhui Li
llvmlistbot at llvm.org
Tue Sep 15 08:56:26 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/216517
>From 2de712ae5ef333e2914b7407c68e526ae45b4169 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 15 Aug 2026 22:37:08 +0000
Subject: [PATCH 1/6] [mlir][linalg] Infer reduction-neutral padding values in
rewriteAsPaddedOp
When `rewriteAsPaddedOp` is asked to pad an op without explicit
`options.paddingValues`, it previously filled every operand with the zero
of its element type. That is wrong for a reduced operand: its padded
elements are combined into the result, so a zero corrupts e.g. a
`maximumf` or `mulf` reduction.
Factor the inference into a new `linalg::inferPaddingValues` utility that
picks a semantics-preserving value per operand:
- operands with no reduction dimension keep the zero default;
- a contraction keeps zero (0 annihilates through the multiply);
- otherwise the reduced operand is padded with the neutral element of the
reduction combiner (`-inf` for `maximumf`, `1` for `mulf`, ...), found
via `matchReduction` so the actual accumulator combiner is used rather
than any combiner-shaped op in the body.
Inference is conservative and returns failure (rather than risk a silent
miscompile) for a non-LinalgOp reduction, a fused op with multiple
reductions, or a reduced operand that reaches the reduction indirectly.
`inferPaddingValues` is public so callers can use it as an up-front check
to decide whether they must supply `padding_values` explicitly, and the
`transform.structured.pad_tiling_interface` failure diagnostic now points
that out.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Dialect/Linalg/Transforms/Transforms.h | 13 +
.../TransformOps/LinalgTransformOps.cpp | 6 +
.../Linalg/Transforms/PadTilingInterface.cpp | 70 ++++-
...orm-op-pad-tiling-interface-reduction.mlir | 296 ++++++++++++++++++
4 files changed, 379 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 130f7f6fd2792..6b31ee353f37d 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -745,6 +745,19 @@ struct PadTilingInterfaceResult {
SmallVector<Value> replacements;
};
+/// Infers a semantics-preserving padding value for every operand of `toPad`
+/// (indexed by operand number). Operands that are reduced are padded with the
+/// neutral element of their reduction combiner (e.g. `-inf` for `maximumf`, `1`
+/// for `mulf`) because a plain zero would corrupt the reduction; every other
+/// operand is padded with the zero value of its element type.
+///
+/// Inference is conservative: it returns failure when a semantics-preserving
+/// value cannot be determined (a non-LinalgOp reduction, or a reduction whose
+/// neutral element is unknown), letting callers decide whether to set
+/// `options.paddingValues` explicitly.
+FailureOr<SmallVector<Attribute>> inferPaddingValues(OpBuilder &builder,
+ TilingInterface toPad);
+
/// Pad the iterator dimensions of `toPad`.
/// * "options.paddingSizes" indicates that each padding dimension should be
/// padded to the specified padding size.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 8c5bd50bd6f8b..bf63844a7f1a1 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2530,6 +2530,12 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
if (failed(maybePadOps)) {
auto diag = emitSilenceableError() << "failed to pad op";
diag.attachNote(target->getLoc()) << "target op";
+ // With no explicit padding values the transform relies on inference,
+ // which fails when a semantics-preserving value cannot be determined
+ // (e.g. a reduction whose neutral element is unknown).
+ if (paddingValues.empty())
+ diag.attachNote(target->getLoc())
+ << "could not infer a padding value; specify `padding_values`";
return diag;
}
const auto &[paddedOperands, paddedOp, slicedResults] = maybePadOps.value();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index ab2629b41a463..d1b517554b901 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -8,8 +8,11 @@
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/UB/IR/UBMatchers.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
@@ -263,6 +266,61 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
paddingValue, /*nofold=*/false, dynDims);
}
+FailureOr<SmallVector<Attribute>>
+linalg::inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
+ Operation *op = toPad.getOperation();
+
+ // Padding acts on operands; default each to the zero of its element type.
+ SmallVector<Attribute> paddingValues;
+ for (Type t : op->getOperandTypes())
+ paddingValues.push_back(builder.getZeroAttr(getElementTypeOrSelf(t)));
+
+ // No reduction: padded elements are never combined, so zero is always safe.
+ SmallVector<utils::IteratorType> iterTypes = toPad.getLoopIteratorTypes();
+ if (!llvm::is_contained(iterTypes, utils::IteratorType::reduction))
+ return paddingValues;
+
+ // A reduction's neutral element requires inspecting the combiner, which is
+ // only possible for a LinalgOp; fail conservatively otherwise.
+ auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+ if (!linalgOp)
+ return failure();
+
+ // In a contraction, zero annihilates through the multiply (0 * x = 0), so the
+ // default zero is already correct for every operand.
+ if (isaContractionOpInterface(linalgOp))
+ return paddingValues;
+
+ // Only a single reduction has an unambiguous per-operand neutral.
+ if (linalgOp.getNumDpsInits() != 1)
+ return failure();
+ SmallVector<Operation *> combiners;
+ if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combiners))
+ return failure();
+ Operation *combiner = combiners.front();
+ std::optional<TypedAttr> neutral = arith::getNeutralElement(combiner);
+ if (!neutral)
+ return failure();
+
+ for (OpOperand *input : linalgOp.getDpsInputOperands()) {
+ // Only operands indexed along a reduction dim need a neutral.
+ AffineMap map = linalgOp.getMatchingIndexingMap(input);
+ bool reduced = llvm::any_of(llvm::enumerate(iterTypes), [&](auto it) {
+ return it.value() == utils::IteratorType::reduction &&
+ map.isFunctionOfDim(it.index());
+ });
+ if (!reduced)
+ continue;
+ // A reduced operand must feed the combiner directly to use its neutral; an
+ // indirect one (e.g. via a math.exp) has no valid pad value -> fail.
+ if (!llvm::is_contained(linalgOp.getMatchingBlockArgument(input).getUsers(),
+ combiner))
+ return failure();
+ paddingValues[input->getOperandNumber()] = *neutral;
+ }
+ return paddingValues;
+}
+
FailureOr<PadTilingInterfaceResult> linalg::rewriteAsPaddedOp(
OpBuilder &builder, TilingInterface toPad,
PadTilingInterfaceOptions options,
@@ -272,14 +330,14 @@ FailureOr<PadTilingInterfaceResult> linalg::rewriteAsPaddedOp(
Location loc = toPad.getLoc();
// Allow inference of pad values if they are not explicitly specified.
- // TODO: be mindful about the value depending on the actual operation.
if (options.paddingValues.empty()) {
- SmallVector<Type> types(toPad->getOperandTypes());
- llvm::append_range(types, toPad->getResultTypes());
- for (Type t : types) {
- options.paddingValues.push_back(
- builder.getZeroAttr(getElementTypeOrSelf(t)));
+ FailureOr<SmallVector<Attribute>> inferred =
+ inferPaddingValues(builder, toPad);
+ if (failed(inferred)) {
+ LLVM_DEBUG(DBGS() << "Could not infer pad values: FAIL\n");
+ return failure();
}
+ options.paddingValues = std::move(*inferred);
}
if (llvm::any_of(toPad->getOperands(),
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
new file mode 100644
index 0000000000000..f3b0c52f40608
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
@@ -0,0 +1,296 @@
+// RUN: mlir-opt --transform-interpreter -canonicalize -split-input-file --verify-diagnostics %s | FileCheck %s
+
+// When padding values are not specified, `pad_tiling_interface` infers them.
+// For an operand that is reduced, a zero pad value would corrupt the result,
+// so the neutral element of the reduction combiner must be inferred instead.
+
+// CHECK-LABEL: @pad_reduce_maximumf
+func.func @pad_reduce_maximumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // maximumf neutral element is -inf.
+ // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 0xFF800000 : f32
+ // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.maximumf %in, %out : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // No padding_values: they are inferred.
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// An operand reduced along several dimensions is padded with the neutral on all
+// of them (here d1 and d2); one neutral suffices since the combiner is the same.
+
+#in3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#out3 = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK-LABEL: @pad_reduce_two_reduction_dims
+func.func @pad_reduce_two_reduction_dims(%in: tensor<8x6x10xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // CHECK-DAG: %[[NINF:.*]] = arith.constant 0xFF800000 : f32
+ // CHECK: tensor.pad %{{.*}} low[0, 0, 0] high[0, 2, 6]
+ // CHECK: tensor.yield %[[NINF]] : f32
+ // CHECK: } : tensor<8x6x10xf32> to tensor<8x8x16xf32>
+ %0 = linalg.generic {indexing_maps = [#in3, #out3],
+ iterator_types = ["parallel", "reduction", "reduction"]}
+ ins(%in : tensor<8x6x10xf32>) outs(%init : tensor<8xf32>) {
+ ^bb0(%a: f32, %o: f32):
+ %m = arith.maximumf %a, %o : f32
+ linalg.yield %m : f32
+ } -> tensor<8xf32>
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %gen to padding_sizes [8, 8, 16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @pad_reduce_mulf
+func.func @pad_reduce_mulf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // mulf neutral element is 1.0 (a zero pad would zero out the product).
+ // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 1.000000e+00 : f32
+ // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.mulf %in, %out : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: @pad_reduce_minimumf
+func.func @pad_reduce_minimumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // minimumf neutral element is +inf.
+ // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 0x7F800000 : f32
+ // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.minimumf %in, %out : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// maxnumf is the NaN-ignoring variant, so its neutral element is NaN. Padding
+// with NaN is correct here (maxnumf(NaN, x) = x), but note it would be unsound
+// under a `nnan` fast-math assumption.
+
+// CHECK-LABEL: @pad_reduce_maxnumf
+func.func @pad_reduce_maxnumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // CHECK-DAG: %[[NAN:.*]] = arith.constant 0xFFC00000 : f32
+ // CHECK: tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NAN]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>) dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.maxnumf %in, %out : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// A contraction is padded with zero: its reduced operands feed the multiply and
+// zero annihilates it (0 * x = 0), so the add-reduction stays correct. The
+// combiner (arith.addf) must NOT be mistaken for the operands' direct combiner
+// (arith.mulf, whose neutral 1 would corrupt the result).
+
+// CHECK-LABEL: @pad_matmul_uses_zero
+func.func @pad_matmul_uses_zero(%A: tensor<8x10xf32>, %B: tensor<10x8xf32>,
+ %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+ // CHECK: tensor.pad %{{.*}} low[0, 0] high[0, 6]
+ // CHECK: tensor.yield %[[ZERO]] : f32
+ // CHECK: } : tensor<8x10xf32> to tensor<8x16xf32>
+ // CHECK: tensor.pad %{{.*}} low[0, 0] high[6, 0]
+ // CHECK: tensor.yield %[[ZERO]] : f32
+ // CHECK: } : tensor<10x8xf32> to tensor<16x8xf32>
+ // CHECK-NOT: arith.constant 1.000000e+00
+ %0 = linalg.matmul ins(%A, %B : tensor<8x10xf32>, tensor<10x8xf32>)
+ outs(%C : tensor<8x8xf32>) -> tensor<8x8xf32>
+ return %0 : tensor<8x8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %mm = transform.structured.match ops{["linalg.matmul"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %mm to padding_sizes [8, 8, 16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Only the reduced input takes the neutral (-inf); the non-reduced init operand
+// keeps zero. Both dims are padded so the init is padded along the parallel dim.
+
+// CHECK-LABEL: @pad_reduce_init_stays_zero
+func.func @pad_reduce_init_stays_zero(%in: tensor<8x30xf32>, %init: tensor<8xf32>)
+ -> tensor<8xf32> {
+ // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+ // CHECK-DAG: %[[NINF:.*]] = arith.constant 0xFF800000 : f32
+ // Reduced input -> maximumf neutral (-inf).
+ // CHECK: tensor.pad %{{.*}}
+ // CHECK: tensor.yield %[[NINF]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<16x32xf32>
+ // Init (non-reduced) operand -> plain zero.
+ // CHECK: tensor.pad %{{.*}}
+ // CHECK: tensor.yield %[[ZERO]] : f32
+ // CHECK: } : tensor<8xf32> to tensor<16xf32>
+ %0 = linalg.reduce ins(%in : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%a: f32, %o: f32) {
+ %m = arith.maximumf %a, %o : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [16, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Negative: the reduced operand is not consumed directly by a recognized
+// reduction combiner (it flows through math.exp first), so no safe neutral
+// element can be inferred and padding fails conservatively rather than
+// silently padding with a wrong value.
+
+func.func @pad_reduce_indirect_combiner_fails(
+ %input: tensor<8x30xf32>, %init: tensor<8xf32>) -> tensor<8xf32> {
+ // expected-note @below {{target op}}
+ // expected-note @below {{could not infer a padding value; specify `padding_values`}}
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %e = math.exp %in : f32
+ %a = arith.addf %e, %out : f32
+ linalg.yield %a : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // expected-error @below {{failed to pad op}}
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Negative: a fused op with two reductions (max and sum) over the same input
+// has no single correct neutral for that input (-inf breaks the sum, 0 breaks
+// the max), so inference fails instead of padding with a wrong value.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#mapr = affine_map<(d0, d1) -> (d0)>
+func.func @pad_multi_reduction_fails(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
+ %s0: tensor<8xf32>)
+ -> (tensor<8xf32>, tensor<8xf32>) {
+ // expected-note @below {{target op}}
+ // expected-note @below {{could not infer a padding value; specify `padding_values`}}
+ %r:2 = linalg.generic {indexing_maps = [#map, #mapr, #mapr],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%in : tensor<8x30xf32>) outs(%m0, %s0 : tensor<8xf32>, tensor<8xf32>) {
+ ^bb0(%x: f32, %mo: f32, %so: f32):
+ %mx = arith.maximumf %x, %mo : f32
+ %sm = arith.addf %x, %so : f32
+ linalg.yield %mx, %sm : f32, f32
+ } -> (tensor<8xf32>, tensor<8xf32>)
+ return %r#0, %r#1 : tensor<8xf32>, tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // expected-error @below {{failed to pad op}}
+ %padded, %pad = transform.structured.pad_tiling_interface %gen to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
>From 95463426280b2607fd77077b5a6be6e8afe11e1d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 28 Aug 2026 00:37:41 +0000
Subject: [PATCH 2/6] [mlir][linalg] Address review comments on padding value
inference
- Make `inferPaddingValues` static and move its docs onto the definition:
it is an implementation detail of `rewriteAsPaddedOp` with no in-tree
caller needing it exposed.
- Replace `isaContractionOpInterface` with a direct body match against the
op pairs for which zero is a valid pad value (`0 * x = 0`, `0 + x = x`).
The interface is meant to accept more bodies over time, and one whose
accumulation is not unchanged by zero -- `max`/`+` -- would silently get
a wrong pad value. This also fixes convolutions, which the interface
rejects for having non-projected-permutation maps and which therefore
failed inference outright.
- Reject a NaN neutral element under `fastmath<nnan>`:
`arith::getNeutralElement` accounts for `ninf` but not `nnan`, so
`maxnumf`/`minnumf` would otherwise be padded with a value the flag
declares impossible.
- Drop the inference note from the `pad_tiling_interface` diagnostic;
`rewriteAsPaddedOp` fails for other reasons too, so it could misattribute
the cause.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Dialect/Linalg/Transforms/Transforms.h | 13 -----
.../TransformOps/LinalgTransformOps.cpp | 6 ---
.../Linalg/Transforms/PadTilingInterface.cpp | 50 +++++++++++++++++--
...orm-op-pad-tiling-interface-reduction.mlir | 34 +++++++++++--
4 files changed, 74 insertions(+), 29 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index f00d5d977576c..9ab50f1136ac1 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -745,19 +745,6 @@ struct PadTilingInterfaceResult {
SmallVector<Value> replacements;
};
-/// Infers a semantics-preserving padding value for every operand of `toPad`
-/// (indexed by operand number). Operands that are reduced are padded with the
-/// neutral element of their reduction combiner (e.g. `-inf` for `maximumf`, `1`
-/// for `mulf`) because a plain zero would corrupt the reduction; every other
-/// operand is padded with the zero value of its element type.
-///
-/// Inference is conservative: it returns failure when a semantics-preserving
-/// value cannot be determined (a non-LinalgOp reduction, or a reduction whose
-/// neutral element is unknown), letting callers decide whether to set
-/// `options.paddingValues` explicitly.
-FailureOr<SmallVector<Attribute>> inferPaddingValues(OpBuilder &builder,
- TilingInterface toPad);
-
/// Pad the iterator dimensions of `toPad`.
/// * "options.paddingSizes" indicates that each padding dimension should be
/// padded to the specified padding size.
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 60e29adc6b2de..8299e1c2edb41 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2539,12 +2539,6 @@ transform::PadTilingInterfaceOp::apply(transform::TransformRewriter &rewriter,
if (failed(maybePadOps)) {
auto diag = emitSilenceableError() << "failed to pad op";
diag.attachNote(target->getLoc()) << "target op";
- // With no explicit padding values the transform relies on inference,
- // which fails when a semantics-preserving value cannot be determined
- // (e.g. a reduction whose neutral element is unknown).
- if (paddingValues.empty())
- diag.attachNote(target->getLoc())
- << "could not infer a padding value; specify `padding_values`";
return diag;
}
const auto &[paddedOperands, paddedOp, slicedResults] = maybePadOps.value();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index d1b517554b901..5659be5afe0a1 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -266,8 +266,38 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
paddingValue, /*nofold=*/false, dynDims);
}
-FailureOr<SmallVector<Attribute>>
-linalg::inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
+/// Returns true if `mul` and `add` are a pair for which `0 * x = 0` and
+/// `0 + x = x`. Zero is then a valid padding value for the operands of a
+/// contraction with such a body: `mul` turns a padded zero into a zero, and
+/// adding that zero leaves the accumulated value unchanged.
+///
+/// Do not grow this list just to match `isaContractionOpInterface`, since not
+/// every contraction can be zero-padded: a `max`/`+` body needs `-inf`.
+static bool isZeroNeutralContractionPair(Operation *mul, Operation *add) {
+ if (isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add))
+ return true;
+ if (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add))
+ return true;
+ if (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add))
+ return true;
+ if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
+ mul->getResult(0).getType().isInteger(1))
+ return true;
+ return false;
+}
+
+/// Infers a semantics-preserving padding value for every operand of `toPad`
+/// (indexed by operand number). Operands that are reduced are padded with the
+/// neutral element of their reduction combiner (e.g. `-inf` for `maximumf`, `1`
+/// for `mulf`); every other operand is padded with the zero value of its
+/// element type.
+///
+/// Inference is conservative: it returns failure when a semantics-preserving
+/// value cannot be determined (a non-LinalgOp reduction, or a reduction whose
+/// neutral element is unknown), letting callers set `options.paddingValues`
+/// explicitly instead.
+static FailureOr<SmallVector<Attribute>>
+inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
Operation *op = toPad.getOperation();
// Padding acts on operands; default each to the zero of its element type.
@@ -286,9 +316,10 @@ linalg::inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
if (!linalgOp)
return failure();
- // In a contraction, zero annihilates through the multiply (0 * x = 0), so the
- // default zero is already correct for every operand.
- if (isaContractionOpInterface(linalgOp))
+ // Keep the default zero for a contraction: the multiply maps a padded zero to
+ // zero, and accumulating a zero changes nothing.
+ if (linalg::detail::isContractionBody(*linalgOp.getBlock(),
+ isZeroNeutralContractionPair))
return paddingValues;
// Only a single reduction has an unambiguous per-operand neutral.
@@ -302,6 +333,15 @@ linalg::inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
if (!neutral)
return failure();
+ if (auto floatNeutral = dyn_cast<FloatAttr>(*neutral);
+ floatNeutral && floatNeutral.getValue().isNaN()) {
+ auto fastMath = dyn_cast<arith::ArithFastMathInterface>(combiner);
+ if (fastMath &&
+ arith::bitEnumContainsAny(fastMath.getFastMathFlagsAttr().getValue(),
+ arith::FastMathFlags::nnan))
+ return failure();
+ }
+
for (OpOperand *input : linalgOp.getDpsInputOperands()) {
// Only operands indexed along a reduction dim need a neutral.
AffineMap map = linalgOp.getMatchingIndexingMap(input);
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
index f3b0c52f40608..a56d25545e6e2 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
@@ -129,9 +129,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// maxnumf is the NaN-ignoring variant, so its neutral element is NaN. Padding
-// with NaN is correct here (maxnumf(NaN, x) = x), but note it would be unsound
-// under a `nnan` fast-math assumption.
+// maxnumf is the NaN-ignoring variant, so its neutral element is NaN.
// CHECK-LABEL: @pad_reduce_maxnumf
func.func @pad_reduce_maxnumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
@@ -160,6 +158,34 @@ module attributes {transform.with_named_sequence} {
// -----
+// Negative: `nnan` declares that NaN operands do not occur, so the NaN neutral
+// element of maxnumf is not a usable pad value and inference fails instead.
+
+func.func @pad_reduce_maxnumf_nnan_fails(%input: tensor<8x30xf32>,
+ %init: tensor<8xf32>) -> tensor<8xf32> {
+ // expected-note @below {{target op}}
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.maxnumf %in, %out fastmath<nnan> : f32
+ linalg.yield %m : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // expected-error @below {{failed to pad op}}
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
// A contraction is padded with zero: its reduced operands feed the multiply and
// zero annihilates it (0 * x = 0), so the add-reduction stays correct. The
// combiner (arith.addf) must NOT be mistaken for the operands' direct combiner
@@ -238,7 +264,6 @@ module attributes {transform.with_named_sequence} {
func.func @pad_reduce_indirect_combiner_fails(
%input: tensor<8x30xf32>, %init: tensor<8xf32>) -> tensor<8xf32> {
// expected-note @below {{target op}}
- // expected-note @below {{could not infer a padding value; specify `padding_values`}}
%0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
@@ -272,7 +297,6 @@ func.func @pad_multi_reduction_fails(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
%s0: tensor<8xf32>)
-> (tensor<8xf32>, tensor<8xf32>) {
// expected-note @below {{target op}}
- // expected-note @below {{could not infer a padding value; specify `padding_values`}}
%r:2 = linalg.generic {indexing_maps = [#map, #mapr, #mapr],
iterator_types = ["parallel", "reduction"]}
ins(%in : tensor<8x30xf32>) outs(%m0, %s0 : tensor<8xf32>, tensor<8xf32>) {
>From e3bd3daf829bba90f0fe58f64e98e53defb3bda3 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 9 Sep 2026 20:23:16 +0000
Subject: [PATCH 3/6] [mlir][linalg] Derive contraction padding value from the
matched body ops
Restructure the contraction path of `inferPaddingValues` so that the
`mul`/`add` pair produces the padding value instead of only answering
whether zero is valid. `isContractionBody` is now used as an extractor:
it validates the body wiring and hands back the two ops, from which
`setContractionPaddingValues` derives a value per reduced input. This
removes the standalone pair predicate that duplicated the one in
Specialize.cpp, and lets bodies whose neutral element is not zero (e.g.
a max-plus contraction, which needs `-inf`) be supported by adding a
single case.
Both paths now share `isReducedOperand`, and the contraction value is
assigned per operand rather than relying on the default zero, so a
non-zero value cannot land on the init operand.
Also:
- Bail out when `matchReduction` reports more than one combiner rather
than assuming `combiners.front()`.
- Document the `nnan` guard on the NaN neutral of `maxnumf`/`minnumf`.
- Pass `/*redPos=*/0` explicitly to `matchReduction`.
- Prefix the negative tests with `negative_` per the testing guide.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Linalg/Transforms/PadTilingInterface.cpp | 101 ++++++++++++------
...orm-op-pad-tiling-interface-reduction.mlir | 6 +-
2 files changed, 74 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 5659be5afe0a1..5a5877e6804bc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -266,24 +266,52 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
paddingValue, /*nofold=*/false, dynDims);
}
-/// Returns true if `mul` and `add` are a pair for which `0 * x = 0` and
-/// `0 + x = x`. Zero is then a valid padding value for the operands of a
-/// contraction with such a body: `mul` turns a padded zero into a zero, and
-/// adding that zero leaves the accumulated value unchanged.
-///
-/// Do not grow this list just to match `isaContractionOpInterface`, since not
-/// every contraction can be zero-padded: a `max`/`+` body needs `-inf`.
-static bool isZeroNeutralContractionPair(Operation *mul, Operation *add) {
- if (isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add))
- return true;
- if (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add))
- return true;
- if (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add))
- return true;
- if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
- mul->getResult(0).getType().isInteger(1))
- return true;
- return false;
+/// Returns true if `operand` is indexed along at least one reduction dimension
+/// of `linalgOp`
+static bool isReducedOperand(linalg::LinalgOp linalgOp, OpOperand *operand,
+ ArrayRef<utils::IteratorType> iterTypes) {
+ AffineMap map = linalgOp.getMatchingIndexingMap(operand);
+ return llvm::any_of(llvm::enumerate(iterTypes), [&](auto it) {
+ return it.value() == utils::IteratorType::reduction &&
+ map.isFunctionOfDim(it.index());
+ });
+}
+
+/// Set `paddingValues` with the padding value of every input of `contractOp`
+/// that is indexed along a reduction dimension. Fails if the pair admits no
+/// padding value.
+static LogicalResult
+setContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp contractOp,
+ Operation *mul, Operation *add,
+ ArrayRef<utils::IteratorType> iterTypes,
+ MutableArrayRef<Attribute> paddingValues) {
+ // Identify the padding value, `p`, for which `mul(p, x)` is the neutral
+ // element of `add` for every `x`.
+ // Not every contraction can be zero-padded: a body with `mul` = `addf` and
+ // `add` = `maximumf` (a max-plus contraction) pads with `-inf` instead.
+ auto getPadValue = [&](Type elementType) -> Attribute {
+ // `0 * x = 0` and `0 + acc = acc`.
+ if ((isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add)) ||
+ (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add)) ||
+ (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add)))
+ return builder.getZeroAttr(elementType);
+ // `false & x = false` and `false | acc = acc`.
+ if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
+ mul->getResult(0).getType().isInteger(1))
+ return builder.getZeroAttr(elementType);
+ return {};
+ };
+
+ for (OpOperand *input : contractOp.getDpsInputOperands()) {
+ if (!isReducedOperand(contractOp, input, iterTypes))
+ continue;
+ Attribute padValue =
+ getPadValue(getElementTypeOrSelf(input->get().getType()));
+ if (!padValue)
+ return failure();
+ paddingValues[input->getOperandNumber()] = padValue;
+ }
+ return success();
}
/// Infers a semantics-preserving padding value for every operand of `toPad`
@@ -316,23 +344,42 @@ inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
if (!linalgOp)
return failure();
- // Keep the default zero for a contraction: the multiply maps a padded zero to
- // zero, and accumulating a zero changes nothing.
- if (linalg::detail::isContractionBody(*linalgOp.getBlock(),
- isZeroNeutralContractionPair))
+ // A contraction-like body accumulates as `acc = add(mul(a, b), acc)`. Use
+ // `isContractionBody` as an extractor: it hands back these two ops, from
+ // which the padding value is derived.
+ Operation *mulOp = nullptr, *addOp = nullptr;
+ auto captureBodyOps = [&](Operation *mul, Operation *add) {
+ mulOp = mul;
+ addOp = add;
+ return true;
+ };
+ if (linalg::detail::isContractionBody(*linalgOp.getBlock(), captureBodyOps)) {
+ if (failed(setContractionPaddingValues(builder, linalgOp, mulOp, addOp,
+ iterTypes, paddingValues)))
+ return failure();
return paddingValues;
+ }
// Only a single reduction has an unambiguous per-operand neutral.
if (linalgOp.getNumDpsInits() != 1)
return failure();
+
+ // For reduction-like bodies, the neutral element of the combiner is used to
+ // pad the reduced operands.
SmallVector<Operation *> combiners;
- if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combiners))
+ if (!matchReduction(linalgOp.getRegionOutputArgs(), /*redPos=*/0, combiners))
+ return failure();
+ if (combiners.size() != 1)
return failure();
Operation *combiner = combiners.front();
+
std::optional<TypedAttr> neutral = arith::getNeutralElement(combiner);
if (!neutral)
return failure();
+ // In case fastMath is enabled with `nnan`, the neutral element of
+ // `maxnumf`/`minnumf`, NaN, becomes an invalid padding value. Bail out and
+ // let the caller supply a value explicitly.
if (auto floatNeutral = dyn_cast<FloatAttr>(*neutral);
floatNeutral && floatNeutral.getValue().isNaN()) {
auto fastMath = dyn_cast<arith::ArithFastMathInterface>(combiner);
@@ -343,13 +390,7 @@ inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
}
for (OpOperand *input : linalgOp.getDpsInputOperands()) {
- // Only operands indexed along a reduction dim need a neutral.
- AffineMap map = linalgOp.getMatchingIndexingMap(input);
- bool reduced = llvm::any_of(llvm::enumerate(iterTypes), [&](auto it) {
- return it.value() == utils::IteratorType::reduction &&
- map.isFunctionOfDim(it.index());
- });
- if (!reduced)
+ if (!isReducedOperand(linalgOp, input, iterTypes))
continue;
// A reduced operand must feed the combiner directly to use its neutral; an
// indirect one (e.g. via a math.exp) has no valid pad value -> fail.
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
index a56d25545e6e2..60837ffe9a617 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
@@ -161,7 +161,7 @@ module attributes {transform.with_named_sequence} {
// Negative: `nnan` declares that NaN operands do not occur, so the NaN neutral
// element of maxnumf is not a usable pad value and inference fails instead.
-func.func @pad_reduce_maxnumf_nnan_fails(%input: tensor<8x30xf32>,
+func.func @negative_pad_reduce_maxnumf_nnan(%input: tensor<8x30xf32>,
%init: tensor<8xf32>) -> tensor<8xf32> {
// expected-note @below {{target op}}
%0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
@@ -261,7 +261,7 @@ module attributes {transform.with_named_sequence} {
// element can be inferred and padding fails conservatively rather than
// silently padding with a wrong value.
-func.func @pad_reduce_indirect_combiner_fails(
+func.func @negative_pad_reduce_indirect_combiner(
%input: tensor<8x30xf32>, %init: tensor<8xf32>) -> tensor<8xf32> {
// expected-note @below {{target op}}
%0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
@@ -293,7 +293,7 @@ module attributes {transform.with_named_sequence} {
#map = affine_map<(d0, d1) -> (d0, d1)>
#mapr = affine_map<(d0, d1) -> (d0)>
-func.func @pad_multi_reduction_fails(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
+func.func @negative_pad_multi_reduction(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
%s0: tensor<8xf32>)
-> (tensor<8xf32>, tensor<8xf32>) {
// expected-note @below {{target op}}
>From 49a6f9e1d8cfb8503d8b954f7fd792bc0118b540 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 11 Sep 2026 16:05:01 +0000
Subject: [PATCH 4/6] [mlir][linalg] Extract the body ops inside padding value
inference
`setContractionPaddingValues` took `mul` and `add` alongside the op they
came from, an invariant the signature could not express. Rename it to
`inferContractionPaddingValues` and have it match the body itself, so it
takes only the op. It now returns the padding values rather than mutating
a `MutableArrayRef`, which also avoids leaving a partially updated vector
behind on the failure path.
The caller no longer distinguishes "body is not contraction-like" from
"pair admits no padding value": both fail and fall through to the plain
reduction path. That is safe because a contraction-like body routes its
inputs through the elemwise op rather than into the combiner, so the
"reduced operand must feed the combiner directly" check rejects it there.
A negative test covering this is added in the following commit.
Also rename `mul`/`add` to `elemwise`/`reduce`, matching the vocabulary
of `isContractionBody`, since either position can hold any op - a
max-plus contraction has `arith.addf` in the elemwise position.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Linalg/Transforms/PadTilingInterface.cpp | 79 +++++++++++--------
1 file changed, 44 insertions(+), 35 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 5a5877e6804bc..7a9dfe2ef32b7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -277,33 +277,50 @@ static bool isReducedOperand(linalg::LinalgOp linalgOp, OpOperand *operand,
});
}
-/// Set `paddingValues` with the padding value of every input of `contractOp`
-/// that is indexed along a reduction dimension. Fails if the pair admits no
-/// padding value.
-static LogicalResult
-setContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp contractOp,
- Operation *mul, Operation *add,
- ArrayRef<utils::IteratorType> iterTypes,
- MutableArrayRef<Attribute> paddingValues) {
- // Identify the padding value, `p`, for which `mul(p, x)` is the neutral
- // element of `add` for every `x`.
- // Not every contraction can be zero-padded: a body with `mul` = `addf` and
- // `add` = `maximumf` (a max-plus contraction) pads with `-inf` instead.
+/// Returns `defaults` with the entry of every input of `linalgOp` that is
+/// indexed along a reduction dimension replaced by its padding value. Fails if
+/// the body of `linalgOp` is not contraction-like, or if its
+/// `elemwise`/`reduce` pair admits no padding value.
+static FailureOr<SmallVector<Attribute>>
+inferContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp linalgOp,
+ ArrayRef<utils::IteratorType> iterTypes,
+ ArrayRef<Attribute> defaults) {
+ // NOTE: For contraction-like Ops, we infer the padding value by looking at
+ // both elemwise and reduce Ops, where (see isContractionBody for details):
+ // %0 = <elemwise>(permutation-of(cu(block-argument-0),
+ // cu(block-argument-1)))
+ // %1 = <reduce>(permutation-of(cu(%0), cu(block-argument-2)))
+ // return-like cu(%1)
+ Operation *elemwise = nullptr, *reduce = nullptr;
+ auto captureBodyOps = [&](Operation *e, Operation *r) {
+ elemwise = e;
+ reduce = r;
+ return true;
+ };
+ if (!linalg::detail::isContractionBody(*linalgOp.getBlock(), captureBodyOps))
+ return failure();
+
+ // Identify the padding value, `p`, for which `elemwise(p, x)` is the neutral
+ // element of `reduce` for every `x`.
+ // Not every contraction can be zero-padded: a body with `elemwise` =
+ // `arith.addf` and `reduce` = `arith.maximumf` (a max-plus contraction) pads
+ // with `-inf` instead.
auto getPadValue = [&](Type elementType) -> Attribute {
// `0 * x = 0` and `0 + acc = acc`.
- if ((isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add)) ||
- (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add)) ||
- (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add)))
+ if ((isa<arith::MulFOp>(elemwise) && isa<arith::AddFOp>(reduce)) ||
+ (isa<arith::MulIOp>(elemwise) && isa<arith::AddIOp>(reduce)) ||
+ (isa<complex::MulOp>(elemwise) && isa<complex::AddOp>(reduce)))
return builder.getZeroAttr(elementType);
// `false & x = false` and `false | acc = acc`.
- if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
- mul->getResult(0).getType().isInteger(1))
+ if (isa<arith::AndIOp>(elemwise) && isa<arith::OrIOp>(reduce) &&
+ elemwise->getResult(0).getType().isInteger(1))
return builder.getZeroAttr(elementType);
return {};
};
- for (OpOperand *input : contractOp.getDpsInputOperands()) {
- if (!isReducedOperand(contractOp, input, iterTypes))
+ SmallVector<Attribute> paddingValues(defaults.begin(), defaults.end());
+ for (OpOperand *input : linalgOp.getDpsInputOperands()) {
+ if (!isReducedOperand(linalgOp, input, iterTypes))
continue;
Attribute padValue =
getPadValue(getElementTypeOrSelf(input->get().getType()));
@@ -311,7 +328,7 @@ setContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp contractOp,
return failure();
paddingValues[input->getOperandNumber()] = padValue;
}
- return success();
+ return paddingValues;
}
/// Infers a semantics-preserving padding value for every operand of `toPad`
@@ -344,21 +361,13 @@ inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
if (!linalgOp)
return failure();
- // A contraction-like body accumulates as `acc = add(mul(a, b), acc)`. Use
- // `isContractionBody` as an extractor: it hands back these two ops, from
- // which the padding value is derived.
- Operation *mulOp = nullptr, *addOp = nullptr;
- auto captureBodyOps = [&](Operation *mul, Operation *add) {
- mulOp = mul;
- addOp = add;
- return true;
- };
- if (linalg::detail::isContractionBody(*linalgOp.getBlock(), captureBodyOps)) {
- if (failed(setContractionPaddingValues(builder, linalgOp, mulOp, addOp,
- iterTypes, paddingValues)))
- return failure();
- return paddingValues;
- }
+ // A contraction-like body accumulates as `acc = reduce(elemwise(a, b), acc)`
+ // and pads with a value derived from that pair.
+ FailureOr<SmallVector<Attribute>> contractionValues =
+ inferContractionPaddingValues(builder, linalgOp, iterTypes,
+ paddingValues);
+ if (succeeded(contractionValues))
+ return *contractionValues;
// Only a single reduction has an unambiguous per-operand neutral.
if (linalgOp.getNumDpsInits() != 1)
>From dd30326c140ac57219c35572e920af548501847d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 11 Sep 2026 16:05:13 +0000
Subject: [PATCH 5/6] [mlir][linalg] Reorganize the padding value inference
tests
Rename transform-op-pad-tiling-interface-reduction.mlir to
pad-neutral.mlir and group the tests by the op under test
(linalg.reduce, linalg.matmul, linalg.generic), with block comments per
the convention in test/Dialect/Linalg/vectorization/linalg-ops.mlir.
Within each group the positive tests come first, so the negative ones
sit together and the covered edge cases are easier to see. Test names
now all carry the op they exercise.
Existing test bodies are unchanged; only their order and names differ.
Also add negative_pad_generic_max_accumulate: a body that is
contraction-like but accumulates with `arith.maximumf`, for which no
padding value of the inputs maps to the neutral element of the
accumulate. It is the only test covering a contraction-like body whose
elemwise/reduce pair admits no padding value.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
...erface-reduction.mlir => pad-neutral.mlir} | 233 +++++++++++-------
1 file changed, 140 insertions(+), 93 deletions(-)
rename mlir/test/Dialect/Linalg/{transform-op-pad-tiling-interface-reduction.mlir => pad-neutral.mlir} (85%)
diff --git a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir b/mlir/test/Dialect/Linalg/pad-neutral.mlir
similarity index 85%
rename from mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
rename to mlir/test/Dialect/Linalg/pad-neutral.mlir
index 60837ffe9a617..7363c104ea07e 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pad-tiling-interface-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/pad-neutral.mlir
@@ -4,6 +4,10 @@
// For an operand that is reduced, a zero pad value would corrupt the result,
// so the neutral element of the reduction combiner must be inferred instead.
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.reduce
+///--------------------------------------------------------------------------------------
+
// CHECK-LABEL: @pad_reduce_maximumf
func.func @pad_reduce_maximumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
-> tensor<8xf32> {
@@ -35,33 +39,29 @@ module attributes {transform.with_named_sequence} {
// -----
-// An operand reduced along several dimensions is padded with the neutral on all
-// of them (here d1 and d2); one neutral suffices since the combiner is the same.
-
-#in3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-#out3 = affine_map<(d0, d1, d2) -> (d0)>
-// CHECK-LABEL: @pad_reduce_two_reduction_dims
-func.func @pad_reduce_two_reduction_dims(%in: tensor<8x6x10xf32>, %init: tensor<8xf32>)
+// CHECK-LABEL: @pad_reduce_minimumf
+func.func @pad_reduce_minimumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
-> tensor<8xf32> {
- // CHECK-DAG: %[[NINF:.*]] = arith.constant 0xFF800000 : f32
- // CHECK: tensor.pad %{{.*}} low[0, 0, 0] high[0, 2, 6]
- // CHECK: tensor.yield %[[NINF]] : f32
- // CHECK: } : tensor<8x6x10xf32> to tensor<8x8x16xf32>
- %0 = linalg.generic {indexing_maps = [#in3, #out3],
- iterator_types = ["parallel", "reduction", "reduction"]}
- ins(%in : tensor<8x6x10xf32>) outs(%init : tensor<8xf32>) {
- ^bb0(%a: f32, %o: f32):
- %m = arith.maximumf %a, %o : f32
- linalg.yield %m : f32
- } -> tensor<8xf32>
+ // minimumf neutral element is +inf.
+ // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 0x7F800000 : f32
+ // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %m = arith.minimumf %in, %out : f32
+ linalg.yield %m : f32
+ }
return %0 : tensor<8xf32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
- %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
: (!transform.any_op) -> !transform.any_op
- %padded, %pad = transform.structured.pad_tiling_interface %gen to padding_sizes [8, 8, 16]
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@@ -99,19 +99,18 @@ module attributes {transform.with_named_sequence} {
// -----
-// CHECK-LABEL: @pad_reduce_minimumf
-func.func @pad_reduce_minimumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+// maxnumf is the NaN-ignoring variant, so its neutral element is NaN.
+
+// CHECK-LABEL: @pad_reduce_maxnumf
+func.func @pad_reduce_maxnumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
-> tensor<8xf32> {
- // minimumf neutral element is +inf.
- // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 0x7F800000 : f32
- // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
- // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK-DAG: %[[NAN:.*]] = arith.constant 0xFFC00000 : f32
+ // CHECK: tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NAN]] : f32
// CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
- // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
- %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
- dimensions = [1]
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>) dimensions = [1]
(%in: f32, %out: f32) {
- %m = arith.minimumf %in, %out : f32
+ %m = arith.maxnumf %in, %out : f32
linalg.yield %m : f32
}
return %0 : tensor<8xf32>
@@ -129,18 +128,26 @@ module attributes {transform.with_named_sequence} {
// -----
-// maxnumf is the NaN-ignoring variant, so its neutral element is NaN.
+// Only the reduced input takes the neutral (-inf); the non-reduced init operand
+// keeps zero. Both dims are padded so the init is padded along the parallel dim.
-// CHECK-LABEL: @pad_reduce_maxnumf
-func.func @pad_reduce_maxnumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+// CHECK-LABEL: @pad_reduce_init_stays_zero
+func.func @pad_reduce_init_stays_zero(%in: tensor<8x30xf32>, %init: tensor<8xf32>)
-> tensor<8xf32> {
- // CHECK-DAG: %[[NAN:.*]] = arith.constant 0xFFC00000 : f32
- // CHECK: tensor.pad %{{.*}} low[0, 0] high[0, 2]
- // CHECK: tensor.yield %[[NAN]] : f32
- // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
- %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>) dimensions = [1]
- (%in: f32, %out: f32) {
- %m = arith.maxnumf %in, %out : f32
+ // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+ // CHECK-DAG: %[[NINF:.*]] = arith.constant 0xFF800000 : f32
+ // Reduced input -> maximumf neutral (-inf).
+ // CHECK: tensor.pad %{{.*}}
+ // CHECK: tensor.yield %[[NINF]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<16x32xf32>
+ // Init (non-reduced) operand -> plain zero.
+ // CHECK: tensor.pad %{{.*}}
+ // CHECK: tensor.yield %[[ZERO]] : f32
+ // CHECK: } : tensor<8xf32> to tensor<16xf32>
+ %0 = linalg.reduce ins(%in : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%a: f32, %o: f32) {
+ %m = arith.maximumf %a, %o : f32
linalg.yield %m : f32
}
return %0 : tensor<8xf32>
@@ -150,7 +157,7 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%red = transform.structured.match ops{["linalg.reduce"]} in %arg1
: (!transform.any_op) -> !transform.any_op
- %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [16, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@@ -186,6 +193,41 @@ module attributes {transform.with_named_sequence} {
// -----
+// Negative: the reduced operand is not consumed directly by a recognized
+// reduction combiner (it flows through math.exp first), so no safe neutral
+// element can be inferred and padding fails conservatively rather than
+// silently padding with a wrong value.
+
+func.func @negative_pad_reduce_indirect_combiner(
+ %input: tensor<8x30xf32>, %init: tensor<8xf32>) -> tensor<8xf32> {
+ // expected-note @below {{target op}}
+ %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %e = math.exp %in : f32
+ %a = arith.addf %e, %out : f32
+ linalg.yield %a : f32
+ }
+ return %0 : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // expected-error @below {{failed to pad op}}
+ %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.matmul
+///--------------------------------------------------------------------------------------
+
// A contraction is padded with zero: its reduced operands feed the multiply and
// zero annihilates it (0 * x = 0), so the add-reduction stays correct. The
// combiner (arith.addf) must NOT be mistaken for the operands' direct combiner
@@ -219,67 +261,37 @@ module attributes {transform.with_named_sequence} {
// -----
-// Only the reduced input takes the neutral (-inf); the non-reduced init operand
-// keeps zero. Both dims are padded so the init is padded along the parallel dim.
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.generic
+///--------------------------------------------------------------------------------------
-// CHECK-LABEL: @pad_reduce_init_stays_zero
-func.func @pad_reduce_init_stays_zero(%in: tensor<8x30xf32>, %init: tensor<8xf32>)
+// An operand reduced along several dimensions is padded with the neutral on all
+// of them (here d1 and d2); one neutral suffices since the combiner is the same.
+
+#in3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#out3 = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK-LABEL: @pad_generic_two_reduction_dims
+func.func @pad_generic_two_reduction_dims(%in: tensor<8x6x10xf32>, %init: tensor<8xf32>)
-> tensor<8xf32> {
- // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[NINF:.*]] = arith.constant 0xFF800000 : f32
- // Reduced input -> maximumf neutral (-inf).
- // CHECK: tensor.pad %{{.*}}
+ // CHECK: tensor.pad %{{.*}} low[0, 0, 0] high[0, 2, 6]
// CHECK: tensor.yield %[[NINF]] : f32
- // CHECK: } : tensor<8x30xf32> to tensor<16x32xf32>
- // Init (non-reduced) operand -> plain zero.
- // CHECK: tensor.pad %{{.*}}
- // CHECK: tensor.yield %[[ZERO]] : f32
- // CHECK: } : tensor<8xf32> to tensor<16xf32>
- %0 = linalg.reduce ins(%in : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
- dimensions = [1]
- (%a: f32, %o: f32) {
- %m = arith.maximumf %a, %o : f32
- linalg.yield %m : f32
- }
- return %0 : tensor<8xf32>
-}
-
-module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
- %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
- : (!transform.any_op) -> !transform.any_op
- %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [16, 32]
- : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
- transform.yield
- }
-}
-
-// -----
-
-// Negative: the reduced operand is not consumed directly by a recognized
-// reduction combiner (it flows through math.exp first), so no safe neutral
-// element can be inferred and padding fails conservatively rather than
-// silently padding with a wrong value.
-
-func.func @negative_pad_reduce_indirect_combiner(
- %input: tensor<8x30xf32>, %init: tensor<8xf32>) -> tensor<8xf32> {
- // expected-note @below {{target op}}
- %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
- dimensions = [1]
- (%in: f32, %out: f32) {
- %e = math.exp %in : f32
- %a = arith.addf %e, %out : f32
- linalg.yield %a : f32
- }
+ // CHECK: } : tensor<8x6x10xf32> to tensor<8x8x16xf32>
+ %0 = linalg.generic {indexing_maps = [#in3, #out3],
+ iterator_types = ["parallel", "reduction", "reduction"]}
+ ins(%in : tensor<8x6x10xf32>) outs(%init : tensor<8xf32>) {
+ ^bb0(%a: f32, %o: f32):
+ %m = arith.maximumf %a, %o : f32
+ linalg.yield %m : f32
+ } -> tensor<8xf32>
return %0 : tensor<8xf32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
- %red = transform.structured.match ops{["linalg.reduce"]} in %arg1
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
: (!transform.any_op) -> !transform.any_op
- // expected-error @below {{failed to pad op}}
- %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+ %padded, %pad = transform.structured.pad_tiling_interface %gen to padding_sizes [8, 8, 16]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
}
@@ -293,7 +305,7 @@ module attributes {transform.with_named_sequence} {
#map = affine_map<(d0, d1) -> (d0, d1)>
#mapr = affine_map<(d0, d1) -> (d0)>
-func.func @negative_pad_multi_reduction(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
+func.func @negative_pad_generic_multi_reduction(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
%s0: tensor<8xf32>)
-> (tensor<8xf32>, tensor<8xf32>) {
// expected-note @below {{target op}}
@@ -318,3 +330,38 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Negative: the body is contraction-like but accumulates with `maximumf`, for
+// which no padding value of the inputs maps to the neutral element of the
+// accumulate, so inference fails rather than zero-padding a max reduction.
+
+#mapA = affine_map<(d0, d1, d2) -> (d0, d2)>
+#mapB = affine_map<(d0, d1, d2) -> (d2, d1)>
+#mapC = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @negative_pad_generic_max_accumulate(%A: tensor<8x10xf32>,
+ %B: tensor<10x8xf32>, %C: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ // expected-note @below {{target op}}
+ %0 = linalg.generic {indexing_maps = [#mapA, #mapB, #mapC],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<8x10xf32>, tensor<10x8xf32>)
+ outs(%C : tensor<8x8xf32>) {
+ ^bb0(%a: f32, %b: f32, %acc: f32):
+ %m = arith.mulf %a, %b : f32
+ %r = arith.maximumf %m, %acc : f32
+ linalg.yield %r : f32
+ } -> tensor<8x8xf32>
+ return %0 : tensor<8x8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %gen = transform.structured.match ops{["linalg.generic"]} in %arg1
+ : (!transform.any_op) -> !transform.any_op
+ // expected-error @below {{failed to pad op}}
+ %padded, %pad = transform.structured.pad_tiling_interface %gen to padding_sizes [8, 8, 16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
>From 7cf426e2935cd9bd6233ecee83b97f2bc71bb67a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 15 Sep 2026 15:56:07 +0000
Subject: [PATCH 6/6] address feedback
---
.../Linalg/Transforms/PadTilingInterface.cpp | 13 ++++++++-----
mlir/test/Dialect/Linalg/pad-neutral.mlir | 5 +++--
2 files changed, 11 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index 7a9dfe2ef32b7..8c6ad2e65bd79 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
@@ -270,6 +270,9 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
/// of `linalgOp`
static bool isReducedOperand(linalg::LinalgOp linalgOp, OpOperand *operand,
ArrayRef<utils::IteratorType> iterTypes) {
+ assert(operand->getOwner() == linalgOp.getOperation() &&
+ "expected operand to belong to the LinalgOp");
+
AffineMap map = linalgOp.getMatchingIndexingMap(operand);
return llvm::any_of(llvm::enumerate(iterTypes), [&](auto it) {
return it.value() == utils::IteratorType::reduction &&
@@ -277,14 +280,13 @@ static bool isReducedOperand(linalg::LinalgOp linalgOp, OpOperand *operand,
});
}
-/// Returns `defaults` with the entry of every input of `linalgOp` that is
-/// indexed along a reduction dimension replaced by its padding value. Fails if
-/// the body of `linalgOp` is not contraction-like, or if its
+/// On success, returns a vector with pad values for every operand of linalgOp.
+/// Fails if the body of `linalgOp` is not contraction-like, or if its
/// `elemwise`/`reduce` pair admits no padding value.
static FailureOr<SmallVector<Attribute>>
inferContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp linalgOp,
ArrayRef<utils::IteratorType> iterTypes,
- ArrayRef<Attribute> defaults) {
+ ArrayRef<Attribute> defaultPadVals) {
// NOTE: For contraction-like Ops, we infer the padding value by looking at
// both elemwise and reduce Ops, where (see isContractionBody for details):
// %0 = <elemwise>(permutation-of(cu(block-argument-0),
@@ -318,7 +320,8 @@ inferContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp linalgOp,
return {};
};
- SmallVector<Attribute> paddingValues(defaults.begin(), defaults.end());
+ SmallVector<Attribute> paddingValues(defaultPadVals.begin(),
+ defaultPadVals.end());
for (OpOperand *input : linalgOp.getDpsInputOperands()) {
if (!isReducedOperand(linalgOp, input, iterTypes))
continue;
diff --git a/mlir/test/Dialect/Linalg/pad-neutral.mlir b/mlir/test/Dialect/Linalg/pad-neutral.mlir
index 7363c104ea07e..c85ad4d73c705 100644
--- a/mlir/test/Dialect/Linalg/pad-neutral.mlir
+++ b/mlir/test/Dialect/Linalg/pad-neutral.mlir
@@ -128,8 +128,9 @@ module attributes {transform.with_named_sequence} {
// -----
-// Only the reduced input takes the neutral (-inf); the non-reduced init operand
-// keeps zero. Both dims are padded so the init is padded along the parallel dim.
+// Only the reduced input takes `-inf`, the reduction Op neutral value (`maximumf`),
+// as the pad value; the non-reduced init operand keeps zero.
+// Both dims are padded so the init is padded along the parallel dim.
// CHECK-LABEL: @pad_reduce_init_stays_zero
func.func @pad_reduce_init_stays_zero(%in: tensor<8x30xf32>, %init: tensor<8xf32>)
More information about the Mlir-commits
mailing list