[Mlir-commits] [mlir] c9c1b3c - [mlir][memref] Fix an invalid dim loop motion crash (#74204)

[llvmlistbot at llvm.org]

Author: Rik Huijzer
Date: 2023-12-04T08:57:59+01:00
New Revision: c9c1b3c37fa5d5c617068afe67afcafacd58a241

URL: https://github.com/llvm/llvm-project/commit/c9c1b3c37fa5d5c617068afe67afcafacd58a241
DIFF: https://github.com/llvm/llvm-project/commit/c9c1b3c37fa5d5c617068afe67afcafacd58a241.diff

LOG: [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
https://github.com/llvm/llvm-project/commit/adabce41185910227ca276a1cfd22e76443dd238
(https://reviews.llvm.org/D135748). According to the enum definition of
`NotSpeculatable`, an op that invokes undefined behavior is
`NotSpeculatable`.

https://github.com/llvm/llvm-project/blob/0c06e8745f131d867c566f4d35a7a04e24b4a075/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.

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    mlir/test/Transforms/loop-invariant-code-motion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index dce96cca016ff..a397506629cfe 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -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;
 }
 

diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8970ea1c73b40..f15695383d34a 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -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;
 }
 

diff  --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index e9b1e22359001..1415583dde9da 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -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) {