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

[Benjamin Maxwell]

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

>From 897d3c6b76bb3c020d706fc3e02eda4cebf1fd3d Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 10 Nov 2023 11:10:23 +0000
Subject: [PATCH 1/2] [mlir][Vector] Don't fully unroll transfer_reads of n-D
 scalable vectors

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).
---
 mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp   |  5 +++++
 .../Conversion/VectorToSCF/vector-to-scf.mlir     | 15 +++++++++++++++
 2 files changed, 20 insertions(+)

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
+}

>From d91aac7bcf6668d092d914198c972f7f2ef9fb64 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 10 Nov 2023 13:42:12 +0000
Subject: [PATCH 2/2] Simplify test case

---
 mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index cff85995e44b078..597cc4f71a63961 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -740,16 +740,11 @@ func.func @cannot_lower_transfer_read_with_leading_scalable(%arg0: memref<?x4xf3
 //  -----
 
 // 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>
+func.func @cannot_fully_unroll_transfer_write_of_nd_scalable_vector(%vec: vector<[4]x[4]xf32>, %memref: memref<?x?xf32>) {
+  // FULL-UNROLL-NOT: vector.extract
   // FULL-UNROLL: vector.transfer_write {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>
+  // FULL-UNROLL-NOT: vector.extract
   %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>
+  vector.transfer_write %vec, %memref[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32>
   return
 }