[Mlir-commits] [mlir] [mlir][Vector] Fix scalable dim guard in broadcast lowering (PR #225021)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Cheng-Chieh Liu (jay87406)

<details>
<summary>Changes</summary>

The scalable outer dim guard in `BroadcastOpLowering` only covered the "stretch not at start" case (`m != 0`). The "stretch at start" case unrolls the same `dstType.getDimSize(0)` trip count and has the same bug: broadcasting e.g. `vector<1x4xf32>` to `vector<[3]x4xf32>` only fills the static minimum lanes via `vector.insert`, leaving the rest as `ub.poison`. 

Widen the guard to cover both cases, and add a regression test.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp (+4-5) 
- (modified) mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index ceeb258805c6a..a16b1812a7313 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -121,11 +121,10 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
         VectorType::get(dstType.getShape().drop_front(), eltType,
                         dstType.getScalableDims().drop_front());
 
-    // For "stretch not at start" with a scalable outer dimension we would need
-    // to emit an scf.for loop, which is not yet supported.  Check before
-    // creating any IR so that returning failure() does not violate the pattern
-    // API contract.
-    if (m != 0 && dstType.getScalableDims()[0]) {
+    // For a scalable outer dimension we would need to emit an scf.for loop,
+    // which is not yet supported. Check before creating any IR so that
+    // returning failure() does not violate the pattern API contract.
+    if (dstType.getScalableDims()[0]) {
       // TODO: For scalable vectors we should emit an scf.for loop.
       return failure();
     }
diff --git a/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir b/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
index c2922cdbe55aa..0b7e6e33e06e7 100644
--- a/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
@@ -236,6 +236,17 @@ func.func @broadcast_vec2d_from_vec1d_scalable_inner(%arg0: vector<2xf32>) -> ve
   return %0 : vector<3x[4]x2xf32>
 }
 
+// TODO: Add support for scalable vectors
+
+// CHECK-LABEL: @broadcast_stretch_at_start_scalable
+// CHECK-SAME:  %[[A:.*]]: vector<1x4xf32>)
+// CHECK: vector.broadcast %[[A]] : vector<1x4xf32> to vector<[3]x4xf32>
+
+func.func @broadcast_stretch_at_start_scalable(%arg0: vector<1x4xf32>) -> vector<[3]x4xf32> {
+  %0 = vector.broadcast %arg0 : vector<1x4xf32> to vector<[3]x4xf32>
+  return %0 : vector<[3]x4xf32>
+}
+
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
     %f = transform.structured.match ops{["func.func"]} in %module_op

``````````

</details>


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