[Mlir-commits] [mlir] [memref] Support non-scalar copies in `reinterpret_cast` elision (PR #203873)
Andrzej Warzyński
llvmlistbot at llvm.org
Thu Jun 25 07:53:46 PDT 2026
================
@@ -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
+ to offset: [0], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x11xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1]>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x1xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_1d_vector_nonzero_offset(
+ %src : memref<1x3x1xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [10], sizes: [1, 3, 1], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX]], %[[C0]]] : memref<1x3x1xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX]], %[[C10]]] : memref<1x3x11xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<1x3x1xf32>
+ to memref<1x3x1xf32, strided<[33, 11, 1], offset: 10>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<42xf32>
+func.func private @copy_to_strided_1d_vector_1d_base_dynamic_offset(
+ %offset : index, %src : memref<4xf32>, %dst : memref<42xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %reinterpret_cast = memref.reinterpret_cast %dst
+ to offset: [%offset], sizes: [4], strides: [1]
+ : memref<42xf32> to memref<4xf32, strided<[1], offset: ?>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[DST_IDX]]] : memref<42xf32>
+ // CHECK: }
+ memref.copy %src, %reinterpret_cast
+ : memref<4xf32> to memref<4xf32, strided<[1], offset: ?>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_to_strided_2d_vector_nonzero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_to_strided_2d_vector_nonzero_offset(
----------------
banach-space wrote:
The name of this function is confusing 😅 Lets look at this example:
```
memref.copy %src, %reinterpret_cast
: memref<1x3x4xf32>
to memref<1x3x4xf32, strided<[33, 11, 1], offset: 7>>
```
There are only rank-3 `memref`s here and no `vector`s (I mean MLIR `vector`s). I suggest:
* `@copy_to_strided_2d_vector_nonzero_offset` -> `@copy_2D_into_2D_strided` (or `@copy_2D_array_into_2D_array_strided` - there's no `array` type in MLIR, so this would be safe, but much longer).
Similar comment for other tests. I would also add a block comment somewhere explaining what `2D` means in this context (i.e. it's effectively a 2D array within a MemRef with rank >= 2).
https://github.com/llvm/llvm-project/pull/203873
More information about the Mlir-commits
mailing list