[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)

[Andrzej Warzyński]

================
@@ -108,61 +91,167 @@ func.func private @concat_dynamic_stride(%stride0: index,
 
   // CHECK-NOT:  memref.copy
   // CHECK:      %[[C0:.*]] = arith.constant 0 : index
-  // CHECK:      %[[C0_0:.*]] = arith.constant 0 : index
   // CHECK:      %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]]] : memref<1x1xf32>
-  /// Dynamic offset used in store
-  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0_0]]] : memref<1x108xf32>
+  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]]] : memref<1x108xf32>
   memref.copy %src, %reinterpret_cast
     : memref<1x1xf32>
       to memref<1x1xf32, strided<[?, ?]>>
+  // CHECK-NOT:  memref.copy
   return
 }
 
-// CHECK-LABEL: func.func private @concat_rank1(
-// CHECK-SAME:   %[[SRC:.*]]: memref<1xf32>
-// CHECK-SAME:   %[[DST:.*]]: memref<108xf32>
-func.func private @concat_rank1(%src : memref<1xf32>, %dst : memref<108xf32>) {
+// CHECK-LABEL: func.func private @copy_to_strided_rank0(
+// CHECK-SAME:   %[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>
+func.func private @copy_to_strided_rank0(%src : memref<f32>, %dst : memref<f32>) {
   // CHECK-NOT:  memref.reinterpret_cast
   %reinterpret_cast = memref.reinterpret_cast %dst
-    to offset: [0], sizes: [1], strides: [1]
-    : memref<108xf32> to memref<1xf32>
+    to offset: [0], sizes: [], strides: []
+    : memref<f32> to memref<f32>
 
   // CHECK-NOT:  memref.copy
-  // CHECK:      %[[C0:.*]] = arith.constant 0 : index
-  // CHECK:      %[[C0_0:.*]] = arith.constant 0 : index
-  // CHECK:      %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]]] : memref<1xf32>
-  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0_0]]] : memref<108xf32>
-  memref.copy %src, %reinterpret_cast
-    : memref<1xf32> to memref<1xf32>
+  // CHECK:      %[[VAL:.*]] = memref.load %[[SRC]][] : memref<f32>
+  // CHECK:      memref.store %[[VAL]], %[[DST]][] : memref<f32>
+  memref.copy %src, %reinterpret_cast : memref<f32> to memref<f32>
+  // CHECK-NOT:  memref.copy
   return
 }
 
-// CHECK-LABEL: func.func private @concat_rank3(
+// CHECK-LABEL: func.func private @copy_to_strided_0d_base_2d(
 // CHECK-SAME:   %[[SRC:.*]]: memref<1x1x1xf32>
-// CHECK-SAME:   %[[DST:.*]]: memref<1x1x108xf32>
-func.func private @concat_rank3(%src : memref<1x1x1xf32>,
-  %dst : memref<1x1x108xf32>) {
+// CHECK-SAME:   %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @copy_to_strided_0d_base_2d(
+  %src : memref<1x1x1xf32>, %dst : memref<1x33x42xf32>) {
   // CHECK-NOT:  memref.reinterpret_cast
   %reinterpret_cast = memref.reinterpret_cast %dst
     to offset: [0], sizes: [1, 1, 1], strides: [1, 1, 1]
-    : memref<1x1x108xf32> to memref<1x1x1xf32>
-
+    : memref<1x33x42xf32>
+      to memref<1x1x1xf32>
   // CHECK-NOT:  memref.copy
   // CHECK:      %[[C0:.*]] = arith.constant 0 : index
-  // CHECK:      %[[C0_0:.*]] = arith.constant 0 : index
   // CHECK:      %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
-  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0_0]]] : memref<1x1x108xf32>
+  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
   memref.copy %src, %reinterpret_cast
     : memref<1x1x1xf32> to memref<1x1x1xf32>
+  // CHECK-NOT:  memref.copy
+  return
+}
+
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_zero_offset(
+// CHECK-SAME:   %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME:   %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_1d_vector_zero_offset(
+  %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+  // CHECK-NOT:  memref.reinterpret_cast
+  %reinterpret_cast = memref.reinterpret_cast %dst
----------------
banach-space wrote:

Here and in other places, please avoid using op name (`memref.reinterpret_cast`) as variable names (`reinterpret_cast`). Given the number of changes in this PR, the cognitive load is very high.

Why not use `%rc` or `%dst`? That would be more readable and self-descriptive.

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