From 66bf2011e7702453b06ac1a13d4f44e29c99d90b Mon Sep 17 00:00:00 2001
From: June Kim <kimjune01@gmail.com>
Date: Wed, 22 Jul 2026 13:20:21 -0700
Subject: [PATCH] Restrict bitwise-op non-negativity to i1 results

NonNegativeResultAnalysis marked and/or/xor/not as unconditionally
non-negative, which is unsound for signed integers: not(0) == -1, and
or/xor/and can set the sign bit. Guarantee them only for boolean (i1)
results, where the value is provably in {0, 1}. Add a regression test:
abs(not(%i32)) must keep its abs.

Fixes #2651
---
 src/enzyme_ad/jax/Utils.cpp     | 11 +++++++++--
 test/lit_tests/abspositive.mlir |  9 +++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/enzyme_ad/jax/Utils.cpp b/src/enzyme_ad/jax/Utils.cpp
index 402c8fa9db..8118b26dca 100644
--- a/src/enzyme_ad/jax/Utils.cpp
+++ b/src/enzyme_ad/jax/Utils.cpp
@@ -952,8 +952,15 @@ NonNegativeResultAnalysis::State NonNegativeResultAnalysis::localGuaranteed(
 
   // integer ops
   if (isa<stablehlo::AbsOp, stablehlo::SqrtOp, stablehlo::ExpOp,
-          stablehlo::IotaOp, stablehlo::AndOp, stablehlo::OrOp,
-          stablehlo::XorOp, stablehlo::NotOp>(op)) {
+          stablehlo::IotaOp>(op)) {
+    return State::GUARANTEED;
+  }
+
+  // Bitwise ops are non-negative only for boolean (i1) results: on a signed
+  // integer not(0) == -1, and or/xor/and can set the sign bit.
+  if (isa<stablehlo::AndOp, stablehlo::OrOp, stablehlo::XorOp,
+          stablehlo::NotOp>(op) &&
+      cast<ShapedType>(val.getType()).getElementType().isInteger(1)) {
     return State::GUARANTEED;
   }
 
diff --git a/test/lit_tests/abspositive.mlir b/test/lit_tests/abspositive.mlir
index 079597880d..e2970b4e8c 100644
--- a/test/lit_tests/abspositive.mlir
+++ b/test/lit_tests/abspositive.mlir
@@ -95,3 +95,12 @@ func.func @test8(%arg0: tensor<4xf64>) -> tensor<4xf64> {
 // CHECK-NEXT:     %0 = stablehlo.logistic %arg0 {enzymexla.non_negative = [#enzymexla<guaranteed GUARANTEED>]} : tensor<4xf64>
 // CHECK-NEXT:     return %0 : tensor<4xf64>
 // CHECK-NEXT: }
+
+// A bitwise op on a signed integer is NOT non-negative: not(0) == -1, so the
+// abs must stay.
+func.func @test_signed_bitwise(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+    %0 = stablehlo.not %arg0 : tensor<4xi32>
+    // CHECK: stablehlo.abs
+    %1 = stablehlo.abs %0 : tensor<4xi32>
+    return %1 : tensor<4xi32>
+}