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

[Rik Huijzer]

https://github.com/rikhuijzer created https://github.com/llvm/llvm-project/pull/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.

>From c11845000eb41f12167f2a966e9bb7c1bb14c997 Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:29:19 +0100
Subject: [PATCH 1/3] [mlir][memref] Fix an invalid dim loop motion crash

---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp        |  2 --
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp        |  2 --
 .../Transforms/loop-invariant-code-motion.mlir  | 17 +++++++++++++++++
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index dce96cca016ff..b9c2e61180652 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,8 +893,6 @@ Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
-  // The verifier rejects operations that violate this assertion.
-  assert(constantIndex < rankedSourceType.getRank());
   return Speculation::Speculatable;
 }
 
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 8970ea1c73b40..c4bbf64381eab 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,8 +686,6 @@ Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
-  // The verifier rejects operations that violate this assertion.
-  assert(constantIndex < rankedSourceType.getRank());
   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..1b2983916dd2a 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: @speculate_tensor_dim_known_rank_known_dim_out_of_bounds
+func.func @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>
+  }
+  spirv.Return
+}
+// CHECK: memref.dim
+// CHECK-NEXT: scf.for
+
+// -----
+
 func.func @no_speculate_divui(
 // CHECK-LABEL: @no_speculate_divui(
     %num: i32, %denom: i32, %lb: index, %ub: index, %step: index) {

>From cd154ac875578a16c18e35ff1e4e840fdfe10544 Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:46:20 +0100
Subject: [PATCH 2/3] Set to `NotSpeculatable`

---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp             | 3 +++
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp             | 3 +++
 mlir/test/Transforms/loop-invariant-code-motion.mlir | 6 +++---
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index b9c2e61180652..e621ea9d5443e 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -893,6 +893,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
+  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 c4bbf64381eab..54fecc0639028 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -686,6 +686,9 @@ Speculation::Speculatability DimOp::getSpeculatability() {
   if (!rankedSourceType)
     return Speculation::NotSpeculatable;
 
+  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 1b2983916dd2a..8013384b31add 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -709,10 +709,10 @@ func.func @speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
   scf.for %arg4 = %c0 to %c22 step %c1 {
     %dim = memref.dim %alloc, %c22 : memref<?xi1>
   }
-  spirv.Return
+  return
 }
-// CHECK: memref.dim
-// CHECK-NEXT: scf.for
+// CHECK: scf.for
+// CHECK-NEXT: memref.dim
 
 // -----
 

>From 5962f83c6f0a10814aac385785fd09c0e0af34ef Mon Sep 17 00:00:00 2001
From: Rik Huijzer <github at huijzer.xyz>
Date: Sat, 2 Dec 2023 19:50:51 +0100
Subject: [PATCH 3/3] Change name

---
 mlir/test/Transforms/loop-invariant-code-motion.mlir | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 8013384b31add..d5a0c4e459e40 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -699,8 +699,8 @@ func.func @speculate_memref_dim_known_rank_known_dim_inbounds(
 
 // -----
 
-// CHECK-LABEL: @speculate_tensor_dim_known_rank_known_dim_out_of_bounds
-func.func @speculate_tensor_dim_known_rank_known_dim_out_of_bounds() {
+// 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