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

[Cullen Rhodes]

================
@@ -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>
----------------
c-rhodes wrote:

is the mask necessary? Could this be simplified by removing it or passing it as an arg (likewise for `%v`)?

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