[Mlir-commits] [mlir] [mlir] Return null from DenseElementsAttr::get on unsupported attribute types (PR #181159)

[Jueon Park]

https://github.com/JueonPark created https://github.com/llvm/llvm-project/pull/181159

DenseElementsAttr::get(ShapedType, ArrayRef<Attribute>) crashed with an unconditional cast<IntegerAttr> when encountering attribute types that are neither FloatAttr nor IntegerAttr (e.g. ub.poison). This can happen when folding ops like tensor.from_elements whose operands include poison values.

As suggested in #179113 by @matthias-springer, this patch fixes the issue at the DenseElementsAttr::get level rather than in individual op folders. The cast<IntegerAttr> is replaced with dyn_cast<IntegerAttr>, and when the attribute is neither FloatAttr nor IntegerAttr, a null DenseElementsAttr is returned. This is a more robust fix because it prevents the same class of crashes in any caller that passes unsupported attributes to DenseElementsAttr::get.

Fixes #178209.

>From beaf93f7682605c0f6aafcccaa0b8ff5fcca68c0 Mon Sep 17 00:00:00 2001
From: rebel-jueonpark <jueonpark at rebellions.ai>
Date: Thu, 12 Feb 2026 23:57:51 +0900
Subject: [PATCH] [mlir] Return null from DenseElementsAttr::get on unsupported
 attribute types

DenseElementsAttr::get(ShapedType, ArrayRef<Attribute>) crashed with an
unconditional cast<IntegerAttr> when encountering attribute types that
are neither FloatAttr nor IntegerAttr (e.g. ub.poison). This can
happen when folding ops like tensor.from_elements whose operands
include poison values.

As suggested in #179113 by @matthias-springer, this patch fixes the issue
at the DenseElementsAttr::get level rather than in individual op
folders. The cast<IntegerAttr> is replaced with dyn_cast<IntegerAttr>,
and when the attribute is neither FloatAttr nor IntegerAttr, a null
DenseElementsAttr is returned. This is a more robust fix because it
prevents the same class of crashes in any caller that passes unsupported
attributes to DenseElementsAttr::get.

Fixes #178209.
---
 mlir/lib/IR/BuiltinAttributes.cpp          |  6 ++++--
 mlir/test/Dialect/Tensor/canonicalize.mlir | 13 +++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index fbbd9d29abe85..1a29fc534b40f 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -947,11 +947,13 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
       assert(floatAttr.getType() == eltType &&
              "expected float attribute type to equal element type");
       intVal = floatAttr.getValue().bitcastToAPInt();
-    } else {
-      auto intAttr = llvm::cast<IntegerAttr>(values[i]);
+    } else if (auto intAttr = llvm::dyn_cast<IntegerAttr>(values[i])) {
       assert(intAttr.getType() == eltType &&
              "expected integer attribute type to equal element type");
       intVal = intAttr.getValue();
+    } else {
+      // Unsupported attribute type.
+      return {};
     }
 
     assert(intVal.getBitWidth() == bitWidth &&
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 7a2d53c0c5850..f30941aa85bd0 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -398,6 +398,19 @@ func.func @extract_from_elements_complex_f() -> tensor<3xcomplex<f32>> {
 
 // -----
 
+// Ensure tensor.from_elements with poison values doesn't crash (issue #178209).
+// CHECK-LABEL: func @from_elements_with_poison
+func.func @from_elements_with_poison() -> tensor<1xindex> {
+  // CHECK: %[[POISON:.*]] = ub.poison : index
+  // CHECK: %[[TENSOR:.*]] = tensor.from_elements %[[POISON]] : tensor<1xindex>
+  // CHECK: return %[[TENSOR]]
+  %0 = ub.poison : index
+  %1 = tensor.from_elements %0 : tensor<1xindex>
+  return %1 : tensor<1xindex>
+}
+
+// -----
+
 // Ensure the optimization doesn't segfault from bad constants
 // CHECK-LABEL: func @extract_negative_from_tensor.from_elements
 func.func @extract_negative_from_tensor.from_elements(%element : index) -> index {