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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+7-7) 
- (modified) mlir/test/Dialect/Shape/canonicalize.mlir (+18) 


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

``````````

</details>


https://github.com/llvm/llvm-project/pull/183931