[Mlir-commits] [mlir] [mlir][shape] Don't crash on poison operands in folders (PR #226671)
Chaitanya Mamatha Ananda
llvmlistbot at llvm.org
Sat Sep 26 03:22:19 PDT 2026
https://github.com/chaitanyaupp18 created https://github.com/llvm/llvm-project/pull/226671
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.
This change checks with `dyn_cast_if_present` first and skips folding if the attribute is anything else, like the other folders in this file already do.
### Folders fixed
- `ConcatOp` and `CstrBroadcastableOp` (including the `hasAtMostSingleNonScalar` helper), reported in the issue
- `NumElementsOp`, `SplitAtOp` and `ToExtentTensorOp`, found while fixing the reported ones
### Notes
- In `hasAtMostSingleNonScalar`, poison is treated like a non-constant operand. Folding to `true` is still valid there, because the result doesn't depend on that operand's value.
- Each new test in `mlir/test/Dialect/Shape/canonicalize.mlir` crashes without this change and passes with it.
Part of #226067 (the affine and scf cases are not addressed here).
Assisted-by: Claude
>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] [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
+}
More information about the Mlir-commits
mailing list