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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: William Moses (wsmoses)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

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


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

``````````

</details>


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