[Mlir-commits] [mlir] [mlir][mem2reg] Promote whole-buffer memref to a vector SSA value (PR #211880)

[Renato Golin]

================
@@ -0,0 +1,192 @@
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+
+// A memref that is only ever accessed as a whole buffer through
+// vector.transfer_read / vector.transfer_write is promoted to a single vector
+// SSA value.
+
+// CHECK-LABEL: func.func @whole_buffer_write_read
+//   CHECK-SAME:   (%[[PAD:.*]]: f32)
+//    CHECK-NOT:   memref.alloca
+//    CHECK-NOT:   vector.transfer_write
+//    CHECK-NOT:   vector.transfer_read
+//        CHECK:   %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
+//        CHECK:   return %[[CST]] : vector<4xf32>
+func.func @whole_buffer_write_read(%pad: f32) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant dense<1.0> : vector<4xf32>
+  %a = memref.alloca() : memref<4xf32>
+  vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+  %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
+  return %r : vector<4xf32>
+}
+
+// -----
+
+// A whole-buffer slot carried across scf.for is threaded as an iter_arg/result.
+
+// CHECK-LABEL: func.func @whole_buffer_in_loop
+//    CHECK-NOT:   memref.alloca
+//    CHECK-NOT:   vector.transfer_write
+//    CHECK-NOT:   vector.transfer_read
+//        CHECK:   %[[RES:.*]] = scf.for {{.*}} iter_args(%[[IT:.*]] = %{{.*}}) -> (vector<4xf32>)
+//        CHECK:     %[[NEXT:.*]] = arith.addf %[[IT]], %[[IT]] : vector<4xf32>
+//        CHECK:     scf.yield %[[NEXT]] : vector<4xf32>
+//        CHECK:   return %[[RES]] : vector<4xf32>
+func.func @whole_buffer_in_loop(%pad: f32, %lb: index, %ub: index, %step: index) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant dense<1.0> : vector<4xf32>
+  %a = memref.alloca() : memref<4xf32>
+  vector.transfer_write %cst, %a[%c0] {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
+  scf.for %i = %lb to %ub step %step {
+    %v = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
----------------
rengolin wrote:

This is a bit of noop, since the `%step` is never used, could probably be a candidate to hoisting / DCE, no?

Do you have an example where the internal read/write is still a subset, and then it becomes a subview?

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