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

[William Moses]

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

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:

```mlir
%w = arith.extui %b : i1 to i8
%r = arith.extui %w nneg : i8 to i32   // satisfied for every bool byte
```

folds to `arith.extui %b nneg : i1 to i32`, which is poison whenever `%b` is true (the signed value of i1 `1` is −1). Downstream, LLVM is then entitled to fold the bool to false everywhere the merged value flows; we hit this as a real end-to-end miscompile where a kernel's boolean-derived launch arguments were all evaluated with the flag's false arm.

Keep `nneg` on the merged extension only if the inner extension carries it — the inner flag is the one that speaks about the surviving source.

>From 668bea35ecee5b5c573e6355cd084348cdb9af45 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. Keep nneg only if the inner extension carries it.

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 | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 1b7ef01ec64ac..451c24c4df6bd 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>()) {
+    if (!lhs.getNonNeg())
+      setNonNeg(false);
     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]]