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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Rik Huijzer (rikhuijzer)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/74204.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+3-2) 
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+3-2) 
- (modified) mlir/test/Transforms/loop-invariant-code-motion.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index dce96cca016ff..e621ea9d5443e 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..54fecc0639028 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..d5a0c4e459e40 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,6 +699,23 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
 
 // -----
 
+// CHECK-LABEL: @no_speculate_tensor_dim_known_rank_known_dim_out_of_bounds
+func.func @no_speculate_tensor_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) {alignment = 64 : i64} : memref<?xi1>
+  %4 = tensor.empty(%c22, %c22, %c22) : tensor<?x?x?xi1>
+  scf.for %arg4 = %c0 to %c22 step %c1 {
+    %dim = memref.dim %alloc, %c22 : 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) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/74204