[Mlir-commits] [mlir] [mlir][shape] Don't crash on poison operands in folders (PR #226671)
Chaitanya Mamatha Ananda
llvmlistbot at llvm.org
Mon Oct 19 19:08:26 PDT 2026
https://github.com/chaitanyaupp18 updated https://github.com/llvm/llvm-project/pull/226671
>From 4657f862dae621cbe743f322935b4ac8be6f0e5e Mon Sep 17 00:00:00 2001
From: Chaitanya Mamatha Ananda <chaitanyaupp18 at gmail.com>
Date: Sat, 26 Sep 2026 03:18:57 -0700
Subject: [PATCH 1/2] [mlir][shape] Don't crash on poison operands in folders
Several shape folders used an unchecked llvm::cast on their folded
operands, assuming any non-null attribute was a DenseIntElementsAttr
(or IntegerAttr). A ub.poison operand folds to #ub.poison, so these
casts asserted. Check with dyn_cast_if_present first and skip folding
if it's anything else, like the other folders in this file already do.
Fixed in ConcatOp, CstrBroadcastableOp (and its
hasAtMostSingleNonScalar helper), NumElementsOp, SplitAtOp and
ToExtentTensorOp. The last three aren't in the issue; I found them
while fixing the reported ones.
Part of #226067 (the affine and scf cases are not addressed here).
Assisted-by: Claude
---
mlir/lib/Dialect/Shape/IR/Shape.cpp | 44 ++++++-----
mlir/test/Dialect/Shape/canonicalize.mlir | 89 +++++++++++++++++++++++
2 files changed, 113 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index e01b705e5610a..71d1ce9b39631 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -877,12 +877,12 @@ void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
//===----------------------------------------------------------------------===//
OpFoldResult ConcatOp::fold(FoldAdaptor adaptor) {
- if (!adaptor.getLhs() || !adaptor.getRhs())
+ auto lhs = llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getLhs());
+ auto rhs = llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getRhs());
+ if (!lhs || !rhs)
return nullptr;
- auto lhsShape = llvm::to_vector<6>(
- llvm::cast<DenseIntElementsAttr>(adaptor.getLhs()).getValues<int64_t>());
- auto rhsShape = llvm::to_vector<6>(
- llvm::cast<DenseIntElementsAttr>(adaptor.getRhs()).getValues<int64_t>());
+ auto lhsShape = llvm::to_vector<6>(lhs.getValues<int64_t>());
+ auto rhsShape = llvm::to_vector<6>(rhs.getValues<int64_t>());
SmallVector<int64_t, 6> resultShape;
resultShape.append(lhsShape.begin(), lhsShape.end());
resultShape.append(rhsShape.begin(), rhsShape.end());
@@ -983,7 +983,9 @@ void CstrBroadcastableOp::getCanonicalizationPatterns(
static bool hasAtMostSingleNonScalar(ArrayRef<Attribute> attributes) {
bool nonScalarSeen = false;
for (Attribute a : attributes) {
- if (!a || llvm::cast<DenseIntElementsAttr>(a).getNumElements() != 0) {
+ // Treat non-constant and poison operands as potentially non-scalar.
+ auto dense = llvm::dyn_cast_if_present<DenseIntElementsAttr>(a);
+ if (!dense || dense.getNumElements() != 0) {
if (nonScalarSeen)
return false;
nonScalarSeen = true;
@@ -1000,10 +1002,10 @@ OpFoldResult CstrBroadcastableOp::fold(FoldAdaptor adaptor) {
if ([&] {
SmallVector<SmallVector<int64_t, 6>, 6> extents;
for (const auto &operand : adaptor.getShapes()) {
- if (!operand)
+ auto dense = llvm::dyn_cast_if_present<DenseIntElementsAttr>(operand);
+ if (!dense)
return false;
- extents.push_back(llvm::to_vector<6>(
- llvm::cast<DenseIntElementsAttr>(operand).getValues<int64_t>()));
+ extents.push_back(llvm::to_vector<6>(dense.getValues<int64_t>()));
}
return OpTrait::util::staticallyKnownBroadcastable(extents);
}())
@@ -1571,12 +1573,13 @@ LogicalResult shape::RankOp::verify() { return verifySizeOrIndexOp(*this); }
OpFoldResult NumElementsOp::fold(FoldAdaptor adaptor) {
// Fold only when argument constant.
- Attribute shape = adaptor.getShape();
+ auto shape =
+ llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getShape());
if (!shape)
return {};
APInt product(64, 1);
- for (auto value : llvm::cast<DenseIntElementsAttr>(shape))
+ for (auto value : shape)
product *= value;
Builder builder(getContext());
return builder.getIndexAttr(product.getLimitedValue());
@@ -1913,13 +1916,14 @@ LogicalResult shape::YieldOp::verify() {
LogicalResult SplitAtOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<OpFoldResult> &results) {
- if (!adaptor.getOperand() || !adaptor.getIndex())
+ auto shapeAttr =
+ llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getOperand());
+ auto indexAttr = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getIndex());
+ if (!shapeAttr || !indexAttr)
return failure();
- auto shapeVec =
- llvm::to_vector<6>(llvm::cast<DenseIntElementsAttr>(adaptor.getOperand())
- .getValues<int64_t>());
+ auto shapeVec = llvm::to_vector<6>(shapeAttr.getValues<int64_t>());
auto shape = llvm::ArrayRef(shapeVec);
- auto splitPoint = llvm::cast<IntegerAttr>(adaptor.getIndex()).getInt();
+ auto splitPoint = indexAttr.getInt();
// Verify that the split point is in the correct range.
// TODO: Constant fold to an "error".
int64_t rank = shape.size();
@@ -1938,12 +1942,12 @@ LogicalResult SplitAtOp::fold(FoldAdaptor adaptor,
//===----------------------------------------------------------------------===//
OpFoldResult ToExtentTensorOp::fold(FoldAdaptor adaptor) {
- if (!adaptor.getInput())
+ auto inputAttr =
+ llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getInput());
+ if (!inputAttr)
return OpFoldResult();
Builder builder(getContext());
- auto shape =
- llvm::to_vector<6>(llvm::cast<DenseIntElementsAttr>(adaptor.getInput())
- .getValues<int64_t>());
+ auto shape = llvm::to_vector<6>(inputAttr.getValues<int64_t>());
auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
builder.getIndexType());
return DenseIntElementsAttr::get(type, shape);
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 89aa167360168..b8399e23ab811 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1685,3 +1685,92 @@ func.func @extract_from_shape_of_memref_no_fold(%arg0: memref<?xf32>) -> index {
%dim = tensor.extract %shape[%c0] : tensor<1xindex>
return %dim : index
}
+
+// -----
+
+// Regression tests for https://github.com/llvm/llvm-project/issues/226067:
+// several shape folders used to crash with an unchecked cast when an operand
+// was ub.poison (a non-DenseIntElementsAttr attribute). The folds must bail
+// out gracefully instead.
+// CHECK-LABEL: func @concat_poison
+func.func @concat_poison() -> tensor<4xindex> {
+ // CHECK: shape.concat
+ %lhs = shape.const_shape [0, 1] : tensor<2xindex>
+ %rhs = ub.poison : tensor<2xindex>
+ %0 = shape.concat %lhs, %rhs : tensor<2xindex>, tensor<2xindex> -> tensor<4xindex>
+ return %0 : tensor<4xindex>
+}
+
+// -----
+
+// A poison operand must not crash the folder.
+// CHECK-LABEL: func @cstr_broadcastable_poison
+func.func @cstr_broadcastable_poison(%arg : tensor<?x4xf32>) {
+ // CHECK: shape.cstr_broadcastable
+ %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex>
+ %1 = ub.poison : tensor<1xindex>
+ %2 = shape.cstr_broadcastable %0, %1 : tensor<2xindex>, tensor<1xindex>
+ "use"(%2) : (!shape.witness) -> ()
+ return
+}
+
+// -----
+
+// A poison operand must not crash the folder.
+// CHECK-LABEL: func @num_elements_poison
+func.func @num_elements_poison() -> index {
+ // CHECK: shape.num_elements
+ %s = ub.poison : tensor<2xindex>
+ %n = shape.num_elements %s : tensor<2xindex> -> index
+ return %n : index
+}
+
+// -----
+
+// A poison operand must not crash the folder.
+// CHECK-LABEL: func @split_at_poison_shape
+// CHECK: shape.split_at
+func.func @split_at_poison_shape() -> (!shape.shape, !shape.shape) {
+ %c2 = arith.constant 2 : index
+ %0 = ub.poison : !shape.shape
+ %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, index) -> (!shape.shape, !shape.shape)
+ return %head, %tail : !shape.shape, !shape.shape
+}
+
+// -----
+
+// A poison operand must not crash the folder.
+// CHECK-LABEL: func @split_at_poison_index
+// CHECK: shape.split_at
+func.func @split_at_poison_index() -> (!shape.shape, !shape.shape) {
+ %c2 = ub.poison : index
+ %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
+ %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, index) -> (!shape.shape, !shape.shape)
+ return %head, %tail : !shape.shape, !shape.shape
+}
+
+// -----
+
+// A poison operand must not crash the folder.
+// CHECK-LABEL: func @to_extent_tensor_poison
+// CHECK: shape.to_extent_tensor
+func.func @to_extent_tensor_poison() -> tensor<2xindex> {
+ %cs = ub.poison : !shape.shape
+ %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
+ return %0 : tensor<2xindex>
+}
+
+// -----
+
+// A poison operand must not crash the folder, even when it is only reached
+// after the scalar check has already given up.
+// CHECK-LABEL: func @cstr_broadcastable_poison_after_non_scalars
+func.func @cstr_broadcastable_poison_after_non_scalars() {
+ // CHECK: shape.cstr_broadcastable
+ %0 = shape.const_shape [1, 3] : tensor<2xindex>
+ %1 = shape.const_shape [3] : tensor<1xindex>
+ %2 = ub.poison : tensor<1xindex>
+ %3 = shape.cstr_broadcastable %0, %1, %2 : tensor<2xindex>, tensor<1xindex>, tensor<1xindex>
+ "use"(%3) : (!shape.witness) -> ()
+ return
+}
>From f5c001ea20660abc1ba813e13091ca51291900a0 Mon Sep 17 00:00:00 2001
From: Chaitanya Mamatha Ananda <chaitanyaupp18 at gmail.com>
Date: Mon, 19 Oct 2026 19:19:50 -0700
Subject: [PATCH 2/2] [mlir][shape] Fold to poison instead of bailing out
As suggested in review: num_elements, to_extent_tensor and split_at
now fold to poison when their input is poison. concat only does this
when both inputs are poison, since with one known input part of the
result is still known.
---
mlir/lib/Dialect/Shape/IR/Shape.cpp | 20 ++++++++++
mlir/test/Dialect/Shape/canonicalize.mlir | 45 ++++++++++++++++++-----
2 files changed, 55 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 71d1ce9b39631..ebb56241e0169 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -877,6 +877,10 @@ void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
//===----------------------------------------------------------------------===//
OpFoldResult ConcatOp::fold(FoldAdaptor adaptor) {
+ // concat(poison, poison) is poison.
+ if (isa_and_present<ub::PoisonAttr>(adaptor.getLhs()) &&
+ isa_and_present<ub::PoisonAttr>(adaptor.getRhs()))
+ return ub::PoisonAttr::get(getContext());
auto lhs = llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getLhs());
auto rhs = llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getRhs());
if (!lhs || !rhs)
@@ -1571,6 +1575,9 @@ LogicalResult shape::RankOp::verify() { return verifySizeOrIndexOp(*this); }
//===----------------------------------------------------------------------===//
OpFoldResult NumElementsOp::fold(FoldAdaptor adaptor) {
+ // num_elements of poison shape is poison.
+ if (isa_and_present<ub::PoisonAttr>(adaptor.getShape()))
+ return ub::PoisonAttr::get(getContext());
// Fold only when argument constant.
auto shape =
@@ -1916,6 +1923,15 @@ LogicalResult shape::YieldOp::verify() {
LogicalResult SplitAtOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<OpFoldResult> &results) {
+ // Splitting a poison shape, or at a poison index, gives two poison halves.
+ if (isa_and_present<ub::PoisonAttr>(adaptor.getOperand()) ||
+ isa_and_present<ub::PoisonAttr>(adaptor.getIndex())) {
+ auto poison = ub::PoisonAttr::get(getContext());
+ results.push_back(poison);
+ results.push_back(poison);
+ return success();
+ }
+
auto shapeAttr =
llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getOperand());
auto indexAttr = llvm::dyn_cast_if_present<IntegerAttr>(adaptor.getIndex());
@@ -1942,6 +1958,10 @@ LogicalResult SplitAtOp::fold(FoldAdaptor adaptor,
//===----------------------------------------------------------------------===//
OpFoldResult ToExtentTensorOp::fold(FoldAdaptor adaptor) {
+ // to_extent_tensor of a poison shape is poison.
+ if (isa_and_present<ub::PoisonAttr>(adaptor.getInput()))
+ return ub::PoisonAttr::get(getContext());
+
auto inputAttr =
llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getInput());
if (!inputAttr)
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index b8399e23ab811..b3408433bbde6 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1690,8 +1690,11 @@ func.func @extract_from_shape_of_memref_no_fold(%arg0: memref<?xf32>) -> index {
// Regression tests for https://github.com/llvm/llvm-project/issues/226067:
// several shape folders used to crash with an unchecked cast when an operand
-// was ub.poison (a non-DenseIntElementsAttr attribute). The folds must bail
-// out gracefully instead.
+// was ub.poison (a non-DenseIntElementsAttr attribute). They now fold to
+// poison when the result depends entirely on the poison operand, and leave
+// the op alone otherwise.
+
+// With only one poison operand, the other extents are still known: no fold.
// CHECK-LABEL: func @concat_poison
func.func @concat_poison() -> tensor<4xindex> {
// CHECK: shape.concat
@@ -1703,6 +1706,20 @@ func.func @concat_poison() -> tensor<4xindex> {
// -----
+// concat of two poison shapes is poison.
+// CHECK-LABEL: func @concat_both_poison
+// CHECK-NOT: shape.concat
+// CHECK: %[[P:.*]] = ub.poison : tensor<4xindex>
+// CHECK: return %[[P]]
+func.func @concat_both_poison() -> tensor<4xindex> {
+ %lhs = ub.poison : tensor<2xindex>
+ %rhs = ub.poison : tensor<2xindex>
+ %0 = shape.concat %lhs, %rhs : tensor<2xindex>, tensor<2xindex> -> tensor<4xindex>
+ return %0 : tensor<4xindex>
+}
+
+// -----
+
// A poison operand must not crash the folder.
// CHECK-LABEL: func @cstr_broadcastable_poison
func.func @cstr_broadcastable_poison(%arg : tensor<?x4xf32>) {
@@ -1716,10 +1733,12 @@ func.func @cstr_broadcastable_poison(%arg : tensor<?x4xf32>) {
// -----
-// A poison operand must not crash the folder.
+// num_elements of a poison shape is poison.
// CHECK-LABEL: func @num_elements_poison
+// CHECK-NOT: shape.num_elements
+// CHECK: %[[P:.*]] = ub.poison : index
+// CHECK: return %[[P]]
func.func @num_elements_poison() -> index {
- // CHECK: shape.num_elements
%s = ub.poison : tensor<2xindex>
%n = shape.num_elements %s : tensor<2xindex> -> index
return %n : index
@@ -1727,9 +1746,11 @@ func.func @num_elements_poison() -> index {
// -----
-// A poison operand must not crash the folder.
+// Splitting a poison shape gives two poison halves.
// CHECK-LABEL: func @split_at_poison_shape
-// CHECK: shape.split_at
+// CHECK-NOT: shape.split_at
+// CHECK: %[[P:.*]] = ub.poison : !shape.shape
+// CHECK: return %[[P]], %[[P]]
func.func @split_at_poison_shape() -> (!shape.shape, !shape.shape) {
%c2 = arith.constant 2 : index
%0 = ub.poison : !shape.shape
@@ -1739,9 +1760,11 @@ func.func @split_at_poison_shape() -> (!shape.shape, !shape.shape) {
// -----
-// A poison operand must not crash the folder.
+// Splitting at a poison index gives two poison halves.
// CHECK-LABEL: func @split_at_poison_index
-// CHECK: shape.split_at
+// CHECK-NOT: shape.split_at
+// CHECK: %[[P:.*]] = ub.poison : !shape.shape
+// CHECK: return %[[P]], %[[P]]
func.func @split_at_poison_index() -> (!shape.shape, !shape.shape) {
%c2 = ub.poison : index
%0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
@@ -1751,9 +1774,11 @@ func.func @split_at_poison_index() -> (!shape.shape, !shape.shape) {
// -----
-// A poison operand must not crash the folder.
+// to_extent_tensor of a poison shape is poison.
// CHECK-LABEL: func @to_extent_tensor_poison
-// CHECK: shape.to_extent_tensor
+// CHECK-NOT: shape.to_extent_tensor
+// CHECK: %[[P:.*]] = ub.poison : tensor<2xindex>
+// CHECK: return %[[P]]
func.func @to_extent_tensor_poison() -> tensor<2xindex> {
%cs = ub.poison : !shape.shape
%0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
More information about the Mlir-commits
mailing list