[Mlir-commits] [mlir] [mlir][Vector] Don't fully unroll transfer_reads of n-D scalable vectors (PR #71924)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Benjamin Maxwell (MacDue)

<details>
<summary>Changes</summary>

It is not possible to unroll a scalable vector at compile time. This currently prevents transfer_writes from being lowered to arm_sme.tile_writes (downstream).

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


2 Files Affected:

- (modified) mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp (+5) 
- (modified) mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir (+15) 


``````````diff
diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index 5fffd9091d2286d..18d6292754f1d71 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -1218,6 +1218,11 @@ struct UnrollTransferWriteConversion
 
     auto vec = getDataVector(xferOp);
     auto xferVecType = xferOp.getVectorType();
+    if (xferVecType.getScalableDims()[0]) {
+      // Cannot unroll a scalable dimension at compile time.
+      return failure();
+    }
+
     int64_t dimSize = xferVecType.getShape()[0];
     Value source = xferOp.getSource(); // memref or tensor to be written to.
     auto sourceType = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();
diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index 4880532c5528cce..cff85995e44b078 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -737,4 +737,19 @@ func.func @cannot_lower_transfer_read_with_leading_scalable(%arg0: memref<?x4xf3
 // CHECK-SAME:                                                                %[[MEMREF:.*]]: memref<?x4xf32>)
 // CHECK: %{{.*}} = vector.transfer_read %[[MEMREF]][%{{.*}}, %{{.*}}], %{{.*}}, %{{.*}} {in_bounds = [true, true]} : memref<?x4xf32>, vector<[4]x4xf32>
 
+//  -----
 
+// FULL-UNROLL-LABEL: @cannot_fully_unroll_transfer_write_of_nd_scalable_vector
+func.func @cannot_fully_unroll_transfer_write_of_nd_scalable_vector(%arg0: memref<?x?xf32>) {
+  // FULL-UNROLL-NOT: vector.extract {{.*}} : vector<[4]xf32> from vector<[4]x[4]xf32>
+  // FULL-UNROLL-NOT: vector.extract {{.*}} : vector<[4]xi1> from vector<[4]x[4]xi1>
+  // FULL-UNROLL: vector.transfer_write {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %v = arith.constant dense<10.0> : vector<[4]x[4]xf32>
+  %dim_a = memref.dim %arg0, %c0 : memref<?x?xf32>
+  %dim_b = memref.dim %arg0, %c1 : memref<?x?xf32>
+  %mask = vector.create_mask %dim_a, %dim_b : vector<[4]x[4]xi1>
+  vector.transfer_write %v, %arg0[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32>
+  return
+}

``````````

</details>


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