Fix: bail out when divisor is zero

[hiraditya]

Fixes: #131279

[llvmbot]

@llvm/pr-subscribers-mlir-shape

Author: AdityaK (hiraditya)

Changes

Fixes: #131279

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+1-1)
• (added) mlir/test/Dialect/Shape/fold-div.mlir (+13)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 2200af0f67a86..10ba808cd26c2 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1124,7 +1124,7 @@ OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
   if (!lhs)
     return nullptr;
   auto rhs = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getRhs());
-  if (!rhs)
+  if (!rhs || rhs.getValue().isZero())
     return nullptr;
 
   // Division in APInt does not follow floor(lhs, rhs) when the result is
diff --git a/mlir/test/Dialect/Shape/fold-div.mlir b/mlir/test/Dialect/Shape/fold-div.mlir
new file mode 100644
index 0000000000000..9deac08bd85cc
--- /dev/null
+++ b/mlir/test/Dialect/Shape/fold-div.mlir
@@ -0,0 +1,13 @@
+// Bug: #131279
+// RUN: mlir-opt --test-scf-pipelining %s | FileCheck %s
+// CHECK: fold_div_index_neg_rhs
+// CHECK-NEXT: %c0 = arith.constant 0 : index
+// CHECK-NEXT: %0 = shape.div %c0, %c0 : index, index -> index
+// CHECK-NEXT: return %0 : index
+module {
+  func.func @fold_div_index_neg_rhs() -> index {
+    %c0 = arith.constant 0 : index
+    %0 = shape.div %c0, %c0 : index, index -> index
+    return %0 : index
+  }
+}

[llvmbot]

@llvm/pr-subscribers-mlir

Author: AdityaK (hiraditya)

Changes

Fixes: #131279

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+1-1)
• (added) mlir/test/Dialect/Shape/fold-div.mlir (+13)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 2200af0f67a86..10ba808cd26c2 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1124,7 +1124,7 @@ OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
   if (!lhs)
     return nullptr;
   auto rhs = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getRhs());
-  if (!rhs)
+  if (!rhs || rhs.getValue().isZero())
     return nullptr;
 
   // Division in APInt does not follow floor(lhs, rhs) when the result is
diff --git a/mlir/test/Dialect/Shape/fold-div.mlir b/mlir/test/Dialect/Shape/fold-div.mlir
new file mode 100644
index 0000000000000..9deac08bd85cc
--- /dev/null
+++ b/mlir/test/Dialect/Shape/fold-div.mlir
@@ -0,0 +1,13 @@
+// Bug: #131279
+// RUN: mlir-opt --test-scf-pipelining %s | FileCheck %s
+// CHECK: fold_div_index_neg_rhs
+// CHECK-NEXT: %c0 = arith.constant 0 : index
+// CHECK-NEXT: %0 = shape.div %c0, %c0 : index, index -> index
+// CHECK-NEXT: return %0 : index
+module {
+  func.func @fold_div_index_neg_rhs() -> index {
+    %c0 = arith.constant 0 : index
+    %0 = shape.div %c0, %c0 : index, index -> index
+    return %0 : index
+  }
+}

[jpienaar]

Thanks!