[mlir][memref] Fix an invalid dim loop motion crash (#74204)

Fixes https://github.com/llvm/llvm-project/issues/73382.

This PR suggests to replace two assertions that were introduced in
adabce4118
(https://reviews.llvm.org/D135748). According to the enum definition of
`NotSpeculatable`, an op that invokes undefined behavior is
`NotSpeculatable`.

0c06e8745f/mlir/include/mlir/Interfaces/SideEffectInterfaces.h (L248-L258)

and both `tensor.dim` and `memref.dim` state that "If the dimension
index is out of bounds, the behavior is undefined."

So therefore it seems to me that `DimOp::getSpeculatability()` should
return `NotSpeculatable` if the dimension index is out of bounds.

The added test is just a simplified version of
https://github.com/llvm/llvm-project/issues/73382.
This commit is contained in:
Rik Huijzer 2023-12-04 08:57:59 +01:00 committed by GitHub
parent d49e9d88a2
commit c9c1b3c37f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 4 deletions

View File

@ -893,8 +893,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
// The verifier rejects operations that violate this assertion.
assert(constantIndex < rankedSourceType.getRank());
if (rankedSourceType.getRank() <= constantIndex)
return Speculation::NotSpeculatable;
return Speculation::Speculatable;
}

View File

@ -686,8 +686,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
if (!rankedSourceType)
return Speculation::NotSpeculatable;
// The verifier rejects operations that violate this assertion.
assert(constantIndex < rankedSourceType.getRank());
if (rankedSourceType.getRank() <= constantIndex)
return Speculation::NotSpeculatable;
return Speculation::Speculatable;
}

View File

@ -699,6 +699,54 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
// -----
// CHECK-LABEL: @speculate_memref_dim_known_rank_known_dim_inbounds
func.func @speculate_memref_dim_known_rank_known_dim_inbounds() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c22 = arith.constant 22 : index
%alloc = memref.alloc(%c22) : memref<?xi1>
scf.for %arg4 = %c0 to %c22 step %c1 {
%dim = memref.dim %alloc, %c0 : memref<?xi1>
}
return
}
// CHECK: memref.dim
// CHECK-NEXT: scf.for
// -----
// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_inbounds
func.func @speculate_tensor_dim_known_rank_known_dim_inbounds() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c22 = arith.constant 22 : index
%t = tensor.empty(%c22, %c22) : tensor<?x?xi1>
scf.for %arg4 = %c0 to %c22 step %c1 {
%dim = tensor.dim %t, %c1 : tensor<?x?xi1>
}
return
}
// CHECK: tensor.dim
// CHECK-NEXT: scf.for
// -----
// CHECK-LABEL: @no_speculate_memref_dim_known_rank_known_dim_out_of_bounds
func.func @no_speculate_memref_dim_known_rank_known_dim_out_of_bounds() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c22 = arith.constant 22 : index
%alloc = memref.alloc(%c22) : memref<?xi1>
scf.for %arg4 = %c0 to %c22 step %c1 {
%dim = memref.dim %alloc, %c1 : memref<?xi1>
}
return
}
// CHECK: scf.for
// CHECK-NEXT: memref.dim
// -----
func.func @no_speculate_divui(
// CHECK-LABEL: @no_speculate_divui(
%num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {