From c349061dc1a8d1278fb91190e760ac594a3a1dd8 Mon Sep 17 00:00:00 2001
From: AdityaK <hiraditya@msn.com>
Date: Fri, 28 Mar 2025 20:53:11 +0000
Subject: [PATCH] Fix: bail out when divisor is zero

---
 mlir/lib/Dialect/Shape/IR/Shape.cpp   |  2 +-
 mlir/test/Dialect/Shape/fold-div.mlir | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Shape/fold-div.mlir

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
+  }
+}