[Mlir-commits] [mlir] [mlir][Vector] Promote buffers accessed via dynamic subviews and memref.copy (PR #217793)

[Jianhui Li]

================
@@ -0,0 +1,255 @@
+// RUN: mlir-opt %s -mem2reg -split-input-file | FileCheck %s
+
+// Reading a static buffer through a DYNAMIC subview: the read returns the value
+// that was written, and takes the transfer's padding wherever it reads past the
+// subview's extent.
+
+// CHECK-LABEL: func.func @read_dyn_subview(
+// CHECK-SAME:      %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[PAD:.*]]: f32
+// CHECK-NOT:     memref.alloca
+// CHECK-NOT:     memref.subview
+// CHECK-NOT:     vector.transfer_read
+// CHECK-NOT:     vector.transfer_write
+// CHECK:         %[[MASK:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK:         %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8x16xf32>
+// CHECK:         %[[SEL:.*]] = arith.select %[[MASK]], %[[V]], %[[PS]]
+// CHECK:         return %[[SEL]]
+func.func @read_dyn_subview(%v: vector<8x16xf32>, %n: index, %pad: f32)
+    -> vector<8x16xf32> {
+  %c0 = arith.constant 0 : index
+  %a = memref.alloca() : memref<8x16xf32>
+  vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]}
+      : vector<8x16xf32>, memref<8x16xf32>
+  %sv = memref.subview %a[0, 0] [8, %n] [1, 1]
+      : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+  %r = vector.transfer_read %sv[%c0, %c0], %pad {in_bounds = [true, false]}
+      : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32>
+  return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Write through a dynamic subview, then read the whole buffer. The write updates
+// only the elements within the dynamic extent, so the read returns the written
+// value for those and the buffer's earlier value for the rest.
+
+// CHECK-LABEL: func.func @write_then_read_dyn_subview(
+// CHECK-SAME:      %[[V:.*]]: vector<8x16xf32>, %[[W:.*]]: vector<8x16xf32>, %[[N:.*]]: index
+// CHECK-NOT:     memref.alloca
+// CHECK:         %[[MASK:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK:         %[[SEL:.*]] = arith.select %[[MASK]], %[[W]], %[[V]]
+// CHECK:         return %[[SEL]]
+func.func @write_then_read_dyn_subview(%v: vector<8x16xf32>, %w: vector<8x16xf32>,
+                                       %n: index, %pad: f32) -> vector<8x16xf32> {
+  %c0 = arith.constant 0 : index
+  %a = memref.alloca() : memref<8x16xf32>
+  vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]}
+      : vector<8x16xf32>, memref<8x16xf32>
+  %sv = memref.subview %a[0, 0] [8, %n] [1, 1]
+      : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+  vector.transfer_write %w, %sv[%c0, %c0] {in_bounds = [true, false]}
+      : vector<8x16xf32>, memref<8x?xf32, strided<[16, 1]>>
+  %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]}
+      : memref<8x16xf32>, vector<8x16xf32>
+  return %r : vector<8x16xf32>
+}
+
+// -----
+
+// A masked read through a dynamic subview combines both masks: the subview
+// extent (create_mask) AND the transfer's own mask.
+// CHECK-LABEL: func.func @masked_read_dyn_subview(
+// CHECK-SAME:      %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[M:.*]]: vector<8x16xi1>, %[[PAD:.*]]: f32
+// CHECK-NOT:     memref.alloca
+// CHECK:         %[[CM:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK:         %[[AND:.*]] = arith.andi %[[CM]], %[[M]]
+// CHECK:         %[[SEL:.*]] = arith.select %[[AND]], %[[V]], %{{.*}}
+// CHECK:         return %[[SEL]]
+func.func @masked_read_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vector<8x16xi1>, %pad: f32) -> vector<8x16xf32> {
+  %c0 = arith.constant 0 : index
+  %a = memref.alloca() : memref<8x16xf32>
+  vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+  %sv = memref.subview %a[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+  %r = vector.transfer_read %sv[%c0, %c0], %pad, %m {in_bounds = [true, false]} : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32>
----------------
Jianhui-Li wrote:

I changed the test names accordingly. 

The masked case is not supported since vector.mask op doesn't support PromotableRegionOpInterface. Added negative tests

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