[Mlir-commits] [mlir] 42605e7 - [mlir][memref] Support truncating copied dims to non-unit in reinterpret_cast elision (#209536)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 15 05:49:34 PDT 2026
Author: ioana ghiban
Date: 2026-07-15T14:49:28+02:00
New Revision: 42605e7b341ccbc6b2486e2ced82386f099c476a
URL: https://github.com/llvm/llvm-project/commit/42605e7b341ccbc6b2486e2ced82386f099c476a
DIFF: https://github.com/llvm/llvm-project/commit/42605e7b341ccbc6b2486e2ced82386f099c476a.diff
LOG: [mlir][memref] Support truncating copied dims to non-unit in reinterpret_cast elision (#209536)
Relax the `memref.copy`through `memref.reinterpret_cast` rewrite to
support same-dimension slices where strictly one non-unit source
dimension size is truncated to a smaller non-unit result dimension.
Previously, the rewrite only supported truncating to unit size:
```mlir
memref<1xMxNxf32>
to memref<1xMx1xf32, strided<[M*N, N, 1], offset: OFF>>
```
This change relaxes constraints to also support:
```mlir
memref<1xMxNxf32>
to memref<1xMxKxf32, strided<[M*N, N, 1], offset: OFF>>
```
where `K < N`.
The rewrite remains conservative:
* source must have identity layout,
* result strides must match the source identity strides,
* offset must be static for non-scalar results, and
* copied slice must be in bounds.
Assisted-by: Codex (refine PR description).
Added:
Modified:
mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index a24a0ccaa3d67..49f30a594e0ee 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -84,7 +84,7 @@ delinearizeStaticRCOffset(memref::ReinterpretCastOp rc) {
return offsetIdxs;
}
-static bool hasExactlyOneCollapsedNonUnitDim(memref::ReinterpretCastOp rc) {
+static bool hasExactlyOneTruncatedNonUnitDim(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
MemRefType resType = dyn_cast<MemRefType>(rc.getType());
assert(srcType.hasStaticShape() && resType.hasStaticShape() &&
@@ -92,25 +92,26 @@ static bool hasExactlyOneCollapsedNonUnitDim(memref::ReinterpretCastOp rc) {
assert(srcType.getRank() == resType.getRank() &&
"expected rank-preserving reinterpret_cast");
- unsigned collapsedDims = 0;
+ unsigned truncatedDims = 0;
for (auto [srcSize, resSize] :
llvm::zip_equal(srcType.getShape(), resType.getShape())) {
if (srcSize == resSize)
continue;
- // Only allow collapsing one non-unit source dim to a unit result dim.
- if (srcSize != 1 && resSize == 1) {
- ++collapsedDims;
+ // Only one non-unit source dimension may be truncated.
+ if (srcSize != 1 && resSize < srcSize) {
+ ++truncatedDims;
continue;
}
- // The sizes
diff er and both of them are non-unit - ATM not supported.
+ // The size change is not a supported single non-unit source dimension
+ // reduction.
return false;
}
- // Make sure there is only one collapsed dimension.
- return collapsedDims == 1;
+ // Make sure there is only one truncated dimension.
+ return truncatedDims == 1;
}
/// Returns the unique non-unit dim or nullopt if # non-unit-dims != 1.
@@ -134,21 +135,19 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// offsets, delinearized offset.
///
/// Supports ranked, static-shape, rank-preserving reinterpret_casts from
-/// identity-layout sources. In addition:
-/// identical to the source identity strides, and exactly one non-unit
-/// source
-/// * Non-scalar results must have static offsets, static result strides
-/// dimension collapsed to unit size
-/// Scalar-shaped results may have arbitrary result strides (i.e. for scalars,
-/// strides are effectively irrelevant).
+/// identity-layout sources.
+/// * Scalar-shaped results may have arbitrary result strides.
+/// * Non-scalar results must have static offsets, static result strides
+/// identical to the source identity strides, and exactly one non-unit
+/// source dimension size truncated.
///
-/// Returns nullopt for unsupported
-/// reinterpret_casts.
+/// Returns nullopt for unsupported reinterpret_casts.
///
/// Examples that return info:
///
/// reinterpret_cast memref<1xMxNxf32, identity-layout>
-/// to memref<1xMx1xf32, strided<[M*N, N, 1], offset: OFF>>
+/// to memref<1xMxKxf32, strided<[M*N, N, 1], offset: OFF>>
+/// where K < N
///
/// reinterpret_cast memref<1xMxf32, identity-layout>
/// to memref<1x1xf32, strided<[?, ?], offset: ?>>
@@ -156,10 +155,11 @@ static std::optional<unsigned> getSingleNonUnitDim(MemRefType type) {
/// Examples that return no info:
///
/// reinterpret_cast memref<1xMxNxf32, identity-layout>
-/// to memref<1xMx1xf32, strided<[?, N, 1]>>
+/// to memref<1xMxKxf32, strided<[?, N, 1]>>
///
/// reinterpret_cast memref<1xMxNxf32, identity-layout>
-/// to memref<1xKx1xf32, strided<[M*N, N, 1], offset: OFF>>
+/// to memref<1xNxPxf32, strided<[M*N, N, 1], offset: OFF>>
+/// where M != N && N != P
static std::optional<ResultNonUnitDimsAndOffsetsForRC>
getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
MemRefType srcType = dyn_cast<MemRefType>(rc.getSource().getType());
@@ -221,7 +221,7 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
}))
return std::nullopt;
- if (!hasExactlyOneCollapsedNonUnitDim(rc))
+ if (!hasExactlyOneTruncatedNonUnitDim(rc))
return std::nullopt;
}
@@ -255,9 +255,9 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
/// the store index is derived from the reinterpret_cast offset.
///
/// 2. Non-scalar reinterpret_cast results that preserve all non-unit source
-/// dimensions except one collapsed-to-unit dimension. Result strides must
-/// be static and identical to the identity strides of the source, and the
-/// static offset selects the collapsed dimension.
+/// dimensions sizes except one. Result
+/// strides must be static and identical to the identity strides of the
+/// source, and the offset must be static.
///
/// // BEFORE (scalar-shaped result)
/// %strided = memref.reinterpret_cast %dst
@@ -268,18 +268,20 @@ getResultNonUnitDimsAndOffsetsForRC(memref::ReinterpretCastOp rc) {
/// %v = memref.load %src[0, ..., 0]
/// memref.store %v, %dst[delinearized(OFF)]
///
-/// // BEFORE (one collapsed non-unit dimension)
+/// // BEFORE (one truncated non-unit dimension)
/// %strided = memref.reinterpret_cast %dst
-/// to offset: [OFF], sizes: [1, M, 1], strides: [M*N, N, 1]
+/// to offset: [OFF], sizes: [1, M, K], strides: [M*N, N, 1]
/// : memref<1xMxNxf32>
-/// to memref<1xMx1xf32, strided<[M*N, N, 1], offset: OFF>>
+/// to memref<1xMxKxf32, strided<[M*N, N, 1], offset: OFF>>
/// memref.copy %src, %strided
///
/// // AFTER
-/// // Assuming OFF delinearizes to [0, 0, OFF]:
+/// // Assuming OFF delinearizes to [0, 0, DELIN_OFF]:
/// scf.for %i = 0 to M step 1 {
-/// %v = memref.load %src[0, %i, 0]
-/// memref.store %v, %dst[0, %i, OFF]
+/// scf.for %k = 0 to K step 1 {
+/// %v = memref.load %src[0, %i, %k]
+/// memref.store %v, %dst[0, %i, DELIN_OFF + %k]
+/// }
/// }
struct CopyToLoadAndStore : public OpRewritePattern<memref::CopyOp> {
public:
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 0955d9638f093..f305de1ccd0d8 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -262,20 +262,27 @@ func.func private @negative_copy_scalar_into_2D_strided_dynamic_offset(
// Non-scalar (ND) copy
//===----------------------------------------------------------------------===//
-/// The result of collapsing dimensions is not a unit dimension.
-// CHECK-LABEL: func.func private @negative_copy_1D_into_1D_strided(
-func.func private @negative_copy_1D_into_1D_strided(
+// CHECK-LABEL: func.func private @copy_1D_into_1D_strided_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<108xf32>
+func.func private @copy_1D_into_1D_strided_zero_offset(
%src : memref<4xf32>, %dst : memref<108xf32>) {
- // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ // CHECK-NOT: memref.reinterpret_cast
%rc = memref.reinterpret_cast %dst
to offset: [0], sizes: [4], strides: [1]
: memref<108xf32> to memref<4xf32, strided<[1]>>
- // CHECK: memref.copy %arg0, %reinterpret_cast
- // CHECK-NOT: memref.load
- // CHECK-NOT: memref.store
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX:.*]] = %[[C0]] to %[[UB]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[IDX]]] : memref<4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[IDX]]] : memref<108xf32>
+ // CHECK: }
memref.copy %src, %rc
: memref<4xf32> to memref<4xf32, strided<[1]>>
+ // CHECK-NOT: memref.copy
return
}
@@ -418,9 +425,8 @@ func.func private @negative_copy_1D_into_2D_strided_zero_offset_dynamic_stride(%
return
}
-/// 4 is collapsed to a non-unit dimension.
-// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_strided_
diff _dim_sizes(
-func.func private @negative_copy_1D_into_2D_strided_
diff _dim_sizes(
+// CHECK-LABEL: func.func private @negative_copy_1D_into_2D_multiple_truncated_dims(
+func.func private @negative_copy_1D_into_2D_multiple_truncated_dims(
%src : memref<1x3x1xf32>, %dst : memref<1x4x11xf32>) {
// CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
%rc = memref.reinterpret_cast %dst
@@ -437,6 +443,113 @@ func.func private @negative_copy_1D_into_2D_strided_
diff _dim_sizes(
return
}
+// CHECK-LABEL: func.func private @negative_copy_into_strided_no_truncated_dims(
+func.func private @negative_copy_into_strided_no_truncated_dims(%src : memref<3x4xf32>,
+ %dst : memref<3x4xf32>) {
+ // CHECK: %reinterpret_cast = memref.reinterpret_cast %arg1
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [3, 4], strides: [12, 1]
+ : memref<3x4xf32> to memref<3x4xf32, strided<[12, 1]>>
+
+ // CHECK: memref.copy %arg0, %reinterpret_cast
+ // CHECK-NOT: memref.load
+ // CHECK-NOT: memref.store
+ memref.copy %src, %rc
+ : memref<3x4xf32> to memref<3x4xf32, strided<[12, 1]>>
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_zero_offset(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_2D_into_2D_strided_zero_offset(
+ %src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [0], sizes: [1, 3, 4], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1]>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x11xf32>
+ // CHECK: }
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1]>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_nonzero_offset_delinearized_v1(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x4xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x3x11xf32>
+func.func private @copy_2D_into_2D_strided_nonzero_offset_delinearized_v1(
+ %src : memref<1x3x4xf32>, %dst : memref<1x3x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [6], sizes: [1, 3, 4], strides: [33, 11, 1]
+ : memref<1x3x11xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: 6>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 6 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX1]] : index
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x4xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[IDX0]], %[[DST_IDX]]] : memref<1x3x11xf32>
+ // CHECK: }
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x3x4xf32>
+ to memref<1x3x4xf32, strided<[33, 11, 1], offset: 6>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
+// CHECK-LABEL: func.func private @copy_2D_into_2D_strided_nonzero_offset_delinearized_v2(
+// CHECK-SAME: %[[SRC:.*]]: memref<1x3x11xf32>
+// CHECK-SAME: %[[DST:.*]]: memref<1x10x11xf32>
+func.func private @copy_2D_into_2D_strided_nonzero_offset_delinearized_v2(
+ %src : memref<1x3x11xf32>, %dst : memref<1x10x11xf32>) {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %dst
+ to offset: [44], sizes: [1, 3, 11], strides: [110, 11, 1]
+ : memref<1x10x11xf32>
+ to memref<1x3x11xf32, strided<[110, 11, 1], offset: 44>>
+
+ // CHECK-NOT: memref.copy
+ // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK-DAG: %[[UB0:.*]] = arith.constant 3 : index
+ // CHECK-DAG: %[[UB1:.*]] = arith.constant 11 : index
+ // CHECK-DAG: %[[OFF:.*]] = arith.constant 4 : index
+ // CHECK: scf.for %[[IDX0:.*]] = %[[C0]] to %[[UB0]] step %[[C1]] {
+ // CHECK: %[[DST_IDX:.*]] = arith.addi %[[OFF]], %[[IDX0]] : index
+ // CHECK: scf.for %[[IDX1:.*]] = %[[C0]] to %[[UB1]] step %[[C1]] {
+ // CHECK: %[[VAL:.*]] = memref.load %[[SRC]][%[[C0]], %[[IDX0]], %[[IDX1]]] : memref<1x3x11xf32>
+ // CHECK: memref.store %[[VAL]], %[[DST]][%[[C0]], %[[DST_IDX]], %[[IDX1]]] : memref<1x10x11xf32>
+ // CHECK: }
+ // CHECK: }
+ memref.copy %src, %rc
+ : memref<1x3x11xf32>
+ to memref<1x3x11xf32, strided<[110, 11, 1], offset: 44>>
+ // CHECK-NOT: memref.copy
+ return
+}
+
// CHECK-LABEL: func.func private @copy_2D_into_3D_strided_zero_offset(
// CHECK-SAME: %[[SRC:.*]]: memref<3x1x4x1xf32>
// CHECK-SAME: %[[DST:.*]]: memref<3x1x4x11xf32>
More information about the Mlir-commits
mailing list