Skip to content

Validate type consistency in reintepret cast sizes #140032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Validate type consistency in reintepret cast sizes #140032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
01fa7a2
Validate type consistency in reintepret cast sizes
yaniv217 May 15, 2025
c9814d9
Validate type consistency in reintepret cast offsets
yaniv217 May 21, 2025
148e20e
Merge branch 'main' into fix/validate-reinterpret-cast-sizes
yaniv217 May 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1057,7 +1057,7 @@ struct DimOfMemRefReshape : public OpRewritePattern<DimOp> {
}
} // else dim.getIndex is a block argument to reshape->getBlock and
// dominates reshape
} // Check condition 2
} // Check condition 2
else if (dim->getBlock() != reshape->getBlock() &&
!dim.getIndex().getParentRegion()->isProperAncestor(
reshape->getParentRegion())) {
Expand Down Expand Up @@ -1840,6 +1840,15 @@ LogicalResult ReinterpretCastOp::verify() {
// Match sizes in result memref type and in static_sizes attribute.
for (auto [idx, resultSize, expectedSize] :
llvm::enumerate(resultType.getShape(), getStaticSizes())) {
// Check that dynamic sizes are not mixed with static sizes
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add the same check for the offset and the strides?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes!

if (ShapedType::isDynamic(resultSize) &&
!ShapedType::isDynamic(expectedSize))
return emitError(
"expected size is static, but result type dimension is dynamic ");
if (!ShapedType::isDynamic(resultSize) &&
ShapedType::isDynamic(expectedSize))
return emitError(
"expected size is dynamic, but result type dimension is static ");
if (!ShapedType::isDynamic(resultSize) && resultSize != expectedSize)
return emitError("expected result type with size = ")
<< (ShapedType::isDynamic(expectedSize)
Expand All @@ -1859,6 +1868,15 @@ LogicalResult ReinterpretCastOp::verify() {

// Match offset in result memref type and in static_offsets attribute.
int64_t expectedOffset = getStaticOffsets().front();
// Check that dynamic offset is not mixed with static offset
if (ShapedType::isDynamic(resultOffset) &&
!ShapedType::isDynamic(expectedOffset))
return emitError(
"expected offset is static, but result type offset is dynamic");
if (!ShapedType::isDynamic(resultOffset) &&
ShapedType::isDynamic(expectedOffset))
return emitError(
"expected offset is dynamic, but result type offset is static");
if (!ShapedType::isDynamic(resultOffset) && resultOffset != expectedOffset)
return emitError("expected result type with offset = ")
<< (ShapedType::isDynamic(expectedOffset)
Expand Down Expand Up @@ -2011,7 +2029,7 @@ struct ReinterpretCastOpExtractStridedMetadataFolder
// Second, check the sizes.
if (!llvm::equal(extractStridedMetadata.getConstifiedMixedSizes(),
op.getConstifiedMixedSizes()))
return false;
return false;

// Finally, check the offset.
assert(op.getMixedOffsets().size() == 1 &&
Expand Down
21 changes: 12 additions & 9 deletions mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -339,30 +339,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], offset: 0>, #spirv.storage_class<CrossWorkgroup>>
return %ret : memref<?xf32, strided<[1], offset: 0>, #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
Expand Down
20 changes: 10 additions & 10 deletions mlir/test/Dialect/MemRef/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -923,13 +923,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<[?, ?], 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<[?, ?], offset: 0>>
return %m2 : memref<?x2xf32, strided<[?, ?], offset: 0>>
}
// -----

Expand All @@ -954,10 +954,10 @@ func.func @reinterpret_of_extract_strided_metadata_same_type(%arg0 : memref<?x?x
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [%[[C0]]], sizes: [4, 2, 2], strides: [1, 1, %[[C1]]]
// CHECK: return %[[RES]]
func.func @reinterpret_of_extract_strided_metadata_w_different_stride(%arg0 : memref<8x2xf32>) -> memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> {
func.func @reinterpret_of_extract_strided_metadata_w_different_stride(%arg0 : memref<8x2xf32>) -> memref<4x2x2xf32, strided<[?, ?, ?], 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<[?, ?, ?], offset: ?>>
return %m2 : memref<4x2x2xf32, strided<[?, ?, ?], offset: ?>>
}
// -----

Expand All @@ -970,10 +970,10 @@ func.func @reinterpret_of_extract_strided_metadata_w_different_stride(%arg0 : me
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[RES:.*]] = memref.reinterpret_cast %[[ARG]] to offset: [1], sizes: [%[[C8]], %[[C2]]], strides: [%[[C2]], %[[C1]]]
// CHECK: return %[[RES]]
func.func @reinterpret_of_extract_strided_metadata_w_different_offset(%arg0 : memref<8x2xf32>) -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
func.func @reinterpret_of_extract_strided_metadata_w_different_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>>
}

// -----
Expand Down
47 changes: 37 additions & 10 deletions mlir/test/Dialect/MemRef/invalid.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,16 +214,6 @@ func.func @memref_reinterpret_cast_no_map_but_offset(%in: memref<?xf32>) {
: memref<?xf32> to memref<10xf32>
return
}

// -----

func.func @memref_reinterpret_cast_offset_mismatch_dynamic(%in: memref<?xf32>, %offset : index) {
// expected-error @+1 {{expected result type with offset = dynamic instead of 0}}
%out = memref.reinterpret_cast %in to offset: [%offset], sizes: [10], strides: [1]
: memref<?xf32> to memref<10xf32>
return
}

// -----

func.func @memref_reinterpret_cast_no_map_but_stride(%in: memref<?xf32>) {
Expand All @@ -245,6 +235,43 @@ func.func @memref_reinterpret_cast_no_map_but_strides(%in: memref<?x?xf32>) {

// -----

func.func @memref_reinterpret_cast_static_dynamic_size_mismatch(%in: memref<1x?x2x1xf32>) {
// expected-error@+1 {{expected size is static, but result type dimension is dynamic }}
%out = memref.reinterpret_cast %in to
offset: [0], sizes: [1, 4672, 1, 1], strides: [4672, 8, 8, 1]
: memref<1x?x2x1xf32> to memref<1x4672x?x1xf32>
}

// -----

func.func @memref_reinterpret_cast_dynamic_static_size_mismatch(%in: memref<1x?x2x1xf32>, %size: index) {
// expected-error@+1 {{expected size is dynamic, but result type dimension is static }}
%out = memref.reinterpret_cast %in to
offset: [0], sizes: [1, %size, 1, 1], strides: [4672, 8, 8, 1]
: memref<1x?x2x1xf32> to memref<1x4672x2x1xf32>
return
}

// -----

func.func @memref_reinterpret_cast_static_dynamic_offset_mismatch(%in: memref<?xf32>) {
// expected-error@+1 {{expected offset is static, but result type offset is dynamic}}
%out = memref.reinterpret_cast %in to
offset: [0], sizes: [10], strides: [1]
: memref<?xf32> to memref<10xf32, strided<[1], offset: ?>>
}

// -----

func.func @memref_reinterpret_cast_dynamic_static_offset_mismatch(%in: memref<?xf32>, %offset: index) {
// expected-error@+1 {{expected offset is dynamic, but result type offset is static}}
%out = memref.reinterpret_cast %in to
offset: [%offset], sizes: [10], strides: [1]
: memref<?xf32> to memref<10xf32, strided<[1], offset: 0>>
return
}

// -----
func.func @memref_reshape_element_type_mismatch(
%buf: memref<*xf32>, %shape: memref<1xi32>) {
// expected-error @+1 {{element types of source and destination memref types should be the same}}
Expand Down
12 changes: 6 additions & 6 deletions mlir/test/Dialect/MemRef/ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,20 +131,20 @@ 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: ?>> {
-> memref<10x10xf32, strided<[?, 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: ?>>
: memref<?xf32> to memref<10x10xf32, strided<[?, 1], offset: 1>>
return %out : memref<10x10xf32, strided<[?, 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], 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], offset: ?>>
return %out : memref<10x10xf32, strided<[?, 1], offset: ?>>
}

// CHECK-LABEL: func @memref_reshape(
Expand Down