[Mlir-commits] [mlir] [mlir][vector] Fix crash in vector.from_elements folding with poison values (PR #158528)

[Jhalak Patel]

https://github.com/jhalakpatel created https://github.com/llvm/llvm-project/pull/158528

The vector.from_elements constant folding was crashing when poison values
were present in the element list. The convertIntegerAttr function was not
properly detecting poison attributes, leading to assertion failures in
dyn_cast operations.

This patch adds explicit poison detection using isa<ub::PoisonAttrInterface>
before attempting any casting operations. The folding now gracefully aborts
when poison attributes are encountered, preventing the crash while preserving
correct folding for legitimate mixed-type constants (int/float).

Fixes assertion: "dyn_cast on a non-existent value" when processing
ub.poison values in vector.from_elements operations.

>From 69730a6e92adbdd8d88a526adb5bbf0e7e26a82f Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Sun, 14 Sep 2025 20:00:11 -0700
Subject: [PATCH] [mlir][vector] Fix crash in vector.from_elements folding with
 poison values

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   | 27 ++++++++++++++++------
 mlir/test/Dialect/Vector/canonicalize.mlir | 11 +++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 85e485c28c74e..8c38305e313bb 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -398,14 +398,21 @@ std::optional<int64_t> vector::getConstantVscaleMultiplier(Value value) {
 
 /// Converts an IntegerAttr to have the specified type if needed.
 /// This handles cases where constant attributes have a different type than the
-/// target element type. If the input attribute is not an IntegerAttr or already
-/// has the correct type, returns it unchanged.
+/// target element type. Returns null if the attribute is poison/invalid or
+/// conversion fails.
 static Attribute convertIntegerAttr(Attribute attr, Type expectedType) {
-  if (auto intAttr = mlir::dyn_cast<IntegerAttr>(attr)) {
-    if (intAttr.getType() != expectedType)
-      return IntegerAttr::get(expectedType, intAttr.getInt());
-  }
-  return attr;
+  // Check for poison attributes before any casting operations
+  if (!attr || isa<ub::PoisonAttrInterface>(attr))
+    return {}; // Poison or invalid attribute
+
+  auto intAttr = mlir::dyn_cast<IntegerAttr>(attr);
+  if (!intAttr)
+    return attr; // Not an IntegerAttr, return unchanged (e.g., FloatAttr)
+
+  if (intAttr.getType() == expectedType)
+    return attr; // Already correct type
+
+  return IntegerAttr::get(expectedType, intAttr.getInt());
 }
 
 //===----------------------------------------------------------------------===//
@@ -2478,6 +2485,12 @@ static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
     return convertIntegerAttr(attr, destEltType);
   });
 
+  // Check if any attributes are poison/invalid (indicated by null attributes).
+  // Note: convertIntegerAttr returns valid non-integer attributes unchanged,
+  // only returns null for poison/invalid attributes.
+  if (llvm::any_of(convertedElements, [](Attribute attr) { return !attr; }))
+    return {};
+
   return DenseElementsAttr::get(destVecType, convertedElements);
 }
 
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index e7381e0c8997e..e300b76b86515 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3740,3 +3740,14 @@ func.func @from_address_of_regression() -> vector<1x!llvm.ptr> {
   %b = vector.from_elements %a : vector<1x!llvm.ptr>
   return %b : vector<1x!llvm.ptr>
 }
+
+// -----
+
+// CHECK-LABEL: @from_elements_poison
+//       CHECK:   %[[VAL:.*]] = ub.poison : vector<2xf32>
+//       CHECK:   return %[[VAL]] : vector<2xf32>
+func.func @from_elements_poison() -> vector<2xf32> {
+  %0 = ub.poison : f32
+  %1 = vector.from_elements %0, %0 : vector<2xf32>
+  return %1 : vector<2xf32>
+}