[Mlir-commits] [mlir] [Mlir][Vector] fix contraction lowering crashes on scalar results after folding a broadcast-like shape_cast (PR #229824)

[llvmlistbot at llvm.org]

https://github.com/StarOne01 updated https://github.com/llvm/llvm-project/pull/229824

>From 2b647f9214a9f0080b897719020bf0de0ae64bad Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Wed, 7 Oct 2026 21:22:07 +0530
Subject: [PATCH 1/4] [Mlir][Vector] Change result type in lowerParallel from
 VectorType to Type

---
 mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
index b3cd4fb1e8f9b0..76e71f76063b0f 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
@@ -1010,7 +1010,7 @@ FailureOr<Value> ContractionOpLowering::lowerParallel(PatternRewriter &rewriter,
                                                       Value mask) const {
   VectorType lhsType = op.getLhsType();
   VectorType rhsType = op.getRhsType();
-  VectorType resType = cast<VectorType>(op.getResultType());
+  Type resType = op.getResultType();
   // Find the iterator type index and result index.
   SmallVector<AffineMap> iMap = op.getIndexingMapsArray();
   int64_t iterIndex = -1;

>From d0d46d1f8838aca6210b230c62c5920ef502d4e7 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Thu, 8 Oct 2026 06:08:54 +0530
Subject: [PATCH 2/4] [mlir][vector] Add test for lowering a scalar-result
 contraction with one-sided unit dims

---
 .../test/Dialect/Vector/transform-vector.mlir | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/mlir/test/Dialect/Vector/transform-vector.mlir b/mlir/test/Dialect/Vector/transform-vector.mlir
index 16011e3041eb55..c5e4115cc3dc29 100644
--- a/mlir/test/Dialect/Vector/transform-vector.mlir
+++ b/mlir/test/Dialect/Vector/transform-vector.mlir
@@ -139,3 +139,38 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// Regression test for a crash in the lowering of a contraction with a scalar
+// result when unit reduction dims appear on only one operand after folding.
+// see https://github.com/llvm/llvm-project/issues/228971
+
+// CHECK-LABEL: func.func @contract_scalar_result_unit_dims_lowering
+//  CHECK-SAME:   %[[X:.*]]: vector<2xbf16>, %{{.*}}: bf16, %[[ACC:.*]]: bf16) -> bf16
+//       CHECK:   %[[M:.*]] = arith.mulf %[[X]], %{{.*}} : vector<2xbf16>
+//       CHECK:   %[[R:.*]] = vector.reduction <add>, %[[M]], %[[ACC]] : vector<2xbf16> into bf16
+//  CHECK-NEXT:   return %[[R]] : bf16
+func.func @contract_scalar_result_unit_dims_lowering(%x: vector<2xbf16>, %scale: bf16, %acc: bf16) -> bf16 {
+  %lhs = vector.shape_cast %x : vector<2xbf16> to vector<1x1x2xbf16>
+  %rhs = vector.broadcast %scale : bf16 to vector<1x1x2xbf16>
+  %out = vector.contract {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                     affine_map<(d0, d1, d2) -> ()>],
+    iterator_types = ["reduction", "reduction", "reduction"],
+    kind = #vector.kind<add>
+  } %lhs, %rhs, %acc : vector<1x1x2xbf16>, vector<1x1x2xbf16> into bf16
+  return %out : bf16
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+    %func = transform.structured.match ops{["func.func"]} in %module_op : (!transform.any_op) -> !transform.any_op
+    transform.apply_patterns to %func {
+      transform.apply_patterns.vector.reduction_to_contract
+      transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+    } : !transform.any_op
+    transform.yield
+  }
+}
\ No newline at end of file

>From 72a99b2edf9a790a592aa93c8f4b29adf10b11a2 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Thu, 8 Oct 2026 06:13:40 +0530
Subject: [PATCH 3/4] [mlir][vector] add missing newline at end of file in
 transform-vector.mlir

---
 mlir/test/Dialect/Vector/transform-vector.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/Vector/transform-vector.mlir b/mlir/test/Dialect/Vector/transform-vector.mlir
index c5e4115cc3dc29..b163c2472e8222 100644
--- a/mlir/test/Dialect/Vector/transform-vector.mlir
+++ b/mlir/test/Dialect/Vector/transform-vector.mlir
@@ -173,4 +173,4 @@ module attributes {transform.with_named_sequence} {
     } : !transform.any_op
     transform.yield
   }
-}
\ No newline at end of file
+}

>From dd8365f470e1f39093ab5d16bcdaf6f2a816c3a0 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Fri, 9 Oct 2026 05:43:35 +0530
Subject: [PATCH 4/4] [mlir][vector] simplify comment of the
 contract_scalar_result_unit_dims_lowering test

---
 mlir/test/Dialect/Vector/transform-vector.mlir | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mlir/test/Dialect/Vector/transform-vector.mlir b/mlir/test/Dialect/Vector/transform-vector.mlir
index b163c2472e8222..f291301e1d4cdb 100644
--- a/mlir/test/Dialect/Vector/transform-vector.mlir
+++ b/mlir/test/Dialect/Vector/transform-vector.mlir
@@ -142,8 +142,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// Regression test for a crash in the lowering of a contraction with a scalar
-// result when unit reduction dims appear on only one operand after folding.
+// Regression test for a crash in the lowering of a contraction with a scalar result 
 // see https://github.com/llvm/llvm-project/issues/228971
 
 // CHECK-LABEL: func.func @contract_scalar_result_unit_dims_lowering