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

[jhalakpatel]

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

This patch refactors convertIntegerAttr to take IntegerAttr directly, moving poison detection to the caller using explicit isaub::PoisonAttrInterface checks. The function signature change provides compile-time type safety while the early poison validation in foldFromElementsToConstant prevents unsafe 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.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Jhalak Patel (jhalakpatel)

Changes

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.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+20-7)
• (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+11)

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>
+}

[yangtetris]

Thanks for fixing the issue. I didn't consider this comprehensively when moving the conversion logic to a separate function.

I'm wondering if we could fix this in another way:

1. Change the parameter type of convertIntegerAttr to IntegerAttr.
2. Expand the !attr check at the beginning of foldFromElementsToConstant to handle poison attrs.

Given that the function name is already convertIntegerAttr, it seems like ensuring the type should be the caller's responsibility.

[banach-space]

Thanks for the fix!

I'm wondering if we could fix this in another way:

1. Change the parameter type of convertIntegerAttr to IntegerAttr.
2. Expand the !attr check at the beginning of foldFromElementsToConstant to handle poison attrs.

Given that the function name is already convertIntegerAttr, it seems like ensuring the type should be the caller's responsibility.

+1

On a related note:

• lets try to keep all tests for foldFromElementsToConstant grouped together,
• lets use the prefix @negative_ to label negative tests.

For more details: https://mlir.llvm.org/getting_started/TestingGuide/

We seem to have already diverged a bit with tests for foldFromElementsToConstant, here's a fix:

• [mlir][vector][nfc] Group all tests for foldFromElementsToConstant #158578

[dcaballe]

Thanks for the fix! I provided some feedback. I think we should handle float attributes in the same way we handle integer attributes. Please, let me know if you have any questions!

[dcaballe]

LGTM % minor comments. I think we can handle the fp case separately. Thanks!

[yangtetris]

LGTM. We can leave the float attribute conversion for another MR.

[jhalakpatel]

LGTM % minor comments. I think we can handle the fp case separately. Thanks!

Filed a follow up issue: #159613

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[github-actions]

@jhalakpatel Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!