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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Prashanth (StarOne01)

<details>
<summary>Changes</summary>

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

fixes #<!-- -->231904

---
Full diff: https://github.com/llvm/llvm-project/pull/232203.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+3) 
- (modified) mlir/test/Dialect/Vector/vector-unroll-options.mlir (+17) 


``````````diff
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

``````````

</details>


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