[Mlir-commits] [mlir] [mlir][vector] Fix crash in vector.from_elements folding with poison values (PR #158528)
Jhalak Patel
llvmlistbot at llvm.org
Mon Sep 15 10:17:44 PDT 2025
https://github.com/jhalakpatel updated https://github.com/llvm/llvm-project/pull/158528
>From 0bcc205ff503f98ba321b9586dbd8e2031084822 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 1/4] [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 05c88b8abfbb0..6e69bb1f761d3 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/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
// +---------------------------------------------------------------------------
>From ab75f64091e8924949246698bfda170baccc5ea4 Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Mon, 15 Sep 2025 09:56:35 -0700
Subject: [PATCH 2/4] Add negative prefix to test
---
mlir/test/Dialect/Vector/canonicalize.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6e69bb1f761d3..8dedde2ade1ce 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3377,10 +3377,10 @@ func.func @negative_from_elements_to_constant() -> vector<1x!llvm.ptr> {
// -----
-// CHECK-LABEL: @from_elements_poison
+// CHECK-LABEL: @negative_from_elements_poison
// CHECK: %[[VAL:.*]] = ub.poison : vector<2xf32>
// CHECK: return %[[VAL]] : vector<2xf32>
-func.func @from_elements_poison() -> vector<2xf32> {
+func.func @negative_from_elements_poison() -> vector<2xf32> {
%0 = ub.poison : f32
%1 = vector.from_elements %0, %0 : vector<2xf32>
return %1 : vector<2xf32>
>From 13af161a1c999c1f7d34f8b005f768a15cb573ed Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Mon, 15 Sep 2025 10:03:46 -0700
Subject: [PATCH 3/4] Address review comments in
https://github.com/llvm/llvm-project/pull/158528\#issuecomment-3290491095
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 57 ++++++++++++------------
1 file changed, 28 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 8c38305e313bb..88925274cc0c0 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -397,20 +397,11 @@ 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. Returns null if the attribute is poison/invalid or
-/// conversion fails.
-static Attribute convertIntegerAttr(Attribute attr, Type expectedType) {
- // 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)
-
+/// This handles cases where integer constant attributes have a different type
+/// than the target element type.
+static IntegerAttr convertIntegerAttr(IntegerAttr intAttr, Type expectedType) {
if (intAttr.getType() == expectedType)
- return attr; // Already correct type
+ return intAttr; // Already correct type
return IntegerAttr::get(expectedType, intAttr.getInt());
}
@@ -2470,7 +2461,10 @@ static OpFoldResult foldFromElementsToElements(FromElementsOp fromElementsOp) {
///
static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
ArrayRef<Attribute> elements) {
- if (llvm::any_of(elements, [](Attribute attr) { return !attr; }))
+ // Check for null or poison attributes before any processing.
+ if (llvm::any_of(elements, [](Attribute attr) {
+ return !attr || isa<ub::PoisonAttrInterface>(attr);
+ }))
return {};
// DenseElementsAttr only supports int/index/float/complex types.
@@ -2479,18 +2473,14 @@ static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
if (!destEltType.isIntOrIndexOrFloat() && !isa<ComplexType>(destEltType))
return {};
- // Constant attributes might have a different type than the return type.
- // Convert them before creating the dense elements attribute.
- auto convertedElements = llvm::map_to_vector(elements, [&](Attribute attr) {
- return convertIntegerAttr(attr, destEltType);
+ // Convert integer attributes to the target type if needed, leave others unchanged.
+ auto convertedElements = llvm::map_to_vector(elements, [&](Attribute attr) -> Attribute {
+ if (auto intAttr = dyn_cast<IntegerAttr>(attr)) {
+ return convertIntegerAttr(intAttr, destEltType);
+ }
+ return attr; // Non-integer attributes (FloatAttr, etc.) returned unchanged
});
- // 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);
}
@@ -3510,13 +3500,22 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
SmallVector<Attribute> insertedValues;
Type destEltType = destTy.getElementType();
- /// Converts the expected type to an IntegerAttr if there's
- /// a mismatch.
+ /// Converts integer attributes to the expected type if there's a mismatch.
+ /// Non-integer attributes are left unchanged.
if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(srcAttr)) {
- for (auto value : denseSource.getValues<Attribute>())
- insertedValues.push_back(convertIntegerAttr(value, destEltType));
+ for (auto value : denseSource.getValues<Attribute>()) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(value)) {
+ insertedValues.push_back(convertIntegerAttr(intAttr, destEltType));
+ } else {
+ insertedValues.push_back(value); // Non-integer attributes unchanged
+ }
+ }
} else {
- insertedValues.push_back(convertIntegerAttr(srcAttr, destEltType));
+ if (auto intAttr = dyn_cast<IntegerAttr>(srcAttr)) {
+ insertedValues.push_back(convertIntegerAttr(intAttr, destEltType));
+ } else {
+ insertedValues.push_back(srcAttr); // Non-integer attributes unchanged
+ }
}
auto allValues = llvm::to_vector(denseDst.getValues<Attribute>());
>From 252112709d243c0d0d1bf350d92fdf53ea25bdef Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Mon, 15 Sep 2025 10:17:12 -0700
Subject: [PATCH 4/4] Remove extra braces
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 13 +++++--------
1 file changed, 5 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 88925274cc0c0..5afeaae2a0d8e 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3503,19 +3503,16 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
/// Converts integer attributes to the expected type if there's a mismatch.
/// Non-integer attributes are left unchanged.
if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(srcAttr)) {
- for (auto value : denseSource.getValues<Attribute>()) {
- if (auto intAttr = dyn_cast<IntegerAttr>(value)) {
+ for (auto value : denseSource.getValues<Attribute>())
+ if (auto intAttr = dyn_cast<IntegerAttr>(value))
insertedValues.push_back(convertIntegerAttr(intAttr, destEltType));
- } else {
+ else
insertedValues.push_back(value); // Non-integer attributes unchanged
- }
- }
} else {
- if (auto intAttr = dyn_cast<IntegerAttr>(srcAttr)) {
+ if (auto intAttr = dyn_cast<IntegerAttr>(srcAttr))
insertedValues.push_back(convertIntegerAttr(intAttr, destEltType));
- } else {
+ else
insertedValues.push_back(srcAttr); // Non-integer attributes unchanged
- }
}
auto allValues = llvm::to_vector(denseDst.getValues<Attribute>());
More information about the Mlir-commits
mailing list