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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Jueon Park (JueonPark)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/IR/BuiltinAttributes.cpp (+4-2) 
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+13) 


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

``````````

</details>


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