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

[ioana ghiban]

================
@@ -149,14 +146,165 @@ func.func private @concat_rank3(%src : 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<1x1x108xf32>
   memref.copy %src, %reinterpret_cast
     : memref<1x1x1xf32> to memref<1x1x1xf32>
   return
 }
 
+// CHECK-LABEL: func.func private @concat_0d(
+// CHECK-SAME:   %[[SRC:.*]]: memref<1x1x1xf32>
+// CHECK-SAME:   %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_0d(
+  %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<1x33x42xf32>
+      to memref<1x1x1xf32>
+  // CHECK-NOT:  memref.copy
+  // CHECK:      %[[C0:.*]] = arith.constant 0 : index
+  // CHECK:      %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x1x1xf32>
+  // CHECK:      memref.store %[[VAL]], %[[DST]][%[[C0]], %[[C0]], %[[C0]]] : memref<1x33x42xf32>
+  memref.copy %src, %reinterpret_cast
+    : memref<1x1x1xf32> to memref<1x1x1xf32>
+  return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_zero_offset(
+// CHECK-SAME:   %[[SRC:.*]]: memref<1x33x1xf32>
+// CHECK-SAME:   %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_1d_vector_zero_offset(
+  %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+  // CHECK-NOT:  memref.reinterpret_cast
+  %reinterpret_cast = memref.reinterpret_cast %dst
+    to offset: [0], sizes: [1, 33, 1], strides: [1386, 42, 1]
+    : memref<1x33x42xf32>
+      to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+
+  // CHECK-NOT:  memref.copy
+  // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG:  %[[C33:.*]] = arith.constant 33 : index
+  // CHECK:      scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+  // CHECK:        %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
+  // CHECK:        memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x42xf32>
+  // CHECK:      }
+  memref.copy %src, %reinterpret_cast
+    : memref<1x33x1xf32>
+      to memref<1x33x1xf32, strided<[1386, 42, 1]>>
+  return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_nonzero_offset(
+// CHECK-SAME:   %[[SRC:.*]]: memref<1x33x1xf32>
+// CHECK-SAME:   %[[DST:.*]]: memref<1x33x42xf32>
+func.func private @concat_1d_vector_nonzero_offset(
+  %src : memref<1x33x1xf32>, %dst : memref<1x33x42xf32>) {
+  // CHECK-NOT:  memref.reinterpret_cast
+  %reinterpret_cast = memref.reinterpret_cast %dst
+    to offset: [41], sizes: [1, 33, 1], strides: [1386, 42, 1]
+    : memref<1x33x42xf32>
+      to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+
+  // CHECK-NOT:  memref.copy
+  // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG:  %[[C33:.*]] = arith.constant 33 : index
+  // CHECK-DAG:  %[[C41:.*]] = arith.constant 41 : index
+  // CHECK:      scf.for %[[IDX:.*]] = %[[C0]] to %[[C33]] step %[[C1]] {
+  // CHECK:        %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x33x1xf32>
+  // CHECK:        memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C41]]] : memref<1x33x42xf32>
+  // CHECK:      }
+  memref.copy %src, %reinterpret_cast
+    : memref<1x33x1xf32>
+      to memref<1x33x1xf32, strided<[1386, 42, 1], offset: 41>>
+  return
+}
+
+// CHECK-LABEL: func.func private @concat_1d_vector_dynamic_offset_same_dim(
+// CHECK-SAME:   %[[OFF:.*]]: index
+// CHECK-SAME:   %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME:   %[[DST:.*]]: memref<42xf32>
+func.func private @concat_1d_vector_dynamic_offset_same_dim(
----------------
ioghiban wrote:

Same dimension as the looped dimension. I'll make that clearer.

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