[Mlir-commits] [mlir] [mlir][tensor] Fix crash in FromElementsOp::fold with poison values (PR #179113)

[Jueon Park]

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

The FromElementsOp::fold method crashed when processing elements that are PoisonAttr because DenseElementsAttr::get only accepts IntegerAttr, FloatAttr, StringAttr, or ArrayAttr (for complex types).

Add validation to check that all elements are supported attribute types before calling DenseElementsAttr::get.

Fixes #178209

>From d4136534f929c6fce16e56adfe24a5b4e387f10a Mon Sep 17 00:00:00 2001
From: rebel-jueonpark <jueonpark at rebellions.ai>
Date: Mon, 2 Feb 2026 00:12:04 +0900
Subject: [PATCH] [mlir][tensor] Fix crash in FromElementsOp::fold with poison
 values

The FromElementsOp::fold method crashed when processing elements that
are PoisonAttr because DenseElementsAttr::get only accepts IntegerAttr,
FloatAttr, StringAttr, or ArrayAttr (for complex types).

Add validation to check that all elements are supported attribute types
before calling DenseElementsAttr::get.

Fixes #178209
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp   | 16 +++++++++++++---
 mlir/test/Dialect/Tensor/canonicalize.mlir | 13 +++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index d885d2c871e3f..8616e6312a366 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1456,9 +1456,19 @@ void FromElementsOp::build(OpBuilder &builder, OperationState &result,
 }
 
 OpFoldResult FromElementsOp::fold(FoldAdaptor adaptor) {
-  if (!llvm::is_contained(adaptor.getElements(), nullptr))
-    return DenseElementsAttr::get(getType(), adaptor.getElements());
-  return {};
+  if (llvm::is_contained(adaptor.getElements(), nullptr))
+    return {};
+
+  // DenseElementsAttr::get only accepts IntegerAttr, FloatAttr, StringAttr,
+  // or ArrayAttr (for complex types). Check that all elements are valid
+  // attribute types before constructing the DenseElementsAttr.
+  // This avoids crashes when e.g. a PoisonAttr is folded in.
+  if (!llvm::all_of(adaptor.getElements(), [](Attribute attr) {
+        return isa<IntegerAttr, FloatAttr, StringAttr, ArrayAttr>(attr);
+      }))
+    return {};
+
+  return DenseElementsAttr::get(getType(), adaptor.getElements());
 }
 
 namespace {
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index f85831b6f4cab..1582fd14b6e45 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 {