[Mlir-commits] [mlir] [mlir][shape] Don't crash on poison operands in folders (PR #226671)

Sat Sep 26 03:26:47 PDT 2026

chaitanyaupp18 wrote:

Thanks! I have read the AI Tool Use Policy, the Code-Review Policy and the
Developer Policy. I used Claude to help me understand the code and review
my changes, this is noted in the PR description. I reproduced each crash,
wrote and verified the fix and tests myself, and can answer questions about
any part of the change.

On Sat, Sep 26, 2026 at 3:23 AM llvmorg-github-actions[bot] <
***@***.***> wrote:

> *llvmorg-github-actions[bot]* left a comment (llvm/llvm-project#226671)
> <https://github.com/llvm/llvm-project/pull/226671#issuecomment-5845449755>
>
> @llvm/pr-subscribers-mlir-shape
> <https://github.com/orgs/llvm/teams/pr-subscribers-mlir-shape>
>
> Author: Chaitanya Mamatha Ananda (chaitanyaupp18)
> Changes
>
> 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 --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+}
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/llvm/llvm-project/pull/226671?email_source=notifications&email_token=AKNR56DE42DVBG4NJLIYKF35Q6KK3A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOBUGU2DIOJXGU22M4TFMFZW63VHNVSW45DJN5XKKZLWMVXHJLDGN5XXIZLSL5RWY2LDNM#issuecomment-5845449755>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AKNR56CDYN2EZLQ47UKJBUL5Q6KK3AVCNFSNUABEKJSXA33TNF2G64TZHM3TKOBSGE2DGMR3JFZXG5LFHM2TKOJTGI2DMNJUHGQXMAQ>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

https://github.com/llvm/llvm-project/pull/226671