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

[Mehdi Amini]

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

>From d05e7ee209055549bcd61a1c5215b92c50719459 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 05:43:37 -0800
Subject: [PATCH] [mlir][Shape] Fix crash in BroadcastOp::fold when operand is
 ub.poison

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
---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 14 +++++++-------
 mlir/test/Dialect/Shape/canonicalize.mlir | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 4d03b7b2b2064..053173bdf8510 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
+}