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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jueon Park (JueonPark)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+13-3) 
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+13) 


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

``````````

</details>


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