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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

 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

---
Full diff: https://github.com/llvm/llvm-project/pull/179140.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+4-2) 
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+10) 


``````````diff
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]]

``````````

</details>


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