[Mlir-commits] [mlir] [mlir][vector] Do not unroll vector.contract with scalar result (PR #232203)

[llvmlistbot at llvm.org]

https://github.com/StarOne01 created https://github.com/llvm/llvm-project/pull/232203

add a check to `UnrollContractionPattern` to bail out if we have a scalar result.

fixes #231904

>From 8842ea15f3590c2fad30a1fe173ba8f9161def21 Mon Sep 17 00:00:00 2001
From: Prashanth <ping at starone01.me>
Date: Thu, 15 Oct 2026 12:39:41 +0530
Subject: [PATCH] [mlir][vector] add a check to UnrollContractionPattern to
 bail out if we have a scalar result and add contract to scalar test case

---
 .../Dialect/Vector/Transforms/VectorUnroll.cpp  |  3 +++
 .../Dialect/Vector/vector-unroll-options.mlir   | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 0ea6bc78d3b0c..89a4604925f86 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -290,6 +290,9 @@ struct UnrollContractionPattern
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
+    if (!isa<VectorType>(contractOp.getResultType()))
+      return rewriter.notifyMatchFailure(contractOp,
+                                         "scalar result not supported");
     auto targetShape = getTargetShape(options, contractOp);
     if (!targetShape)
       return failure();
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 61a74744e7a72..22fcb483ea1e5 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -935,3 +935,20 @@ func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf3
 // CHECK-LABEL: func @masked_vector_multi_reduction
 //  CHECK-NEXT:   %[[R:.*]] = vector.mask %{{.*}} { vector.multi_reduction <add>, {{.*}} : vector<4x6xf32> to vector<4xf32> }
 //  CHECK-NEXT:   return %[[R]]
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0) -> ()>
+
+func.func @vector_contract_to_scalar(%a: vector<4xf32>, %b: vector<4xf32>, %acc: f32) -> f32 {
+  %r = vector.contract {indexing_maps = [#map, #map, #map0], iterator_types = ["reduction"], kind = #vector.kind<add>} %a, %b, %acc : vector<4xf32>, vector<4xf32> into f32
+  return %r : f32
+}
+
+// CHECK-LABEL: func @vector_contract_to_scalar(
+// CHECK-SAME:    %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[ACC:.*]]: f32) -> f32 {
+// CHECK-NOT:     vector.extract_strided_slice
+// CHECK:         %[[R:.*]] = vector.contract
+// CHECK-SAME:      %[[A]], %[[B]], %[[ACC]] : vector<4xf32>, vector<4xf32> into f32
+// CHECK:         return %[[R]] : f32