[Mlir-commits] [mlir] [mlir][arith] Do not preserve nneg when folding extui of extui (PR #216315)

[William Moses]

https://github.com/wsmoses updated https://github.com/llvm/llvm-project/pull/216315

>From 0ecf7ab2f5c7d7169055d553cfcde7f0eb015858 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Fri, 14 Aug 2026 08:12:10 -0500
Subject: [PATCH] [mlir][arith] Do not preserve nneg when folding extui of
 extui

The extui(extui(x)) fold reassigns the outer extension's source while
keeping its nneg flag. The flag asserts the source is non-negative as a
signed value, about which the outer flag says nothing once the source
changes: extui nneg i8 to i32 of extui i1 to i8 is satisfied for every
boolean byte, but the merged extui nneg i1 to i32 is poison whenever the
bool is true. Take the nneg flag from the inner extension, which is the
one that speaks about the surviving source.

Co-Authored-By: Claude Fable 5 <noreply at anthropic.com>
Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
---
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp    |  1 +
 mlir/test/Dialect/Arith/canonicalize.mlir | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 1b7ef01ec64ac..d06a93876c582 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1822,6 +1822,7 @@ convertFloatValue(APFloat sourceValue,
 
 OpFoldResult arith::ExtUIOp::fold(FoldAdaptor adaptor) {
   if (auto lhs = getIn().getDefiningOp<ExtUIOp>()) {
+    setNonNeg(lhs.getNonNeg());
     getInMutable().assign(lhs.getIn());
     return getResult();
   }
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 3761d434b5254..dd3c3abc45bf2 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -451,6 +451,24 @@ func.func @extUIOfExtUI(%arg0: i1) -> i64 {
   return %ext2 : i64
 }
 
+// CHECK-LABEL: @extUIOfExtUI_outer_nneg
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 : i1 to i64
+//       CHECK:   return %[[res]]
+func.func @extUIOfExtUI_outer_nneg(%arg0: i1) -> i64 {
+  %ext1 = arith.extui %arg0 : i1 to i8
+  %ext2 = arith.extui %ext1 nneg : i8 to i64
+  return %ext2 : i64
+}
+
+// CHECK-LABEL: @extUIOfExtUI_nneg
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 nneg : i8 to i64
+//       CHECK:   return %[[res]]
+func.func @extUIOfExtUI_nneg(%arg0: i8) -> i64 {
+  %ext1 = arith.extui %arg0 nneg : i8 to i32
+  %ext2 = arith.extui %ext1 nneg : i32 to i64
+  return %ext2 : i64
+}
+
 // CHECK-LABEL: @extSIOfExtSI
 //       CHECK:   %[[res:.+]] = arith.extsi %arg0 : i1 to i64
 //       CHECK:   return %[[res]]