[mlir][vector] Fix crash in vector.from_elements folding with poison values

Author: jhalakpatel

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);
 }
 

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -3375,6 +3375,17 @@ func.func @negative_from_elements_to_constant() -> 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>
+}
+
 // +---------------------------------------------------------------------------
 // End of  Tests for foldFromElementsToConstant
 // +---------------------------------------------------------------------------