[Mlir-commits] [mlir] [mlir][linalg] Remap linalg.index when decomposing unit convolutions (PR #234882)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 02:25:10 PDT 2026
https://github.com/Anoymoguy created https://github.com/llvm/llvm-project/pull/234882
`downscaleSizeOneWindowedConvolution` removes unit output/filter spatial loops and compresses the indexing maps, but moves the scalar region without updating `linalg.index`. For example, after `(oh, ow, f, kh, kw, c)` becomes `(ow, f, kw, c)`, the original `index 1` reads `f` instead of `ow`. References to later dimensions can also become invalid.
Replace indices of removed loops with constant zero and renumber surviving indices by the number of preceding removed loops. Restrict the rewrite to the new generic's immediate body so nested Linalg operations retain their own index semantics; captured outer indices are updated through their SSA uses.
Add regression coverage for all six original indices, removal of either spatial pair, and a nested generic that uses both its own index and a captured outer index.
Part of #220505 (the `transform.structured.decompose` case).
Validation:
- `git diff --check` and `git-clang-format --diff`: clean (clang-format 23.1.4; the upstream formatting image specifies 23.1.3).
- Native Windows validation using fixed/unfixed drivers linked against the LLVM `416f328b939f` SDK: the minimal reproducer goes from six loops to four, `index 1` becomes `index 0`, and the result reparses successfully. All convolution FileCheck checks in this file, including both new regressions, pass. The function in this PR matches the compiled SDK port apart from whitespace.
- A full build of the PR head and `check-mlir` have not been run locally. The older SDK fails the unchanged softmax expectation updated by #232257; this PR preserves that upstream expectation. No numerical execution was performed.
Draft pending contributor self-review and upstream CI.
Assisted-by: Codex
>From df3afa8285274ad09acc0c2f8bfc3fe1b402c730 Mon Sep 17 00:00:00 2001
From: Anoymoguy <1553486465 at qq.com>
Date: Fri, 23 Oct 2026 17:23:55 +0800
Subject: [PATCH] [mlir][linalg] Remap linalg.index when decomposing unit
convolutions
Preserve region-visible loop coordinates when downscaling a convolution.
Replace indices of removed unit loops with zero and compress surviving
indices, without rewriting indices owned by nested Linalg operations.
Cover both spatial removal cases, all original loop indices, and captured
outer indices in a nested generic.
Part of #220505 (transform.structured.decompose).
Assisted-by: Codex
---
.../Dialect/Linalg/Transforms/Transforms.cpp | 18 +++
.../Linalg/transform-op-decompose.mlir | 127 ++++++++++++++++++
2 files changed, 145 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index f6dcb4e93565be..8288d0d5113540 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1625,6 +1625,24 @@ linalg::downscaleSizeOneWindowedConvolution(RewriterBase &rewriter,
rewriter.inlineRegionBefore(op->getRegion(0), newOp.getRegion(),
newOp.getRegion().begin());
+ // Keep index semantics consistent with the compressed iteration domain. Only
+ // visit this body's indices: nested LinalgOps have their own loop dimensions.
+ for (IndexOp indexOp :
+ llvm::make_early_inc_range(newOp.getBody()->getOps<IndexOp>())) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(indexOp);
+ if (llvm::is_contained(loopDimsToRemove, indexOp.getDim())) {
+ rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
+ } else {
+ unsigned droppedDims =
+ llvm::count_if(loopDimsToRemove,
+ [&](unsigned dim) { return dim < indexOp.getDim(); });
+ if (droppedDims != 0)
+ rewriter.replaceOpWithNewOp<IndexOp>(indexOp,
+ indexOp.getDim() - droppedDims);
+ }
+ }
+
// Try to specialize the generic back to a named op only if the input was
// already a specialized (named) op.
LinalgOp resultOp = newOp;
diff --git a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
index 05a01b0a20eb36..a5caaf43ee61ff 100644
--- a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
@@ -272,6 +272,58 @@ func.func @batchless_conv_2d_hwc_hwcf(%input: tensor<1x14x8xf32>, %filter: tenso
return %0 : tensor<1x12x16xf32>
}
+// Dropping oh and kh must replace their indices by zero and renumber every
+// remaining index from (oh, ow, f, kh, kw, c) to (ow, f, kw, c).
+// CHECK-LABEL: @batchless_conv_2d_hwc_hwcf_indices
+func.func @batchless_conv_2d_hwc_hwcf_indices(%input: tensor<1x3x1xi32>, %filter: tensor<1x1x1x2xi32>, %init: tensor<1x3x2xi32>) -> tensor<1x3x2xi32> {
+ // CHECK: %[[CONV:.+]] = linalg.generic
+ // CHECK-SAME: indexing_maps = [#[[$CONV_I]], #[[$CONV_F]], #[[$CONV_O]]]
+ // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "reduction"]
+ // CHECK: ^bb0(%[[IN:.+]]: i32, %[[FILTER:.+]]: i32, %[[ACC:.+]]: i32):
+ // CHECK-NEXT: %[[OH:.+]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[OW:.+]] = linalg.index 0 : index
+ // CHECK-NEXT: %[[F:.+]] = linalg.index 1 : index
+ // CHECK-NEXT: %[[KH:.+]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[KW:.+]] = linalg.index 2 : index
+ // CHECK-NEXT: %[[C:.+]] = linalg.index 3 : index
+ // CHECK-NEXT: %[[S0:.+]] = arith.addi %[[OH]], %[[OW]] : index
+ // CHECK-NEXT: %[[S1:.+]] = arith.addi %[[S0]], %[[F]] : index
+ // CHECK-NEXT: %[[S2:.+]] = arith.addi %[[S1]], %[[KH]] : index
+ // CHECK-NEXT: %[[S3:.+]] = arith.addi %[[S2]], %[[KW]] : index
+ // CHECK-NEXT: %[[S4:.+]] = arith.addi %[[S3]], %[[C]] : index
+ // CHECK-NEXT: %[[POSITION:.+]] = arith.index_cast %[[S4]] : index to i32
+ // CHECK-NEXT: %[[PRODUCT:.+]] = arith.muli %[[IN]], %[[FILTER]] : i32
+ // CHECK-NEXT: %[[SUM:.+]] = arith.addi %[[ACC]], %[[PRODUCT]] : i32
+ // CHECK-NEXT: %[[VALUE:.+]] = arith.addi %[[SUM]], %[[POSITION]] : i32
+ // CHECK-NEXT: linalg.yield %[[VALUE]] : i32
+ // CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[CONV]]
+ // CHECK: return %[[RESULT]] : tensor<1x3x2xi32>
+ %result = linalg.generic {
+ indexing_maps = [#map_conv_i, #map_conv_f, #map_conv_o],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]
+ } ins(%input, %filter : tensor<1x3x1xi32>, tensor<1x1x1x2xi32>)
+ outs(%init : tensor<1x3x2xi32>) {
+ ^bb0(%a: i32, %b: i32, %acc: i32):
+ %oh = linalg.index 0 : index
+ %ow = linalg.index 1 : index
+ %f = linalg.index 2 : index
+ %kh = linalg.index 3 : index
+ %kw = linalg.index 4 : index
+ %c = linalg.index 5 : index
+ %s0 = arith.addi %oh, %ow : index
+ %s1 = arith.addi %s0, %f : index
+ %s2 = arith.addi %s1, %kh : index
+ %s3 = arith.addi %s2, %kw : index
+ %s4 = arith.addi %s3, %c : index
+ %position = arith.index_cast %s4 : index to i32
+ %product = arith.muli %a, %b : i32
+ %sum = arith.addi %acc, %product : i32
+ %value = arith.addi %sum, %position : i32
+ linalg.yield %value : i32
+ } -> tensor<1x3x2xi32>
+ return %result : tensor<1x3x2xi32>
+}
+
#map_dw_i = affine_map<(oh, ow, c, kh, kw) -> (oh + kh, ow + kw, c)>
#map_dw_f = affine_map<(oh, ow, c, kh, kw) -> (kh, kw, c)>
#map_dw_o = affine_map<(oh, ow, c, kh, kw) -> (oh, ow, c)>
@@ -412,3 +464,78 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+#input_map = affine_map<(oh, ow, f, kh, kw, c) -> (oh + kh, ow + kw, c)>
+#filter_map = affine_map<(oh, ow, f, kh, kw, c) -> (kh, kw, c, f)>
+#output_map = affine_map<(oh, ow, f, kh, kw, c) -> (oh, ow, f)>
+#inner_map = affine_map<(i, j) -> (i, j)>
+
+// Here ow and kw are removed instead. The nested generic's own index 1 must
+// remain unchanged, while its use of the captured outer f index is updated.
+// CHECK-LABEL: @batchless_conv_2d_nested_indices
+func.func @batchless_conv_2d_nested_indices(%input: tensor<3x1x1xi32>, %filter: tensor<1x1x1x2xi32>, %init: tensor<3x1x2xi32>, %nested_init: tensor<2x2xi32>) -> tensor<3x1x2xi32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK: %[[CONV:.+]] = linalg.generic
+ // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "reduction"]
+ // CHECK: ^bb0(%[[IN:.+]]: i32, %[[FILTER:.+]]: i32, %[[ACC:.+]]: i32):
+ // CHECK-NEXT: %[[OH:.+]] = linalg.index 0 : index
+ // CHECK-NEXT: %[[OW:.+]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[F:.+]] = linalg.index 1 : index
+ // CHECK-NEXT: %[[SPATIAL:.+]] = arith.addi %[[OH]], %[[OW]] : index
+ // CHECK-NEXT: %[[POSITION:.+]] = arith.index_cast %[[SPATIAL]] : index to i32
+ // CHECK-NEXT: %[[INNER:.+]] = linalg.generic
+ // CHECK-SAME: iterator_types = ["parallel", "parallel"]
+ // CHECK: ^bb0
+ // CHECK-NEXT: %[[INNER_INDEX:.+]] = linalg.index 1 : index
+ // CHECK-NEXT: %[[WITH_F:.+]] = arith.addi %[[INNER_INDEX]], %[[F]] : index
+ // CHECK-NEXT: %[[INNER_VALUE:.+]] = arith.index_cast %[[WITH_F]] : index to i32
+ // CHECK-NEXT: linalg.yield %[[INNER_VALUE]] : i32
+ // CHECK: %[[EXTRACTED:.+]] = tensor.extract %[[INNER]]
+ // CHECK-NEXT: %[[PRODUCT:.+]] = arith.muli %[[IN]], %[[FILTER]] : i32
+ // CHECK-NEXT: %[[SUM:.+]] = arith.addi %[[ACC]], %[[PRODUCT]] : i32
+ // CHECK-NEXT: %[[INDEXED:.+]] = arith.addi %[[SUM]], %[[POSITION]] : i32
+ // CHECK-NEXT: %[[VALUE:.+]] = arith.subi %[[INDEXED]], %[[EXTRACTED]] : i32
+ // CHECK-NEXT: linalg.yield %[[VALUE]] : i32
+ // CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[CONV]]
+ // CHECK: return %[[RESULT]] : tensor<3x1x2xi32>
+ %result = linalg.generic {
+ decompose_target,
+ indexing_maps = [#input_map, #filter_map, #output_map],
+ iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]
+ } ins(%input, %filter : tensor<3x1x1xi32>, tensor<1x1x1x2xi32>)
+ outs(%init : tensor<3x1x2xi32>) {
+ ^bb0(%a: i32, %b: i32, %acc: i32):
+ %oh = linalg.index 0 : index
+ %ow = linalg.index 1 : index
+ %f = linalg.index 2 : index
+ %spatial = arith.addi %oh, %ow : index
+ %position = arith.index_cast %spatial : index to i32
+ %inner = linalg.generic {
+ indexing_maps = [#inner_map], iterator_types = ["parallel", "parallel"]
+ } outs(%nested_init : tensor<2x2xi32>) {
+ ^bb0(%unused: i32):
+ %inner_index = linalg.index 1 : index
+ %with_f = arith.addi %inner_index, %f : index
+ %inner_value = arith.index_cast %with_f : index to i32
+ linalg.yield %inner_value : i32
+ } -> tensor<2x2xi32>
+ %extracted = tensor.extract %inner[%c0, %c1] : tensor<2x2xi32>
+ %product = arith.muli %a, %b : i32
+ %sum = arith.addi %acc, %product : i32
+ %indexed = arith.addi %sum, %position : i32
+ %value = arith.subi %indexed, %extracted : i32
+ linalg.yield %value : i32
+ } -> tensor<3x1x2xi32>
+ return %result : tensor<3x1x2xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %target = transform.structured.match attributes{decompose_target} in %root : (!transform.any_op) -> !transform.any_op
+ %decomposed = transform.structured.decompose %target : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list