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

[Samarth Narang]

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

 In the presence of poison (e.g., ub.poison), constant folding may produce a non-TypedAttr attribute, causing an assertion failure. Guard the cast and bail out of folding when the folded attribute is not a TypedAttr. 
Also adds a test exercising addui_extended with ub.poison.

Fixes https://github.com/llvm/llvm-project/issues/179080

>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] [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]]