[Mlir-commits] [mlir] [MLIR][Vector] Broadcast folding creates invalid scalar-result vector.contract (PR #230256)
Atharva Ajmera
llvmlistbot at llvm.org
Thu Oct 29 04:13:29 PDT 2026
https://github.com/atharvaajmera updated https://github.com/llvm/llvm-project/pull/230256
>From f07c831fef17be0e4e25a42741b970a5b2fff1a4 Mon Sep 17 00:00:00 2001
From: atharvaajmera <atharvaajmera06 at gmail.com>
Date: Fri, 9 Oct 2026 00:54:14 +0530
Subject: [PATCH 1/6] [MLIR][Vector] Broadcast folding creates invalid
scalar-result vector.contract
Fixes #230084
---
.../Vector/Transforms/VectorTransforms.cpp | 7 ++---
.../Vector/vector-reduce-to-contract.mlir | 28 +++++++++++++++++++
2 files changed, 31 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 80b2b017404e3c..87117c9a1d92ce 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -382,10 +382,9 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
for (unsigned i = 0; i < iterators.size(); ++i) {
if (!isReductionIterator(iterators[i]))
continue;
- if (getResultIndex(maps[0], i) && getResultIndex(maps[1], i)) {
- hasReductionIteratorApplyingOnBothSides = true;
- break;
- }
+ if (!getResultIndex(maps[0], i) || !getResultIndex(maps[1], i))
+ return failure();
+ hasReductionIteratorApplyingOnBothSides = true;
}
if (!hasReductionIteratorApplyingOnBothSides)
return failure();
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index 8e53a815f70f88..b08b42773998e1 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -526,3 +526,31 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
%resT = vector.transpose %contract, [0, 2, 1] : vector<2x4x8xf32> to vector<2x8x4xf32>
return %resT : vector<2x8x4xf32>
}
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
+// CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
+// CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
+// CHECK: %[[RHS:.+]] = vector.broadcast %[[ARG1]] : bf16 to vector<1x1x2xbf16>
+// CHECK: vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
+// CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
+// CHECK-SAME: %[[LHS]], %[[RHS]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
+ %zero = arith.constant 0.0 : bf16
+ %lhs = vector.shape_cast %arg0 : vector<2xbf16> to vector<1x1x2xbf16>
+ %rhs = vector.broadcast %arg1 : bf16 to vector<1x1x2xbf16>
+ %result = vector.contract {
+ indexing_maps = [#map0, #map0, #map1],
+ iterator_types = ["reduction", "reduction", "reduction"],
+ kind = #vector.kind<add>
+ } %lhs, %rhs, %zero : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+ return %result : bf16
+}
>From 58d5224647a8bc5722807175a9389f4138f5e7ab Mon Sep 17 00:00:00 2001
From: Atharva Ajmera <atharvaajmera06 at gmail.com>
Date: Fri, 9 Oct 2026 16:21:31 +0530
Subject: [PATCH 2/6] [MLIR][Vector] Add missing accumulator CHECK to broadcast
fold test
---
mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index b08b42773998e1..22678c056bc8f6 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -537,12 +537,13 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
// CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
+// CHECK: %[[CST:.+]] = arith.constant 0.0{{.*}}: bf16
// CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
// CHECK: %[[RHS:.+]] = vector.broadcast %[[ARG1]] : bf16 to vector<1x1x2xbf16>
// CHECK: vector.contract
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
// CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
-// CHECK-SAME: %[[LHS]], %[[RHS]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
%zero = arith.constant 0.0 : bf16
%lhs = vector.shape_cast %arg0 : vector<2xbf16> to vector<1x1x2xbf16>
>From 2df0aca1da8d09f697b18d61c2fec91bb81d1655 Mon Sep 17 00:00:00 2001
From: atharvaajmera <atharvaajmera06 at gmail.com>
Date: Sat, 10 Oct 2026 03:26:59 +0530
Subject: [PATCH 3/6] Add comments to one-sided broadcast fold test
---
mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir | 7 +++++++
1 file changed, 7 insertions(+)
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index 22678c056bc8f6..f07cff2753c1c1 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -529,6 +529,13 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
// -----
+// Regression test for the one-sided broadcast fold in CombineContractBroadcast.
+// Folding a broadcast-like shape_cast into only one operand of a contraction
+// must be rejected: the pass used to accept it whenever a single reduction pair
+// remained, leaving the other retained reductions unpaired. Those became free
+// dimensions and produced an invalid scalar-result vector.contract. Here the
+// fold is refused, so the shape_cast, broadcast and contract are preserved.
+
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> ()>
>From 769938c3d6047917423af7db957f680cf4802a0d Mon Sep 17 00:00:00 2001
From: Atharva Ajmera <atharvaajmera06 at gmail.com>
Date: Sat, 10 Oct 2026 13:43:44 +0530
Subject: [PATCH 4/6] Change the test comment
---
mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir | 9 +++------
1 file changed, 3 insertions(+), 6 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index f07cff2753c1c1..522e79ca1b001f 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -529,12 +529,9 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
// -----
-// Regression test for the one-sided broadcast fold in CombineContractBroadcast.
-// Folding a broadcast-like shape_cast into only one operand of a contraction
-// must be rejected: the pass used to accept it whenever a single reduction pair
-// remained, leaving the other retained reductions unpaired. Those became free
-// dimensions and produced an invalid scalar-result vector.contract. Here the
-// fold is refused, so the shape_cast, broadcast and contract are preserved.
+// Test that CombineContractBroadcast is not combining this case, as that would
+// result in the reduction iterators of the other operand becoming unpaired and
+// produce an invalid scalar-result vector.contract.
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> ()>
>From eda4910959874dbb7762b8dbcc2bad60f64f63aa Mon Sep 17 00:00:00 2001
From: atharvaajmera <atharvaajmera06 at gmail.com>
Date: Mon, 12 Oct 2026 19:57:44 +0530
Subject: [PATCH 5/6] fixes for the latest review
---
.../Vector/Transforms/VectorTransforms.cpp | 8 +--
.../Vector/vector-reduce-to-contract.mlir | 65 +++++++++----------
2 files changed, 35 insertions(+), 38 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 87117c9a1d92ce..d86989eb1fd1c4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -373,11 +373,9 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
}
}
- // Check that compressing unused dims isn't removing all reduction dimension
- // pairs. For example, if the vector.contract had only one reduction
- // iterator and that was a unit-dimension created by a broadcast,
- // then we should bail here, otherwise we would create a contract without
- // a reduction dimension pair.
+ // Bail if no reduction iterator survives on both LHS and RHS: folding a
+ // one-sided unit-dim broadcast would leave an unpaired reduction iterator
+ // and produce an invalid contract.
bool hasReductionIteratorApplyingOnBothSides = false;
for (unsigned i = 0; i < iterators.size(); ++i) {
if (!isReductionIterator(iterators[i]))
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index 522e79ca1b001f..1d6901db1f5d89 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -493,6 +493,38 @@ func.func @contract_broadcast_would_have_no_reduction_dim_pair(%arg0 : vector<1x
return %result : vector<1xf32>
}
+// -----
+
+// Test that CombineContractBroadcast is not combining this case, as that would
+// result in the reduction iterators of the other operand becoming unpaired and
+// produce an invalid scalar-result vector.contract.
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> ()>
+
+// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
+// CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
+// CHECK: %[[CST:.+]] = arith.constant 0.0{{.*}}: bf16
+// CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
+// CHECK: %[[RHS:.+]] = vector.broadcast %[[ARG1]] : bf16 to vector<1x1x2xbf16>
+// CHECK: vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
+// CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
+ %zero = arith.constant 0.0 : bf16
+ %lhs = vector.shape_cast %arg0 : vector<2xbf16> to vector<1x1x2xbf16>
+ %rhs = vector.broadcast %arg1 : bf16 to vector<1x1x2xbf16>
+ %result = vector.contract {
+ indexing_maps = [#map0, #map0, #map1],
+ iterator_types = ["reduction", "reduction", "reduction"],
+ kind = #vector.kind<add>
+ } %lhs, %rhs, %zero : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+ return %result : bf16
+}
// -----
@@ -526,36 +558,3 @@ func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8x
%resT = vector.transpose %contract, [0, 2, 1] : vector<2x4x8xf32> to vector<2x8x4xf32>
return %resT : vector<2x8x4xf32>
}
-
-// -----
-
-// Test that CombineContractBroadcast is not combining this case, as that would
-// result in the reduction iterators of the other operand becoming unpaired and
-// produce an invalid scalar-result vector.contract.
-
-#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-#map1 = affine_map<(d0, d1, d2) -> ()>
-
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> ()>
-
-// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
-// CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
-// CHECK: %[[CST:.+]] = arith.constant 0.0{{.*}}: bf16
-// CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
-// CHECK: %[[RHS:.+]] = vector.broadcast %[[ARG1]] : bf16 to vector<1x1x2xbf16>
-// CHECK: vector.contract
-// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
-// CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
-// CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
-func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
- %zero = arith.constant 0.0 : bf16
- %lhs = vector.shape_cast %arg0 : vector<2xbf16> to vector<1x1x2xbf16>
- %rhs = vector.broadcast %arg1 : bf16 to vector<1x1x2xbf16>
- %result = vector.contract {
- indexing_maps = [#map0, #map0, #map1],
- iterator_types = ["reduction", "reduction", "reduction"],
- kind = #vector.kind<add>
- } %lhs, %rhs, %zero : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
- return %result : bf16
-}
>From 0877abb6cd3f066afdd50753da95b3b6bc34915b Mon Sep 17 00:00:00 2001
From: atharvaajmera <atharvaajmera06 at gmail.com>
Date: Thu, 29 Oct 2026 16:22:27 +0530
Subject: [PATCH 6/6] [MLIR][Vector] Rename negative broadcast-fold regression
test
Follow the MLIR test naming convention requested in review.
---
mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index 1d6901db1f5d89..eec1a06e2ea7a7 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -505,7 +505,7 @@ func.func @contract_broadcast_would_have_no_reduction_dim_pair(%arg0 : vector<1x
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> ()>
-// CHECK-LABEL: contract_broadcast_one_sided_unit_dim_reduction
+// CHECK-LABEL: negative_contract_broadcast_one_sided_unit_dim_reduction
// CHECK-SAME: (%[[ARG0:.+]]: vector<2xbf16>, %[[ARG1:.+]]: bf16)
// CHECK: %[[CST:.+]] = arith.constant 0.0{{.*}}: bf16
// CHECK: %[[LHS:.+]] = vector.shape_cast %[[ARG0]] : vector<2xbf16> to vector<1x1x2xbf16>
@@ -514,7 +514,7 @@ func.func @contract_broadcast_would_have_no_reduction_dim_pair(%arg0 : vector<1x
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]]
// CHECK-SAME: iterator_types = ["reduction", "reduction", "reduction"]
// CHECK-SAME: %[[LHS]], %[[RHS]], %[[CST]] : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
-func.func @contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
+func.func @negative_contract_broadcast_one_sided_unit_dim_reduction(%arg0 : vector<2xbf16>, %arg1 : bf16) -> bf16 {
%zero = arith.constant 0.0 : bf16
%lhs = vector.shape_cast %arg0 : vector<2xbf16> to vector<1x1x2xbf16>
%rhs = vector.broadcast %arg1 : bf16 to vector<1x1x2xbf16>
More information about the Mlir-commits
mailing list