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

[llvmlistbot at llvm.org]

Author: Benjamin Maxwell
Date: 2023-11-13T11:22:55Z
New Revision: ced97ffd0899a3faff796170ed778723efcf8a5b

URL: https://github.com/llvm/llvm-project/commit/ced97ffd0899a3faff796170ed778723efcf8a5b
DIFF: https://github.com/llvm/llvm-project/commit/ced97ffd0899a3faff796170ed778723efcf8a5b.diff

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

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).

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
    mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

Removed: 
    


################################################################################
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..597cc4f71a63961 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -737,4 +737,14 @@ 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(%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
+  vector.transfer_write %vec, %memref[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32>
+  return
+}