[Mlir-commits] [mlir] [mlir][vector] Disallow vector broadcast along scalable dim (PR #212197)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Bryth (Brythzz)

<details>
<summary>Changes</summary>

Broadcasts from vectors to higher rank vectors get lowered recursively by iterating over leading dimensions to create "size(dimension)" insert ops. This cannot be done for scalable dimensions whose size is unknown. This PR prevents such transforms from occurring by stopping if the leading dimension is scalable.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp (+2) 
- (modified) mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir (+12) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index a78c31367ee28..0e31b6934e239 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -73,6 +73,8 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
     //   %x = [%b,%b,%b,%b] : n-D
     if (srcRank < dstRank) {
       // Duplication.
+      if (dstType.getScalableDims()[0])
+        return failure();
       VectorType resType = VectorType::Builder(dstType).dropDim(0);
       Value bcst =
           vector::BroadcastOp::create(rewriter, loc, resType, op.getSource());
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
index 094353119e493..3da33dad91efc 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -242,6 +242,18 @@ func.func @broadcast_vec2d_from_vec1d_scalable(%arg0: vector<[2]xf32>) -> vector
 
 // -----
 
+// TODO: Add support for scalable vectors
+
+func.func @broadcast_vec2d_from_vec1d_scalable_leading(%arg0: vector<2xf32>) -> vector<[4]x2xf32> {
+  %0 = vector.broadcast %arg0 : vector<2xf32> to vector<[4]x2xf32>
+  return %0 : vector<[4]x2xf32>
+}
+// CHECK-LABEL: @broadcast_vec2d_from_vec1d_scalable_leading
+// CHECK-SAME:  %[[A:.*]]: vector<2xf32>)
+// CHECK: vector.broadcast %[[A]] : vector<2xf32> to vector<[4]x2xf32>
+
+// -----
+
 func.func @broadcast_vec2d_from_index_vec1d(%arg0: vector<2xindex>) -> vector<3x2xindex> {
   %0 = vector.broadcast %arg0 : vector<2xindex> to vector<3x2xindex>
   return %0 : vector<3x2xindex>

``````````

</details>


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