[Mlir-commits] [mlir] a115e6b - [mlir][Shape] Fix crash in BroadcastOp::fold when operand is ub.poison (#183931)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-05T13:23:53+01:00
New Revision: a115e6bd096da9655df3541fb8161949d69cf562

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

LOG: [mlir][Shape] Fix crash in BroadcastOp::fold when operand is ub.poison (#183931)

BroadcastOp::fold used an unchecked llvm::cast<DenseIntElementsAttr> on
each operand's folded attribute. The existing null-check only guarded
against a missing (unset) attribute, not against a non-null attribute of
a different type such as PoisonAttr (produced when an operand is
ub.poison).

Replace the unchecked casts with dyn_cast_or_null, bailing out with
nullptr (i.e. no fold) when any operand does not provide a
DenseIntElementsAttr.

Add a regression test with a ub.poison operand.

Fixes #179679

Added: 
    

Modified: 
    mlir/lib/Dialect/Shape/IR/Shape.cpp
    mlir/test/Dialect/Shape/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 38cac53be3022..43a6816c6d863 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -652,18 +652,18 @@ OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
     return getShapes().front();
   }
 
-  if (!adaptor.getShapes().front())
+  auto firstAttr =
+      dyn_cast_or_null<DenseIntElementsAttr>(adaptor.getShapes().front());
+  if (!firstAttr)
     return nullptr;
 
-  SmallVector<int64_t, 6> resultShape(
-      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes().front())
-          .getValues<int64_t>());
+  SmallVector<int64_t, 6> resultShape(firstAttr.getValues<int64_t>());
 
   for (auto next : adaptor.getShapes().drop_front()) {
-    if (!next)
+    auto nextAttr = dyn_cast_or_null<DenseIntElementsAttr>(next);
+    if (!nextAttr)
       return nullptr;
-    auto nextShape = llvm::to_vector<6>(
-        llvm::cast<DenseIntElementsAttr>(next).getValues<int64_t>());
+    auto nextShape = llvm::to_vector<6>(nextAttr.getValues<int64_t>());
 
     SmallVector<int64_t, 6> tmpShape;
     // If the shapes are not compatible, we can't fold it.

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 22add87ff3ed4..d1b5e7bb035bf 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1636,3 +1636,21 @@ func.func @shape_of_static_with_shape_result(%arg0: tensor<3xf32>) -> !shape.sha
   %0 = shape.shape_of %arg0 : tensor<3xf32> -> !shape.shape
   return %0 : !shape.shape
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/179679:
+// shape.broadcast fold used to crash with an unchecked cast when one of the
+// operands was ub.poison (a non-DenseIntElementsAttr attribute). The fold
+// must bail out gracefully instead.
+
+// CHECK-LABEL: @broadcast_no_crash_on_poison
+// CHECK-NOT: shape.broadcast
+// CHECK: return
+func.func @broadcast_no_crash_on_poison() {
+  %0 = shape.const_shape [1, 2, 3] : tensor<3xindex>
+  %1 = ub.poison : tensor<3xindex>
+  %2 = shape.broadcast %0, %1 : tensor<3xindex>, tensor<3xindex> -> tensor<3xindex>
+  %3 = tensor.rank %2 : tensor<3xindex>
+  return
+}