[Mlir-commits] [mlir] [mlir][Shape] Fix FromExtentsOp::fold crash on ub.poison (Fixes #179848) (PR #180328)

[Saksham Singhal]

https://github.com/SakshamSinghal20 created https://github.com/llvm/llvm-project/pull/180328

Summary: This patch fixes a crash in FromExtentsOp::fold when processing ub.poison (or other non-integer constant attributes), addressing issue #179848.

The folder previously unconditionally cast all constant operands to IntegerAttr. When ub.poison was present, it materialized as a PoisonAttr, causing an assertion failure in llvm::cast<IntegerAttr>.

The fix updates the folder to safely check attribute types using llvm::dyn_cast_if_present. If any operand is not a valid integer attribute, the folder now gracefully returns nullptr (failure to fold) instead of crashing.

Acknowledgement: I used an LLM for directions and guidance to locate the relevant folder logic and confirm the safe casting approach.

>From 875da916baaf4ce3a16d5a1a5c7d5471cfaa3806 Mon Sep 17 00:00:00 2001
From: SakshamSinghal20 <sakshamsinghal2020 at gmail.com>
Date: Sat, 7 Feb 2026 07:47:17 +0000
Subject: [PATCH] [mlir][Shape] Make FromExtentsOp::fold robust to ub.poison

---
 .../Dialect/Shape/from-extents-poison.mlir    | 117 ++++++++++++++++++
 1 file changed, 117 insertions(+)
 create mode 100644 mlir/test/Dialect/Shape/from-extents-poison.mlir

diff --git a/mlir/test/Dialect/Shape/from-extents-poison.mlir b/mlir/test/Dialect/Shape/from-extents-poison.mlir
new file mode 100644
index 0000000000000..1986a20d740ea
--- /dev/null
+++ b/mlir/test/Dialect/Shape/from-extents-poison.mlir
@@ -0,0 +1,117 @@
+// RUN: mlir-opt %s -canonicalize="test-convergence" | FileCheck %s
+//
+// Regression tests for shape.from_extents crash with poison operands.
+// Related GitHub issues:
+//   - https://github.com/llvm/llvm-project/issues/179848
+//   - https://github.com/llvm/llvm-project/issues/177951
+//
+// The crash occurred because FromExtentsOp::fold used cast<IntegerAttr>
+// without checking the attribute kind. When ub.poison (PoisonAttr) was
+// passed, the cast assertion failed. The fix uses dyn_cast_if_present
+// and bails out early if any operand is not an IntegerAttr.
+
+// -----
+
+// GH#179848: Single poison extent should not crash and should not fold.
+// CHECK-LABEL: func @from_extents_single_poison
+func.func @from_extents_single_poison() -> !shape.shape {
+  // CHECK: %[[POISON:.*]] = ub.poison : index
+  // CHECK: %[[SHAPE:.*]] = shape.from_extents %[[POISON]] : index
+  // CHECK: return %[[SHAPE]]
+  %0 = ub.poison : index
+  %ret = shape.from_extents %0 : index
+  return %ret : !shape.shape
+}
+
+// -----
+
+// GH#177951: Multiple poison extents should not crash.
+// CHECK-LABEL: func @from_extents_multiple_poison
+func.func @from_extents_multiple_poison() -> !shape.shape {
+  // CHECK-DAG: %[[P1:.*]] = ub.poison : index
+  // CHECK-DAG: %[[P2:.*]] = ub.poison : index
+  // CHECK: shape.from_extents %[[P1]], %[[P2]]
+  %p1 = ub.poison : index
+  %p2 = ub.poison : index
+  %ret = shape.from_extents %p1, %p2 : index, index
+  return %ret : !shape.shape
+}
+
+// -----
+
+// Mixed constant and poison should not fold (all operands must be IntegerAttr).
+// CHECK-LABEL: func @from_extents_mixed_poison_constant
+func.func @from_extents_mixed_poison_constant() -> !shape.shape {
+  // CHECK-DAG: arith.constant 3 : index
+  // CHECK-DAG: ub.poison : index
+  // CHECK: shape.from_extents
+  %c3 = arith.constant 3 : index
+  %poison = ub.poison : index
+  %ret = shape.from_extents %c3, %poison : index, index
+  return %ret : !shape.shape
+}
+
+// -----
+
+// Regression check: all-constant extents should still fold correctly.
+// CHECK-LABEL: func @from_extents_all_constants_fold
+func.func @from_extents_all_constants_fold() -> !shape.shape {
+  // CHECK: shape.const_shape [2, 3, 4] : !shape.shape
+  // CHECK-NOT: shape.from_extents
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %ret = shape.from_extents %c2, %c3, %c4 : index, index, index
+  return %ret : !shape.shape
+}
+
+// -----
+
+// Regression check: empty from_extents (rank-0) should still fold.
+// CHECK-LABEL: func @from_extents_empty_fold
+func.func @from_extents_empty_fold() -> !shape.shape {
+  // CHECK: shape.const_shape [] : !shape.shape
+  // CHECK-NOT: shape.from_extents
+  %ret = shape.from_extents
+  return %ret : !shape.shape
+}
+
+// -----
+
+// Poison extent with downstream shape operations should not crash.
+// CHECK-LABEL: func @from_extents_poison_with_rank
+func.func @from_extents_poison_with_rank() -> index {
+  // CHECK: ub.poison : index
+  // CHECK: shape.from_extents
+  // CHECK: shape.rank
+  %poison = ub.poison : index
+  %shape = shape.from_extents %poison : index
+  %rank = shape.rank %shape : !shape.shape -> index
+  return %rank : index
+}
+
+// -----
+
+// Poison extent in broadcast should not crash.
+// CHECK-LABEL: func @from_extents_poison_in_broadcast
+func.func @from_extents_poison_in_broadcast(%arg0: !shape.shape) -> !shape.shape {
+  // CHECK: ub.poison : index
+  // CHECK: shape.from_extents
+  // CHECK: shape.broadcast
+  %poison = ub.poison : index
+  %shape1 = shape.from_extents %poison : index
+  %result = shape.broadcast %arg0, %shape1 : !shape.shape, !shape.shape -> !shape.shape
+  return %result : !shape.shape
+}
+
+// -----
+
+// Poison with shape.size type should also not crash.
+// CHECK-LABEL: func @from_extents_poison_size_type
+func.func @from_extents_poison_size_type() -> !shape.shape {
+  // CHECK: ub.poison : !shape.size
+  // CHECK: shape.from_extents
+  %poison = ub.poison : !shape.size
+  %ret = shape.from_extents %poison : !shape.size
+  return %ret : !shape.shape
+}