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

[llvmlistbot at llvm.org]

Author: William Moses
Date: 2026-08-14T15:13:45Z
New Revision: ceca5336c5d57555bd098ee630f140bf22502c17

URL: https://github.com/llvm/llvm-project/commit/ceca5336c5d57555bd098ee630f140bf22502c17
DIFF: https://github.com/llvm/llvm-project/commit/ceca5336c5d57555bd098ee630f140bf22502c17.diff

LOG: [mlir][arith] Do not preserve nneg when folding extui of extui (#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.

Assisted-By: Claude

---------

Co-authored-by: Claude Fable 5 <noreply at anthropic.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/IR/ArithOps.cpp
    mlir/test/Dialect/Arith/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 1b7ef01ec64ac..07c7f99999424 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -1822,6 +1822,9 @@ 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 3761d434b5254..e8e72926cb563 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -451,6 +451,33 @@ 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_inner_nneg
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 nneg : i8 to i64
+//       CHECK:   return %[[res]]
+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
+}
+
 // CHECK-LABEL: @extSIOfExtSI
 //       CHECK:   %[[res:.+]] = arith.extsi %arg0 : i1 to i64
 //       CHECK:   return %[[res]]