From 3433af7c8330187fed3a935d50780ad6fccc1edd Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 11:50:04 -0700 Subject: [PATCH] Fix NonNegativeResultAnalysis square rule to compare Values not defining ops The (mul a a) non-negativity rule compared the operands' defining ops. Two distinct block arguments both have a null defining op, so null == null wrongly proved mul(%arg0, %arg1) non-negative even when the product can be negative. Compare the operand Values directly instead, and add a regression test. Fixes #2648 --- src/enzyme_ad/jax/Utils.cpp | 8 ++++---- test/lit_tests/abspositive.mlir | 9 +++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/enzyme_ad/jax/Utils.cpp b/src/enzyme_ad/jax/Utils.cpp index 402c8fa9db..ef630bb906 100644 --- a/src/enzyme_ad/jax/Utils.cpp +++ b/src/enzyme_ad/jax/Utils.cpp @@ -987,10 +987,10 @@ NonNegativeResultAnalysis::State NonNegativeResultAnalysis::localGuaranteed( // (mul a a) is always non-negative if (auto mulOp = dyn_cast(op)) { - auto lhsOp = mulOp.getLhs().getDefiningOp(); - auto rhsOp = mulOp.getRhs().getDefiningOp(); - - if (lhsOp == rhsOp) { + // Compare the operand Values, not their defining ops: two distinct block + // arguments both have a null defining op, so a defining-op comparison + // would wrongly prove mul(%arg0, %arg1) non-negative. + if (mulOp.getLhs() == mulOp.getRhs()) { return State::GUARANTEED; } } diff --git a/test/lit_tests/abspositive.mlir b/test/lit_tests/abspositive.mlir index 079597880d..f2afaad02f 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]} : tensor<4xf64> // CHECK-NEXT: return %0 : tensor<4xf64> // CHECK-NEXT: } + +// mul of two distinct block args must NOT be proven non-negative: the product +// is negative whenever the args have opposite signs, so the abs has to stay. +func.func @test9(%arg0: tensor<12xf64>, %arg1: tensor<12xf64>) -> tensor<12xf64> { + %0 = stablehlo.multiply %arg0, %arg1 : tensor<12xf64> + // CHECK: stablehlo.abs + %1 = stablehlo.abs %0 : tensor<12xf64> + return %1 : tensor<12xf64> +}