[Mlir-commits] [mlir] [MLIR][Vector] Broadcast folding creates invalid scalar-result vector.contract (PR #230256)

[Atharva Ajmera]

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/3] [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/3] [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/3] 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) -> ()>