[Mlir-commits] [mlir] [mlir][shape] Don't crash on poison operands in folders (PR #226671)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 26 03:23:24 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Chaitanya Mamatha Ananda (chaitanyaupp18)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/226671.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+24-20)
- (modified) mlir/test/Dialect/Shape/canonicalize.mlir (+89)
``````````diff
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
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/226671
More information about the Mlir-commits
mailing list