[Mlir-commits] [mlir] [mlir][Arith] Fix crash when folding operations with dynamic-shaped tensors (PR #178428)

[llvmlistbot at llvm.org]

https://github.com/mugiwaraluffy56 updated https://github.com/llvm/llvm-project/pull/178428

>From 824540f61a343b428c481e47912dfa47b8b601ca Mon Sep 17 00:00:00 2001
From: mugiwaraluffy56 <myakampuneeth at gmail.com>
Date: Wed, 28 Jan 2026 19:39:45 +0530
Subject: [PATCH] [mlir][Arith] Fix crash when folding operations with
 dynamic-shaped tensors

DenseElementsAttr::get requires a static shape. When folding comparison
or select operations with dynamic-shaped tensor types, we cannot create
a DenseElementsAttr for the result. This patch adds checks to return
early (no fold) when the type has a dynamic shape.

Fixes #178415.

Co-Authored-By: Claude Opus 4.5 <noreply at anthropic.com>
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp    |  6 ++++++
 mlir/test/Dialect/Arith/canonicalize.mlir | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 565a537616971..0b631a0f0c2d0 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -151,6 +151,9 @@ static Attribute getBoolAttribute(Type type, bool value) {
   ShapedType shapedType = dyn_cast_or_null<ShapedType>(type);
   if (!shapedType)
     return boolAttr;
+  // DenseElementsAttr requires a static shape.
+  if (!shapedType.hasStaticShape())
+    return {};
   return DenseElementsAttr::get(shapedType, boolAttr);
 }
 
@@ -2527,6 +2530,9 @@ OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
   // select %cst_vec, %cst0, %cst1 => %cst2
   if (auto cond =
           dyn_cast_if_present<DenseElementsAttr>(adaptor.getCondition())) {
+    // DenseElementsAttr::get requires a static shape.
+    if (!cond.getType().hasStaticShape())
+      return nullptr;
     if (auto lhs =
             dyn_cast_if_present<DenseElementsAttr>(adaptor.getTrueValue())) {
       if (auto rhs =
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 3ad1530248809..dff3ef0386a95 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -3403,3 +3403,26 @@ func.func @unreachable() {
   cf.br ^unreachable
 }
 
+// -----
+
+// Verify that cmpi with dynamic-shaped tensors does not crash during folding.
+// The fold cannot create a DenseElementsAttr for dynamic shapes.
+// CHECK-LABEL: @cmpi_dynamic_shape_no_fold
+//       CHECK:   arith.cmpi eq
+func.func @cmpi_dynamic_shape_no_fold(%arg0: tensor<?xi32>) -> tensor<?xi1> {
+  %0 = arith.cmpi eq, %arg0, %arg0 : tensor<?xi32>
+  return %0 : tensor<?xi1>
+}
+
+// -----
+
+// Verify that select with constant condition and dynamic-shaped tensor operands
+// folds correctly without crashing.
+// CHECK-LABEL: @select_dynamic_shape_const_condition
+//       CHECK:   return %arg0
+func.func @select_dynamic_shape_const_condition(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
+  %true = arith.constant true
+  %0 = arith.select %true, %arg0, %arg1 : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+