[Mlir-commits] [mlir] [mlir][shape] Fix crash in FromExtentsOp::fold with poison operands (PR #178844)

[llvmlistbot at llvm.org]

https://github.com/mugiwaraluffy56 created https://github.com/llvm/llvm-project/pull/178844

## Summary
- Fix assertion failure in `shape.from_extents` fold when operands include `ub.poison`
- The fold assumed all non null attributes were `IntegerAttr`, but poison produces a different attribute type
- Use `dyn_cast_if_present` to safely handle non integer attributes

Fixes #178820

## Test plan
- Added regression test `@from_extents_poison` in canonicalize.mlir


>From 56d6a2422660a1c36de9a7bb677ce5e77f9e6dbe Mon Sep 17 00:00:00 2001
From: mugiwaraluffy56 <myakampuneeth at gmail.com>
Date: Fri, 30 Jan 2026 12:34:26 +0530
Subject: [PATCH] [mlir][shape] Fix crash in FromExtentsOp::fold with poison
 operands

The fold implementation assumed all non-null operand attributes were
IntegerAttr, but ub.poison produces a different attribute type. Use
dyn_cast_if_present to safely handle non-integer attributes.

Fixes #178820
---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 10 ++++++----
 mlir/test/Dialect/Shape/canonicalize.mlir | 12 ++++++++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 95546b4a4deb7..66cb360cecb63 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1203,11 +1203,13 @@ void IndexToSizeOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
 //===----------------------------------------------------------------------===//
 
 OpFoldResult FromExtentsOp::fold(FoldAdaptor adaptor) {
-  if (llvm::any_of(adaptor.getExtents(), [](Attribute a) { return !a; }))
-    return nullptr;
   SmallVector<int64_t, 6> extents;
-  for (auto attr : adaptor.getExtents())
-    extents.push_back(llvm::cast<IntegerAttr>(attr).getInt());
+  for (Attribute attr : adaptor.getExtents()) {
+    auto intAttr = llvm::dyn_cast_if_present<IntegerAttr>(attr);
+    if (!intAttr)
+      return nullptr;
+    extents.push_back(intAttr.getInt());
+  }
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 32a7c0186f282..f3c25b8c8100e 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -272,6 +272,18 @@ func.func @no_fold(%arg0: index) -> !shape.shape {
 
 // -----
 
+// GH#178820: Verify from_extents doesn't crash on poison operands.
+// CHECK-LABEL: func @from_extents_poison
+func.func @from_extents_poison() -> !shape.shape {
+  // CHECK: %[[POISON:.*]] = ub.poison : index
+  // CHECK: shape.from_extents %[[POISON]]
+  %0 = ub.poison : index
+  %ret = shape.from_extents %0 : index
+  return %ret : !shape.shape
+}
+
+// -----
+
 // Cast constant size to index and fold it away.
 // CHECK-LABEL: func @const_size_to_index
 func.func @const_size_to_index() -> index {