[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 1/2] [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]]

>From 433aa229d3ba8a78af381c87f0e2af928a525144 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Fri, 14 Aug 2026 10:00:28 -0500
Subject: [PATCH 2/2] Address review: inner-only test, rename both_nneg, inline
 comment

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    |  2 ++
 mlir/test/Dialect/Arith/canonicalize.mlir | 13 +++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index d06a93876c582..07c7f99999424 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1822,6 +1822,8 @@ convertFloatValue(APFloat sourceValue,
 
 OpFoldResult arith::ExtUIOp::fold(FoldAdaptor adaptor) {
   if (auto lhs = getIn().getDefiningOp<ExtUIOp>()) {
+    // Only the inner extension's nneg speaks about the surviving source; the
+    // outer flag described the already-extended value.
     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 dd3c3abc45bf2..e8e72926cb563 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -460,10 +460,19 @@ func.func @extUIOfExtUI_outer_nneg(%arg0: i1) -> i64 {
   return %ext2 : i64
 }
 
-// CHECK-LABEL: @extUIOfExtUI_nneg
+// CHECK-LABEL: @extUIOfExtUI_inner_nneg
 //       CHECK:   %[[res:.+]] = arith.extui %arg0 nneg : i8 to i64
 //       CHECK:   return %[[res]]
-func.func @extUIOfExtUI_nneg(%arg0: i8) -> i64 {
+func.func @extUIOfExtUI_inner_nneg(%arg0: i8) -> i64 {
+  %ext1 = arith.extui %arg0 nneg : i8 to i32
+  %ext2 = arith.extui %ext1 : i32 to i64
+  return %ext2 : i64
+}
+
+// CHECK-LABEL: @extUIOfExtUI_both_nneg
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 nneg : i8 to i64
+//       CHECK:   return %[[res]]
+func.func @extUIOfExtUI_both_nneg(%arg0: i8) -> i64 {
   %ext1 = arith.extui %arg0 nneg : i8 to i32
   %ext2 = arith.extui %ext1 nneg : i32 to i64
   return %ext2 : i64