[Mlir-commits] [mlir] 2a96d57 - [mlir][memref] Enforce consistent reinterpret_cast metadata (#217338)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 28 02:26:01 PDT 2026
Author: ioana ghiban
Date: 2026-09-28T11:25:47+02:00
New Revision: 2a96d5737524012839e786699a0b327f22d1cfee
URL: https://github.com/llvm/llvm-project/commit/2a96d5737524012839e786699a0b327f22d1cfee
DIFF: https://github.com/llvm/llvm-project/commit/2a96d5737524012839e786699a0b327f22d1cfee.diff
LOG: [mlir][memref] Enforce consistent reinterpret_cast metadata (#217338)
This PR implements [[RFC] Clarify `memref.reinterpret_cast` Verification
of Dynamic
Metadata](https://discourse.llvm.org/t/rfc-clarify-memref-reinterpret-cast-verification-of-dynamic-metadata).
## Context
`memref.reinterpret_cast` describes its result through both its mixed
offset/size/stride metadata and its result `MemRefType`. Currently,
verification allowed a dynamic result-type correspond to static
operation metadata, introducing inconsistent descriptors.
For example, the verifier permits:
```mlir
%r = memref.reinterpret_cast %base
to offset: [0], sizes: [%n, 2], strides: [%s, 1] : memref<f32>
to memref<?x?xf32, strided<[?, ?], offset: ?>>
```
where the operation constructs a descriptor with a statically known
second size and stride, but the result type does not expose those facts.
## Changes
This PR consolidates the verifier to require `reinterpret_cast` result
type’s sizes, offset, and strides to exactly match the static/dynamic
representation and values of the operation metadata.
Transforms that create `memref.reinterpret_cast`, including GPU memref
decomposition and strided-metadata expansion for `subview` and
`collapse_shape`, now derive the result type from the final, potentially
folded descriptor instead of inheriting the original operation’s result
type.
When this descriptor-derived type differs from the original result type,
the rewrite inserts a `memref.cast` back to the original type. These
rewrites are not type-conversion patterns, so their existing
users—including function returns and calls—must retain their original
types. Identity stride layouts are canonicalized when possible to avoid
unnecessary explicit layouts and casts.
Tests are updated to use result types consistent with their
`reinterpret_cast` metadata. In the SPIR-V tests, the size `10` is
represented by an SSA value because those cases intentionally remain
dynamically shaped and lower to element pointers; making the size static
would change the SPIR-V pointee type and prevent the existing
`reinterpret_cast` lowering from applying.
A regression test also verifies that folded `subview` offsets and
strides are reflected in the generated `reinterpret_cast` type.
---
Assisted-by: Codex (refine implementation + tests). I reviewed all code
and tests before submission.
Added:
Modified:
mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
mlir/lib/Dialect/GPU/CMakeLists.txt
mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
mlir/lib/Dialect/Utils/StaticValueUtils.cpp
mlir/lib/IR/BuiltinTypes.cpp
mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
mlir/test/Dialect/GPU/decompose-memrefs.mlir
mlir/test/Dialect/MemRef/canonicalize.mlir
mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
mlir/test/Dialect/MemRef/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index a0ceee6e543eb..d5d3faa5ba0ce 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -285,6 +285,14 @@ struct SaturatedInteger {
int64_t v = 0;
};
+/// Returns a memref type matching the provided offset, size, and stride
+/// metadata. Static attributes become static type metadata; SSA values remain
+/// dynamic. The caller remains responsible for ensuring the metadata is
+/// semantically correct.
+MemRefType updateTypeFromMetadata(MemRefType type, OpFoldResult offset,
+ ArrayRef<OpFoldResult> sizes,
+ ArrayRef<OpFoldResult> strides);
+
} // namespace mlir
#endif // MLIR_DIALECT_UTILS_STATICVALUEUTILS_H
diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 87e09f1dbf47f..6726bf0cab1d2 100644
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/lib/Dialect/GPU/CMakeLists.txt
@@ -67,6 +67,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
MLIRAsyncDialect
MLIRBufferizationDialect
MLIRDataLayoutInterfaces
+ MLIRDialectUtils
MLIRExecutionEngineUtils
MLIRGPUDialect
MLIRGPUUtils
diff --git a/mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp b/mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp
index 9fed82a5a8419..63a5ab85704a9 100644
--- a/mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/DecomposeMemRefs.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/PatternMatch.h"
@@ -216,8 +217,18 @@ struct FlattenSubview : public OpRewritePattern<memref::SubViewOp> {
finalStrides.push_back(strides[i]);
}
- rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
- op, resultType, base, finalOffset, finalSizes, finalStrides);
+ resultType = updateTypeFromMetadata(resultType, finalOffset, finalSizes,
+ finalStrides);
+ auto flattenedSubview = memref::ReinterpretCastOp::create(
+ rewriter, op.getLoc(), resultType, base, finalOffset, finalSizes,
+ finalStrides);
+ if (resultType == op.getType()) {
+ rewriter.replaceOp(op, flattenedSubview);
+ return success();
+ }
+ // Preserve the original result type expected by existing users.
+ rewriter.replaceOpWithNewOp<memref::CastOp>(op, op.getType(),
+ flattenedSubview);
return success();
}
};
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 217e5ce841292..245b2ba6babf0 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2027,6 +2027,20 @@ OpFoldResult RankOp::fold(FoldAdaptor adaptor) {
// ReinterpretCastOp
//===----------------------------------------------------------------------===//
+namespace {
+
+struct PrintDynamicOrValue {
+ int64_t value;
+};
+
+Diagnostic &operator<<(Diagnostic &diag, PrintDynamicOrValue printed) {
+ if (ShapedType::isDynamic(printed.value))
+ return diag << "dynamic";
+ return diag << printed.value;
+}
+
+} // namespace
+
void ReinterpretCastOp::getAsmResultNames(
function_ref<void(Value, StringRef)> setNameFn) {
setNameFn(getResult(), "reinterpret_cast");
@@ -2112,12 +2126,10 @@ LogicalResult ReinterpretCastOp::verify() {
// Match sizes in result memref type and in static_sizes attribute.
for (auto [idx, resultSize, expectedSize] :
llvm::enumerate(resultType.getShape(), getStaticSizes())) {
- if (ShapedType::isStatic(resultSize) && resultSize != expectedSize)
+ if (resultSize != expectedSize)
return emitError("expected result type with size = ")
- << (ShapedType::isDynamic(expectedSize)
- ? std::string("dynamic")
- : std::to_string(expectedSize))
- << " instead of " << resultSize << " in dim = " << idx;
+ << PrintDynamicOrValue{expectedSize} << " instead of "
+ << PrintDynamicOrValue{resultSize} << " in dim = " << idx;
}
// Match offset and strides in static_offset and static_strides attributes. If
@@ -2131,22 +2143,18 @@ LogicalResult ReinterpretCastOp::verify() {
// Match offset in result memref type and in static_offsets attribute.
int64_t expectedOffset = getStaticOffsets().front();
- if (ShapedType::isStatic(resultOffset) && resultOffset != expectedOffset)
+ if (resultOffset != expectedOffset)
return emitError("expected result type with offset = ")
- << (ShapedType::isDynamic(expectedOffset)
- ? std::string("dynamic")
- : std::to_string(expectedOffset))
- << " instead of " << resultOffset;
+ << PrintDynamicOrValue{expectedOffset} << " instead of "
+ << PrintDynamicOrValue{resultOffset};
// Match strides in result memref type and in static_strides attribute.
for (auto [idx, resultStride, expectedStride] :
llvm::enumerate(resultStrides, getStaticStrides())) {
- if (ShapedType::isStatic(resultStride) && resultStride != expectedStride)
+ if (resultStride != expectedStride)
return emitError("expected result type with stride = ")
- << (ShapedType::isDynamic(expectedStride)
- ? std::string("dynamic")
- : std::to_string(expectedStride))
- << " instead of " << resultStride << " in dim = " << idx;
+ << PrintDynamicOrValue{expectedStride} << " instead of "
+ << PrintDynamicOrValue{resultStride} << " in dim = " << idx;
}
return success();
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
index 20d543b7210b1..f95b2529d21e5 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
@@ -19,8 +19,8 @@
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
@@ -197,10 +197,20 @@ struct SubviewFolder : public OpRewritePattern<memref::SubViewOp> {
"failed to resolve subview metadata");
}
- rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
- subview, subview.getType(), stridedMetadata->basePtr,
+ MemRefType resultType = updateTypeFromMetadata(
+ subview.getType(), stridedMetadata->offset, stridedMetadata->sizes,
+ stridedMetadata->strides);
+ auto foldedSubview = memref::ReinterpretCastOp::create(
+ rewriter, subview.getLoc(), resultType, stridedMetadata->basePtr,
stridedMetadata->offset, stridedMetadata->sizes,
stridedMetadata->strides);
+ if (resultType == subview.getType()) {
+ rewriter.replaceOp(subview, foldedSubview);
+ return success();
+ }
+ // Preserve the original result type expected by existing users.
+ rewriter.replaceOpWithNewOp<memref::CastOp>(subview, subview.getType(),
+ foldedSubview);
return success();
}
};
@@ -600,10 +610,22 @@ struct ReshapeFolder : public OpRewritePattern<ReassociativeReshapeLikeOp> {
"failed to resolve reshape metadata");
}
- rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
- reshape, reshape.getType(), stridedMetadata->basePtr,
+ MemRefType resultType = reshape.getResultType();
+ if (isa<memref::CollapseShapeOp>(reshape.getOperation()))
+ resultType = updateTypeFromMetadata(resultType, stridedMetadata->offset,
+ stridedMetadata->sizes,
+ stridedMetadata->strides);
+ auto foldedReshape = memref::ReinterpretCastOp::create(
+ rewriter, reshape.getLoc(), resultType, stridedMetadata->basePtr,
stridedMetadata->offset, stridedMetadata->sizes,
stridedMetadata->strides);
+ if (resultType == reshape.getResultType()) {
+ rewriter.replaceOp(reshape, foldedReshape);
+ return success();
+ }
+ // Preserve the original result type expected by existing users.
+ rewriter.replaceOpWithNewOp<memref::CastOp>(
+ reshape, reshape.getResultType(), foldedReshape);
return success();
}
};
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index b5aac62d5c3a4..ee1086e70597c 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -472,4 +472,23 @@ LogicalResult foldDynamicStrideList(SmallVectorImpl<OpFoldResult> &strides) {
/*onlyNonZero=*/true);
}
+MemRefType updateTypeFromMetadata(MemRefType type, OpFoldResult offset,
+ ArrayRef<OpFoldResult> sizes,
+ ArrayRef<OpFoldResult> strides) {
+ SmallVector<OpFoldResult> offsets{offset};
+ SmallVector<int64_t> staticOffsets = decomposeMixedValues(offsets).first;
+ SmallVector<int64_t> staticSizes = decomposeMixedValues(sizes).first;
+ SmallVector<int64_t> staticStrides = decomposeMixedValues(strides).first;
+ auto layout = StridedLayoutAttr::get(type.getContext(), staticOffsets.front(),
+ staticStrides);
+ // Build a MemRefType using the original element type and memory space,
+ // but derive its shape, offset, and strides from the supplied metadata.
+ MemRefType updatedType = MemRefType::get(staticSizes, type.getElementType(),
+ layout, type.getMemorySpace());
+ if (!type.getLayout().isIdentity())
+ return updatedType;
+ MemRefType canonicalType = updatedType.canonicalizeStridedLayout();
+ return canonicalType.getLayout().isIdentity() ? canonicalType : updatedType;
+}
+
} // namespace mlir
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 36353edf15bc7..845143f863d94 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -8,6 +8,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "TypeDetail.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
diff --git a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
index 1095798483575..6bbbd9ff69761 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
@@ -493,30 +493,33 @@ func.func @reinterpret_cast(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgr
// CHECK: %[[RET:.*]] = spirv.InBoundsPtrAccessChain %[[MEM1]][%[[OFF1]]] : !spirv.ptr<f32, CrossWorkgroup>, i32
// CHECK: %[[RET1:.*]] = builtin.unrealized_conversion_cast %[[RET]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
// CHECK: return %[[RET1]]
- %ret = memref.reinterpret_cast %arg to offset: [%arg1], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+ %c10 = arith.constant 10 : index
+ %ret = memref.reinterpret_cast %arg to offset: [%arg1], sizes: [%c10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
}
// CHECK-LABEL: func.func @reinterpret_cast_0
// CHECK-SAME: (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>)
-func.func @reinterpret_cast_0(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>> {
+func.func @reinterpret_cast_0(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1]>, #spirv.storage_class<CrossWorkgroup>> {
// CHECK-DAG: %[[MEM1:.*]] = builtin.unrealized_conversion_cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to !spirv.ptr<f32, CrossWorkgroup>
-// CHECK-DAG: %[[RET:.*]] = builtin.unrealized_conversion_cast %[[MEM1]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+// CHECK-DAG: %[[RET:.*]] = builtin.unrealized_conversion_cast %[[MEM1]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1]>, #spirv.storage_class<CrossWorkgroup>>
// CHECK: return %[[RET]]
- %ret = memref.reinterpret_cast %arg to offset: [0], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
- return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+ %c10 = arith.constant 10 : index
+ %ret = memref.reinterpret_cast %arg to offset: [0], sizes: [%c10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1]>, #spirv.storage_class<CrossWorkgroup>>
+ return %ret : memref<?xf32, strided<[1]>, #spirv.storage_class<CrossWorkgroup>>
}
// CHECK-LABEL: func.func @reinterpret_cast_5
// CHECK-SAME: (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>)
-func.func @reinterpret_cast_5(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>> {
+func.func @reinterpret_cast_5(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<?xf32, strided<[1], offset: 5>, #spirv.storage_class<CrossWorkgroup>> {
// CHECK: %[[MEM1:.*]] = builtin.unrealized_conversion_cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to !spirv.ptr<f32, CrossWorkgroup>
// CHECK: %[[OFF:.*]] = spirv.Constant 5 : i32
// CHECK: %[[RET:.*]] = spirv.InBoundsPtrAccessChain %[[MEM1]][%[[OFF]]] : !spirv.ptr<f32, CrossWorkgroup>, i32
-// CHECK: %[[RET1:.*]] = builtin.unrealized_conversion_cast %[[RET]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+// CHECK: %[[RET1:.*]] = builtin.unrealized_conversion_cast %[[RET]] : !spirv.ptr<f32, CrossWorkgroup> to memref<?xf32, strided<[1], offset: 5>, #spirv.storage_class<CrossWorkgroup>>
// CHECK: return %[[RET1]]
- %ret = memref.reinterpret_cast %arg to offset: [5], sizes: [10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
- return %ret : memref<?xf32, strided<[1], offset: ?>, #spirv.storage_class<CrossWorkgroup>>
+ %c10 = arith.constant 10 : index
+ %ret = memref.reinterpret_cast %arg to offset: [5], sizes: [%c10], strides: [1] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<?xf32, strided<[1], offset: 5>, #spirv.storage_class<CrossWorkgroup>>
+ return %ret : memref<?xf32, strided<[1], offset: 5>, #spirv.storage_class<CrossWorkgroup>>
}
} // end module
diff --git a/mlir/test/Dialect/GPU/decompose-memrefs.mlir b/mlir/test/Dialect/GPU/decompose-memrefs.mlir
index 12d59dd4f98d5..d7423c7cfe538 100644
--- a/mlir/test/Dialect/GPU/decompose-memrefs.mlir
+++ b/mlir/test/Dialect/GPU/decompose-memrefs.mlir
@@ -88,7 +88,8 @@ func.func @decompose_load(%arg0 : memref<?x?x?xf32>) {
// CHECK-SAME: threads(%[[TX:.*]], %[[TY:.*]], %[[TZ:.*]]) in
// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[TX]], %[[STRIDES]]#0, %[[TY]], %[[STRIDES]]#1, %[[TZ]]]
// CHECK: %[[PTR:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[IDX]]], sizes: [%{{.*}}, %{{.*}}, %{{.*}}], strides: [%[[STRIDES]]#0, %[[STRIDES]]#1, 1]
-// CHECK: "test.test"(%[[PTR]]) : (memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>) -> ()
+// CHECK: %[[CAST:.*]] = memref.cast %[[PTR]] : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+// CHECK: "test.test"(%[[CAST]]) : (memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>) -> ()
func.func @decompose_subview(%arg0 : memref<?x?x?xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index 89efdc31f26e1..b0261c665754a 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1235,13 +1235,13 @@ func.func @reinterpret_of_extract_strided_metadata_w_type_mistach(%arg0 : memref
// same constant value, the match is valid.
// CHECK-LABEL: func @reinterpret_of_extract_strided_metadata_w_constants
// CHECK-SAME: (%[[ARG:.*]]: memref<8x2xf32>)
-// CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] : memref<8x2xf32> to memref<?x?xf32,
+// CHECK: %[[CAST:.*]] = memref.cast %[[ARG]] : memref<8x2xf32> to memref<?x2xf32,
// CHECK: return %[[CAST]]
-func.func @reinterpret_of_extract_strided_metadata_w_constants(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+func.func @reinterpret_of_extract_strided_metadata_w_constants(%arg0 : memref<8x2xf32>) -> memref<?x2xf32, strided<[2, ?], offset: 0>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
%c8 = arith.constant 8: index
- %m2 = memref.reinterpret_cast %base to offset: [0], sizes: [%c8, 2], strides: [2, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
- return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
+ %m2 = memref.reinterpret_cast %base to offset: [0], sizes: [%c8, 2], strides: [2, %strides#1] : memref<f32> to memref<?x2xf32, strided<[2, ?], offset: 0>>
+ return %m2 : memref<?x2xf32, strided<[2, ?], offset: 0>>
}
// -----
@@ -1265,10 +1265,10 @@ func.func @reinterpret_of_extract_strided_metadata_same_type(%arg0 : memref<?x?x
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [0], sizes: [4, 2, 2], strides: [1, 1, 1]
// CHECK: %[[CAST:.*]] = memref.cast %[[RES]]
// CHECK: return %[[CAST]]
-func.func @reinterpret_of_extract_strided_metadata_w_
diff erent_stride(%arg0 : memref<8x2xf32>) -> memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> {
+func.func @reinterpret_of_extract_strided_metadata_w_
diff erent_stride(%arg0 : memref<8x2xf32>) -> memref<4x2x2xf32, strided<[1, 1, ?], offset: ?>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
- %m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [4, 2, 2], strides: [1, 1, %strides#1] : memref<f32> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
- return %m2 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+ %m2 = memref.reinterpret_cast %base to offset: [%offset], sizes: [4, 2, 2], strides: [1, 1, %strides#1] : memref<f32> to memref<4x2x2xf32, strided<[1, 1, ?], offset: ?>>
+ return %m2 : memref<4x2x2xf32, strided<[1, 1, ?], offset: ?>>
}
// -----
@@ -1279,10 +1279,10 @@ func.func @reinterpret_of_extract_strided_metadata_w_
diff erent_stride(%arg0 : me
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [1], sizes: [8, 2], strides: [2, 1]
// CHECK: %[[CAST:.*]] = memref.cast %[[RES]]
// CHECK: return %[[CAST]]
-func.func @reinterpret_of_extract_strided_metadata_w_
diff erent_offset(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+func.func @reinterpret_of_extract_strided_metadata_w_
diff erent_offset(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: 1>> {
%base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<8x2xf32> -> memref<f32>, index, index, index, index, index
- %m2 = memref.reinterpret_cast %base to offset: [1], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
- return %m2 : memref<?x?xf32, strided<[?, ?], offset: ?>>
+ %m2 = memref.reinterpret_cast %base to offset: [1], sizes: [%sizes#0, %sizes#1], strides: [%strides#0, %strides#1] : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: 1>>
+ return %m2 : memref<?x?xf32, strided<[?, ?], offset: 1>>
}
// -----
@@ -1348,12 +1348,12 @@ func.func @reinterpret_cast_with_negative_size_and_offset(%arg0: memref<2x3xf32>
// CHECK-SAME: (%[[ARG:.*]]: memref<2x3xf32>)
// CHECK: %[[NEG:.*]] = arith.constant -1 : index
// CHECK: memref.reinterpret_cast %[[ARG]] to offset: [%[[NEG]]], sizes: [%[[NEG]], %[[NEG]]], strides: [2, 1]
-func.func @reinterpret_cast_no_fold_with_all_negative_size_and_offset(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+func.func @reinterpret_cast_no_fold_with_all_negative_size_and_offset(%arg0: memref<2x3xf32>) -> memref<?x?xf32, strided<[2, 1], offset: ?>> {
%neg = arith.constant -1 : index
%output = memref.reinterpret_cast %arg0 to
offset: [%neg], sizes: [%neg, %neg], strides: [2, 1]
- : memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
- return %output : memref<?x?xf32, strided<[?, ?], offset: ?>>
+ : memref<2x3xf32> to memref<?x?xf32, strided<[2, 1], offset: ?>>
+ return %output : memref<?x?xf32, strided<[2, 1], offset: ?>>
}
// -----
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
index 6db97cfcca0c1..0987111e8e928 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast-load.mlir
@@ -300,9 +300,10 @@ func.func private @negative_dynamic_shape(%dim : index,
// CHECK: %[[RC:.*]] = memref.reinterpret_cast %[[SRC]]
%reinterpret_cast = memref.reinterpret_cast %src
to offset: [0], sizes: [1, %dim], strides: [1, 1]
- : memref<?xf32> to memref<1x?xf32>
+ : memref<?xf32> to memref<1x?xf32, strided<[1, 1]>>
// CHECK: memref.load %[[RC]]
- %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] : memref<1x?xf32>
+ %0 = memref.load %reinterpret_cast[%idx_1, %idx_2] :
+ memref<1x?xf32, strided<[1, 1]>>
return
}
diff --git a/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir b/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
index 70c5e1aee85dc..85ef1b9b515c6 100644
--- a/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
+++ b/mlir/test/Dialect/MemRef/expand-strided-metadata.mlir
@@ -72,6 +72,31 @@ func.func @simplify_subview_all_dynamic(
// -----
+// Check that folded offsets and strides are reflected in the result type of
+// the reinterpret_cast. Sizes remain dynamic because they are forwarded
+// without folding.
+// CHECK-LABEL: func @simplify_subview_folded_metadata
+// CHECK: %[[CAST:.*]] = memref.reinterpret_cast %{{[^ ]+}} to
+// CHECK-SAME: offset: [10],
+// CHECK-SAME: sizes: [%{{[^ ]+}}, %{{[^ ]+}}],
+// CHECK-SAME: strides: [8, 2]
+// CHECK-SAME: memref<f32> to memref<?x?xf32, strided<[8, 2], offset: 10>>
+// CHECK: memref.load %[[CAST]]
+// CHECK-SAME: memref<?x?xf32, strided<[8, 2], offset: 10>>
+func.func @simplify_subview_folded_metadata(%base: memref<4x8xf32>) -> f32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %subview = memref.subview %base[%c1, %c2] [%c2, %c3] [%c1, %c2] :
+ memref<4x8xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
+ %result = memref.load %subview[%c0, %c0] :
+ memref<?x?xf32, strided<[?, ?], offset: ?>>
+ return %result : f32
+}
+
+// -----
+
// Check that we simplify extract_strided_metadata of subview to
// base_buf, base_offset, base_sizes, base_strides = extract_strided_metadata
// strides = base_stride_i * subview_stride_i
diff --git a/mlir/test/Dialect/MemRef/ops.mlir b/mlir/test/Dialect/MemRef/ops.mlir
index b09a5b95e663d..2081d2c8ef1ea 100644
--- a/mlir/test/Dialect/MemRef/ops.mlir
+++ b/mlir/test/Dialect/MemRef/ops.mlir
@@ -130,21 +130,21 @@ func.func @memref_reinterpret_cast(%in: memref<?xf32>)
}
// CHECK-LABEL: func @memref_reinterpret_cast_static_to_dynamic_sizes
-func.func @memref_reinterpret_cast_static_to_dynamic_sizes(%in: memref<?xf32>)
- -> memref<10x?xf32, strided<[?, 1], offset: ?>> {
+func.func @memref_reinterpret_cast_static_to_dynamic_sizes(%in: memref<?xf32>, %size: index)
+ -> memref<10x?xf32, strided<[1, 1], offset: 1>> {
%out = memref.reinterpret_cast %in to
- offset: [1], sizes: [10, 10], strides: [1, 1]
- : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>
- return %out : memref<10x?xf32, strided<[?, 1], offset: ?>>
+ offset: [1], sizes: [10, %size], strides: [1, 1]
+ : memref<?xf32> to memref<10x?xf32, strided<[1, 1], offset: 1>>
+ return %out : memref<10x?xf32, strided<[1, 1], offset: 1>>
}
// CHECK-LABEL: func @memref_reinterpret_cast_dynamic_offset
func.func @memref_reinterpret_cast_dynamic_offset(%in: memref<?xf32>, %offset: index)
- -> memref<10x?xf32, strided<[?, 1], offset: ?>> {
+ -> memref<10x10xf32, strided<[1, 1], offset: ?>> {
%out = memref.reinterpret_cast %in to
offset: [%offset], sizes: [10, 10], strides: [1, 1]
- : memref<?xf32> to memref<10x?xf32, strided<[?, 1], offset: ?>>
- return %out : memref<10x?xf32, strided<[?, 1], offset: ?>>
+ : memref<?xf32> to memref<10x10xf32, strided<[1, 1], offset: ?>>
+ return %out : memref<10x10xf32, strided<[1, 1], offset: ?>>
}
// CHECK-LABEL: func @memref_reshape(
More information about the Mlir-commits
mailing list