[Mlir-commits] [mlir] a3d2f4b - [[mlir][linalg] Infer reduction-neutral padding values in rewriteAsPaddedOp (#216517)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 15 11:04:46 PDT 2026
Author: Jianhui Li
Date: 2026-09-15T11:04:40-07:00
New Revision: a3d2f4b4a0aee08f26aa218e4e4d9dae1afa1d7a
URL: https://github.com/llvm/llvm-project/commit/a3d2f4b4a0aee08f26aa218e4e4d9dae1afa1d7a
DIFF: https://github.com/llvm/llvm-project/commit/a3d2f4b4a0aee08f26aa218e4e4d9dae1afa1d7a.diff
LOG: [[mlir][linalg] Infer reduction-neutral padding values in rewriteAsPaddedOp (#216517)
**Problem**:
with no explicit options.paddingValues, rewriteAsPaddedOp padded every
operand with zero — silently corrupting reductions (a padded
maximumf/mulf element gets combined).
**Change**:
Added a new public linalg::inferPaddingValues(builder, toPad) to pick a
semantics-preserving value per operand:
- no reduction dim → zero;
- contraction → zero (0 annihilates through the multiply);
- other reductions → the combiner's neutral (-inf for maximumf, 1 for
mulf, …), found via matchReduction so the real accumulator combiner is
used.
rewriteAsPaddedOp calls it when options.paddingValues is empty. If no
value can be inferred, it returns failure() with the error "could not
infer a padding value"; the caller is then expected to determine the
value and supply it explicitly through options.paddingValues.
assited-by-claude
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
Added:
mlir/test/Dialect/Linalg/pad-neutral.mlir
Modified:
mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp b/mlir/lib/Dialect/Linalg/Transforms/PadTilingInterface.cpp
index ab2629b41a463..8c6ad2e65bd79 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,154 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
paddingValue, /*nofold=*/false, dynDims);
}
+/// 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) {
+ 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 &&
+ map.isFunctionOfDim(it.index());
+ });
+}
+
+/// 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> 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),
+ // 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>(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>(elemwise) && isa<arith::OrIOp>(reduce) &&
+ elemwise->getResult(0).getType().isInteger(1))
+ return builder.getZeroAttr(elementType);
+ return {};
+ };
+
+ SmallVector<Attribute> paddingValues(defaultPadVals.begin(),
+ defaultPadVals.end());
+ for (OpOperand *input : linalgOp.getDpsInputOperands()) {
+ if (!isReducedOperand(linalgOp, input, iterTypes))
+ continue;
+ Attribute padValue =
+ getPadValue(getElementTypeOrSelf(input->get().getType()));
+ if (!padValue)
+ return failure();
+ paddingValues[input->getOperandNumber()] = padValue;
+ }
+ return paddingValues;
+}
+
+/// 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.
+ 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();
+
+ // 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)
+ 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(), /*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);
+ if (fastMath &&
+ arith::bitEnumContainsAny(fastMath.getFastMathFlagsAttr().getValue(),
+ arith::FastMathFlags::nnan))
+ return failure();
+ }
+
+ for (OpOperand *input : linalgOp.getDpsInputOperands()) {
+ 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.
+ 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 +423,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/pad-neutral.mlir b/mlir/test/Dialect/Linalg/pad-neutral.mlir
new file mode 100644
index 0000000000000..c85ad4d73c705
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/pad-neutral.mlir
@@ -0,0 +1,368 @@
+// 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.
+
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.reduce
+///--------------------------------------------------------------------------------------
+
+// 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
+ }
+}
+
+// -----
+
+// 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
+ }
+}
+
+// -----
+
+// 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
+ }
+}
+
+// -----
+
+// 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> {
+ // 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
+ }
+}
+
+// -----
+
+// 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>)
+ -> 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: `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 @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>)
+ 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
+ }
+}
+
+// -----
+
+// 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
+// (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
+ }
+}
+
+// -----
+
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.generic
+///--------------------------------------------------------------------------------------
+
+// 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: %[[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
+ }
+}
+
+// -----
+
+// 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 @negative_pad_generic_multi_reduction(%in: tensor<8x30xf32>, %m0: tensor<8xf32>,
+ %s0: tensor<8xf32>)
+ -> (tensor<8xf32>, tensor<8xf32>) {
+ // expected-note @below {{target op}}
+ %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
+ }
+}
+
+// -----
+
+// 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
+ }
+}
More information about the Mlir-commits
mailing list