[Mlir-commits] [mlir] [mlir][arith] Fix addui_extended fold assert on non-TypedAttr operands (PR #179140)

[Samarth Narang]

https://github.com/snarang181 updated https://github.com/llvm/llvm-project/pull/179140

>From 179d89274ba1b8d5ed148f1f58de5bb2b23355aa Mon Sep 17 00:00:00 2001
From: Samarth Narang <snarang at utexas.edu>
Date: Sun, 1 Feb 2026 16:29:41 -0500
Subject: [PATCH 1/2] [mlir][arith] Fix addui_extended fold assert on
 non-TypedAttr operands

---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp    |  6 ++++--
 mlir/test/Dialect/Arith/canonicalize.mlir | 10 ++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 83b3cec8d41af..d8dcf2f6c088e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -459,10 +459,12 @@ arith::AddUIExtendedOp::fold(FoldAdaptor adaptor,
   if (Attribute sumAttr = constFoldBinaryOp<IntegerAttr>(
           adaptor.getOperands(),
           [](APInt a, const APInt &b) { return std::move(a) + b; })) {
+    auto typedSumAttr = llvm::dyn_cast<TypedAttr>(sumAttr);
+    if (!typedSumAttr)
+      return failure();
     Attribute overflowAttr = constFoldBinaryOp<IntegerAttr>(
         ArrayRef({sumAttr, adaptor.getLhs()}),
-        getI1SameShape(llvm::cast<TypedAttr>(sumAttr).getType()),
-        calculateUnsignedOverflow);
+        getI1SameShape(typedSumAttr.getType()), calculateUnsignedOverflow);
     if (!overflowAttr)
       return failure();
 
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 18e0d2d2ea3c4..00e2bad408b42 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -1471,6 +1471,16 @@ func.func @adduiExtendedConstantsSplatVector() -> (vector<4xi32>, vector<4xi1>)
   return %sum, %overflow : vector<4xi32>, vector<4xi1>
 }
 
+// CHECK-LABEL: @adduiExtendedDoesNotAssertOnPoison
+// CHECK: %[[SUM:.+]], %[[OV:.+]] = arith.addui_extended %{{.+}}, %{{.+}} : tensor<1xi32>, tensor<1xi1>
+// CHECK: return %[[SUM]], %[[OV]] : tensor<1xi32>, tensor<1xi1>
+func.func @adduiExtendedDoesNotAssertOnPoison() -> (tensor<1xi32>, tensor<1xi1>) {
+  %c0 = arith.constant dense<0> : tensor<1xi32>
+  %p = ub.poison : tensor<1xi32>
+  %sum, %overflow = arith.addui_extended %c0, %p : tensor<1xi32>, tensor<1xi1>
+  return %sum, %overflow : tensor<1xi32>, tensor<1xi1>
+}
+
 // CHECK-LABEL: @mulsiExtendedZeroRhs
 //  CHECK-NEXT:   %[[zero:.+]] = arith.constant 0 : i32
 //  CHECK-NEXT:   return %[[zero]], %[[zero]]

>From 5d9c1db66c572fa1b26a3768e4d720fe9d6ee28a Mon Sep 17 00:00:00 2001
From: Samarth Narang <70980689+snarang181 at users.noreply.github.com>
Date: Tue, 17 Feb 2026 09:14:17 -0500
Subject: [PATCH 2/2] Apply suggestion from @kuhar

Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index d8dcf2f6c088e..c446ccc672b9f 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -456,7 +456,7 @@ arith::AddUIExtendedOp::fold(FoldAdaptor adaptor,
   // Let the `constFoldBinaryOp` utility attempt to fold the sum of both
   // operands. If that succeeds, calculate the overflow bit based on the sum
   // and the first (constant) operand, `lhs`.
-  if (Attribute sumAttr = constFoldBinaryOp<IntegerAttr>(
+  if (auto sumAttr = dyn_cast_if_present<TypedAttr>(constFoldBinaryOp<IntegerAttr>(
           adaptor.getOperands(),
           [](APInt a, const APInt &b) { return std::move(a) + b; })) {
     auto typedSumAttr = llvm::dyn_cast<TypedAttr>(sumAttr);