[Mlir-commits] [mlir] [mlir][vector] Fix crash in vector.from_elements folding with poison values (PR #158528)
Jhalak Patel
llvmlistbot at llvm.org
Thu Sep 18 13:01:12 PDT 2025
https://github.com/jhalakpatel updated https://github.com/llvm/llvm-project/pull/158528
>From c5b8798c7b8d6e97b10eea97bcc1a4ddaa09aa7d 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/6] [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 8d6e263934fb4..6dbd2897f546f 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 a368667366bb402bf95aa7ef13d5ae8bbec70e93 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/6] 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 ddd6040099f800b5f292f7e20451c8300f1d4474 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/6] 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 6dbd2897f546f..c302d2b6d956c 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 c5e38a7c26239e2c806cc015e72eb28ff5b186b9 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/6] 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 c302d2b6d956c..13378f7ac446e 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>());
>From b77006fa6dfa7008a7c44b05c21745765259781e Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Wed, 17 Sep 2025 10:49:29 -0700
Subject: [PATCH 5/6] Add more tests
---
mlir/test/Dialect/Vector/canonicalize.mlir | 27 +++++++++++++++++++++-
1 file changed, 26 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 8dedde2ade1ce..08d28be3f8f73 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3380,12 +3380,37 @@ func.func @negative_from_elements_to_constant() -> vector<1x!llvm.ptr> {
// CHECK-LABEL: @negative_from_elements_poison
// CHECK: %[[VAL:.*]] = ub.poison : vector<2xf32>
// CHECK: return %[[VAL]] : vector<2xf32>
-func.func @negative_from_elements_poison() -> vector<2xf32> {
+func.func @negative_from_elements_poison_f32() -> vector<2xf32> {
%0 = ub.poison : f32
%1 = vector.from_elements %0, %0 : vector<2xf32>
return %1 : vector<2xf32>
}
+// -----
+
+// CHECK-LABEL: @negative_from_elements_poison_i32
+// CHECK: %[[VAL:.*]] = ub.poison : vector<2xi32>
+// CHECK: return %[[VAL]] : vector<2xi32>
+func.func @negative_from_elements_poison_i32() -> vector<2xi32> {
+ %0 = ub.poison : i32
+ %1 = vector.from_elements %0, %0 : vector<2xi32>
+ return %1 : vector<2xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @negative_from_elements_poison_constant_mix
+// CHECK: %[[POISON:.*]] = ub.poison : f32
+// CHECK: %[[CONST:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK: %[[RES:.*]] = vector.from_elements %[[POISON]], %[[CONST]] : vector<2xf32>
+// CHECK: return %[[RES]] : vector<2xf32>
+func.func @negative_from_elements_poison_constant_mix() -> vector<2xf32> {
+ %0 = ub.poison : f32
+ %c = arith.constant 1.0 : f32
+ %1 = vector.from_elements %0, %c : vector<2xf32>
+ return %1 : vector<2xf32>
+}
+
// +---------------------------------------------------------------------------
// End of Tests for foldFromElementsToConstant
// +---------------------------------------------------------------------------
>From 6113875c90e686a4e8cb699551b3e0860003be4f Mon Sep 17 00:00:00 2001
From: Jhalak Patel <jhalakp at nvidia.com>
Date: Thu, 18 Sep 2025 13:00:51 -0700
Subject: [PATCH 6/6] Fix formatting issues
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 25 +++++++++++++-----------
1 file changed, 14 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 13378f7ac446e..347141e2773b8 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -397,7 +397,7 @@ std::optional<int64_t> vector::getConstantVscaleMultiplier(Value value) {
}
/// Converts an IntegerAttr to have the specified type if needed.
-/// This handles cases where integer constant attributes have a different type
+/// 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)
@@ -2462,9 +2462,9 @@ static OpFoldResult foldFromElementsToElements(FromElementsOp fromElementsOp) {
static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
ArrayRef<Attribute> elements) {
// Check for null or poison attributes before any processing.
- if (llvm::any_of(elements, [](Attribute attr) {
- return !attr || isa<ub::PoisonAttrInterface>(attr);
- }))
+ if (llvm::any_of(elements, [](Attribute attr) {
+ return !attr || isa<ub::PoisonAttrInterface>(attr);
+ }))
return {};
// DenseElementsAttr only supports int/index/float/complex types.
@@ -2473,13 +2473,16 @@ static OpFoldResult foldFromElementsToConstant(FromElementsOp fromElementsOp,
if (!destEltType.isIntOrIndexOrFloat() && !isa<ComplexType>(destEltType))
return {};
- // 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
- });
+ // 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
+ });
return DenseElementsAttr::get(destVecType, convertedElements);
}
More information about the Mlir-commits
mailing list