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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: puneeth_aditya_5656 (mugiwaraluffy56)

<details>
<summary>Changes</summary>

## 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


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+6-4) 
- (modified) mlir/test/Dialect/Shape/canonicalize.mlir (+12) 


``````````diff
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 {

``````````

</details>


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