From 6b76b03d518b54eb74c8aeaaa614b5a5f86c105f Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski@arm.com>
Date: Wed, 7 Aug 2024 21:12:16 +0100
Subject: [PATCH 1/2] [mlir][vector] Add tests for scalable vectors in
 bufferize-invalid.mlir

---
 mlir/test/Dialect/Vector/bufferize-invalid.mlir | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlir/test/Dialect/Vector/bufferize-invalid.mlir b/mlir/test/Dialect/Vector/bufferize-invalid.mlir
index bcca50a0fe79a6..12ee0d57281705 100644
--- a/mlir/test/Dialect/Vector/bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Vector/bufferize-invalid.mlir
@@ -6,3 +6,12 @@ func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vect
   %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @mask_scalable(
+func.func @mask_scalable(%t0: tensor<?xf32>, %val: vector<[16]xf32>, %idx: index, %m0: vector<[16]xi1>) -> tensor<?xf32> {
+  // expected-error @+1 {{'vector.mask' op body must bufferize in-place}}
+  %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<[16]xf32>, tensor<?xf32> } : vector<[16]xi1> -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}

From 61e1705708c24c04cbd9ad9f6577d027973de7e2 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski@arm.com>
Date: Fri, 9 Aug 2024 10:50:54 +0100
Subject: [PATCH 2/2] fixup! [mlir][vector] Add tests for scalable vectors in
 bufferize-invalid.mlir

Remove negative test, add positive test
---
 mlir/test/Dialect/Vector/bufferize-invalid.mlir  |  9 ---------
 mlir/test/Dialect/Vector/one-shot-bufferize.mlir | 11 +++++++++++
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/mlir/test/Dialect/Vector/bufferize-invalid.mlir b/mlir/test/Dialect/Vector/bufferize-invalid.mlir
index 12ee0d57281705..bcca50a0fe79a6 100644
--- a/mlir/test/Dialect/Vector/bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Vector/bufferize-invalid.mlir
@@ -6,12 +6,3 @@ func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vect
   %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
-
-// -----
-
-// CHECK-LABEL: func @mask_scalable(
-func.func @mask_scalable(%t0: tensor<?xf32>, %val: vector<[16]xf32>, %idx: index, %m0: vector<[16]xi1>) -> tensor<?xf32> {
-  // expected-error @+1 {{'vector.mask' op body must bufferize in-place}}
-  %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<[16]xf32>, tensor<?xf32> } : vector<[16]xi1> -> tensor<?xf32>
-  return %0 : tensor<?xf32>
-}
diff --git a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
index 64238c3c08a6f1..c2d699b9b013a1 100644
--- a/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Vector/one-shot-bufferize.mlir
@@ -12,6 +12,17 @@ func.func @mask(%t0: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vect
   return %0 : tensor<?xf32>
 }
 
+// CHECK-LABEL: func @mask_scalable(
+//  CHECK-SAME:     %[[t0:.*]]: memref<?xf32, strided<[?], offset: ?>>
+func.func @mask_scalable(%t0: tensor<?xf32>, %val: vector<[16]xf32>, %idx: index, %m0: vector<[16]xi1>) -> tensor<?xf32> {
+  // CHECK-NOT: alloc
+  // CHECK-NOT: copy
+  //     CHECK: vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[t0]][%{{.*}}] : vector<[16]xf32>, memref<?xf32, strided<[?], offset: ?>> } : vector<[16]xi1>
+  %0 = vector.mask %m0 { vector.transfer_write %val, %t0[%idx] : vector<[16]xf32>, tensor<?xf32> } : vector<[16]xi1> -> tensor<?xf32>
+  //     CHECK: return %[[t0]]
+  return %0 : tensor<?xf32>
+}
+
 // -----
 
 // CHECK-ANALYSIS-LABEL: func @non_reading_xfer_write(