From 86eef28b59ebb6a20f63debbb36c239fba467995 Mon Sep 17 00:00:00 2001 From: June Kim Date: Mon, 22 Jun 2026 12:36:22 -0700 Subject: [PATCH 01/14] Guard LogSimplify rewrites against domain narrowing (#2570) Four LogSimplify sub-cases applied real-analysis identities valid only for positive bases, with no domain guard, so the optimized program returned NaN where the original computed a finite real: log(a*a) -> 2*log(a) NaN for any a < 0 log(pow(x,y)) -> y*log(x) NaN for x < 0 (even y), or y = 0 log(a*b) -> log(a)+log(b) NaN when the constant operand < 0 log(a/b) -> log(a)-log(b) NaN when the constant operand < 0 These are not a fast-math/precision question: as partial functions the two sides have different domains (log(a*a) is real for a != 0, 2*log(a) only for a > 0), so the divergence appears before floating point is involved. - log(a*a) -> 2*log(abs(a)): unconditionally domain-sound (a*a and abs(a) share a domain), so no guard needed; just take log of |a|. - log(pow(x,y)): fire only when x is provably non-negative. - log(a*b) / log(a/b): fire only when the constant operand is non-negative (a positive constant agrees, NaN-for-NaN, on negative inputs; a negative constant manufactures the NaN). Reuses the existing guaranteedNonNegativeResult analysis. Updates the two log(a*a) cases in logsimplify.mlir to the abs form and adds negative-constant / unknown-sign-base cases asserting the guarded rewrites no longer fire. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 44 ++++++++++++---- test/lit_tests/logsimplify.mlir | 64 ++++++++++++++++++++--- 2 files changed, 90 insertions(+), 18 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 3ad4efd519..c28ec3c2e5 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27390,7 +27390,11 @@ struct LogSimplify final { // log(pow(x, y)) -> y * log(x) auto defOp = op.getOperand().getDefiningOp(); - if (defOp) { + // Only sound when the base x is provably non-negative. For x < 0, + // log(pow(x, y)) can be a finite real (e.g. even integer y, where + // pow(x, y) > 0) while y * log(x) is NaN; likewise y = 0 gives + // log(pow(x, 0)) = 0 but 0 * log(x) = NaN. See issue #2570. + if (defOp && guaranteedNonNegativeResult(defOp.getLhs(), rewriter)) { rewriter.replaceOpWithNewOp( op, defOp.getRhs(), stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs())); @@ -27403,23 +27407,35 @@ struct LogSimplify final if (defOp) { auto lhs = defOp.getLhs(); auto rhs = defOp.getRhs(); - if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(a) + if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(abs(a)) + // a*a is non-negative for every real a, but 2 * log(a) is NaN for + // a < 0. Take the log of |a| so the rewrite keeps log(a*a)'s domain + // (log(a*a) = 2 * log|a|). See issue #2570. rewriter.replaceOpWithNewOp( op, stablehlo::ConstantOp::create( rewriter, op.getLoc(), lhs.getType(), cast(makeAttr(lhs.getType(), 2))), - stablehlo::LogOp::create(rewriter, op.getLoc(), lhs)); + stablehlo::LogOp::create( + rewriter, op.getLoc(), + stablehlo::AbsOp::create(rewriter, op.getLoc(), + lhs.getType(), lhs))); return success(); } if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(mul(a, b)) -> log(a) + // log(b) if a or b is constant - rewriter.replaceOpWithNewOp( - op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), - stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); - return success(); + // Only sound when the constant operand is non-negative: a negative + // constant makes log(const) NaN where log(a*b) was a finite real + // (a < 0 makes both sides NaN, so they still agree). Issue #2570. + Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; + if (guaranteedNonNegativeResult(cst, rewriter)) { + rewriter.replaceOpWithNewOp( + op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), + stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); + return success(); + } } } } @@ -27469,10 +27485,16 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(div(a, b)) -> log(a) - // log(b) if a or b is constant - rewriter.replaceOpWithNewOp( - op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), - stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); - return success(); + // Only sound when the constant operand is non-negative: a negative + // constant makes log(const) NaN where log(a/b) was a finite real. + // Issue #2570. + Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; + if (guaranteedNonNegativeResult(cst, rewriter)) { + rewriter.replaceOpWithNewOp( + op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), + stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); + return success(); + } } } } diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 215e350d37..2c8db51c3f 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -19,10 +19,11 @@ func.func @main2(%arg0: tensor) -> tensor { // CHECK: func.func @main2(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor -// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor -// CHECK-NEXT: return %2 : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.abs %0 : tensor +// CHECK-NEXT: %2 = stablehlo.log %1 : tensor +// CHECK-NEXT: %3 = stablehlo.multiply %cst, %2 : tensor +// CHECK-NEXT: return %3 : tensor // CHECK-NEXT: } func.func @main3(%arg0: tensor) -> tensor { @@ -74,9 +75,10 @@ func.func @main6(%arg0: tensor) -> tensor { // CHECK: func.func @main6(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor -// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor +// CHECK-NEXT: return %2 : tensor // CHECK-NEXT: } func.func @main7(%arg0: tensor) -> tensor { @@ -145,3 +147,51 @@ func.func @main11(%arg0: tensor) -> tensor { // CHECK-NEXT: %1 = stablehlo.divide %0, %cst : tensor // CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } + +// Domain guards (issue #2570): the split rewrites must NOT fire when they would +// narrow the domain and manufacture a NaN on inputs the original handles. + +// log(a * c) with a negative constant must stay as-is: splitting to +// log(a) + log(c) makes log(c) NaN where log(a*c) is a finite real for a < 0. +func.func @main_neg_mul(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<-4.000000e+00> : tensor + %0 = stablehlo.multiply %arg0, %cst : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_neg_mul(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// log(a / c) with a negative constant must stay as-is for the same reason. +func.func @main_neg_div(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<-4.000000e+00> : tensor + %0 = stablehlo.divide %arg0, %cst : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_neg_div(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.divide %arg0, %cst : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// log(pow(x, y)) must stay as-is when x is not provably non-negative: for x < 0, +// log(pow(x, y)) can be a finite real while y * log(x) is NaN. +func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor { + %0 = stablehlo.power %arg0, %arg1 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor { +// CHECK-NEXT: %0 = stablehlo.power %arg0, %arg1 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } From 4f2c333ab933e91b5fd30281539823a033cb621d Mon Sep 17 00:00:00 2001 From: June Kim Date: Mon, 22 Jun 2026 12:36:22 -0700 Subject: [PATCH 02/14] Guard LogSimplify rewrites against domain narrowing (#2570) Four LogSimplify sub-cases applied real-analysis identities valid only for positive bases, with no domain guard, so the optimized program returned NaN where the original computed a finite real: log(a*a) -> 2*log(a) NaN for any a < 0 log(pow(x,y)) -> y*log(x) NaN for x < 0 (even y), or y = 0 log(a*b) -> log(a)+log(b) NaN when the constant operand < 0 log(a/b) -> log(a)-log(b) NaN when the constant operand < 0 These are not a fast-math/precision question: as partial functions the two sides have different domains (log(a*a) is real for a != 0, 2*log(a) only for a > 0), so the divergence appears before floating point is involved. - log(a*a) -> 2*log(abs(a)): unconditionally domain-sound (a*a and abs(a) share a domain), so no guard needed; just take log of |a|. - log(pow(x,y)): fire only when x is provably non-negative. - log(a*b) / log(a/b): fire only when the constant operand is non-negative (a positive constant agrees, NaN-for-NaN, on negative inputs; a negative constant manufactures the NaN). Reuses the existing guaranteedNonNegativeResult analysis. Updates the two log(a*a) cases in logsimplify.mlir to the abs form and adds negative-constant / unknown-sign-base cases asserting the guarded rewrites no longer fire. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 44 ++++++++++++---- test/lit_tests/logsimplify.mlir | 64 ++++++++++++++++++++--- 2 files changed, 90 insertions(+), 18 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 25d2c1cbb1..cf62e1b9c3 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27760,7 +27760,11 @@ struct LogSimplify final { // log(pow(x, y)) -> y * log(x) auto defOp = op.getOperand().getDefiningOp(); - if (defOp) { + // Only sound when the base x is provably non-negative. For x < 0, + // log(pow(x, y)) can be a finite real (e.g. even integer y, where + // pow(x, y) > 0) while y * log(x) is NaN; likewise y = 0 gives + // log(pow(x, 0)) = 0 but 0 * log(x) = NaN. See issue #2570. + if (defOp && guaranteedNonNegativeResult(defOp.getLhs(), rewriter)) { rewriter.replaceOpWithNewOp( op, defOp.getRhs(), stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs())); @@ -27773,23 +27777,35 @@ struct LogSimplify final if (defOp) { auto lhs = defOp.getLhs(); auto rhs = defOp.getRhs(); - if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(a) + if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(abs(a)) + // a*a is non-negative for every real a, but 2 * log(a) is NaN for + // a < 0. Take the log of |a| so the rewrite keeps log(a*a)'s domain + // (log(a*a) = 2 * log|a|). See issue #2570. rewriter.replaceOpWithNewOp( op, stablehlo::ConstantOp::create( rewriter, op.getLoc(), lhs.getType(), cast(makeAttr(lhs.getType(), 2))), - stablehlo::LogOp::create(rewriter, op.getLoc(), lhs)); + stablehlo::LogOp::create( + rewriter, op.getLoc(), + stablehlo::AbsOp::create(rewriter, op.getLoc(), + lhs.getType(), lhs))); return success(); } if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(mul(a, b)) -> log(a) + // log(b) if a or b is constant - rewriter.replaceOpWithNewOp( - op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), - stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); - return success(); + // Only sound when the constant operand is non-negative: a negative + // constant makes log(const) NaN where log(a*b) was a finite real + // (a < 0 makes both sides NaN, so they still agree). Issue #2570. + Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; + if (guaranteedNonNegativeResult(cst, rewriter)) { + rewriter.replaceOpWithNewOp( + op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), + stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); + return success(); + } } } } @@ -27839,10 +27855,16 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(div(a, b)) -> log(a) - // log(b) if a or b is constant - rewriter.replaceOpWithNewOp( - op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), - stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); - return success(); + // Only sound when the constant operand is non-negative: a negative + // constant makes log(const) NaN where log(a/b) was a finite real. + // Issue #2570. + Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; + if (guaranteedNonNegativeResult(cst, rewriter)) { + rewriter.replaceOpWithNewOp( + op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), + stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); + return success(); + } } } } diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 215e350d37..2c8db51c3f 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -19,10 +19,11 @@ func.func @main2(%arg0: tensor) -> tensor { // CHECK: func.func @main2(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor -// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor -// CHECK-NEXT: return %2 : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.abs %0 : tensor +// CHECK-NEXT: %2 = stablehlo.log %1 : tensor +// CHECK-NEXT: %3 = stablehlo.multiply %cst, %2 : tensor +// CHECK-NEXT: return %3 : tensor // CHECK-NEXT: } func.func @main3(%arg0: tensor) -> tensor { @@ -74,9 +75,10 @@ func.func @main6(%arg0: tensor) -> tensor { // CHECK: func.func @main6(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor -// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor +// CHECK-NEXT: return %2 : tensor // CHECK-NEXT: } func.func @main7(%arg0: tensor) -> tensor { @@ -145,3 +147,51 @@ func.func @main11(%arg0: tensor) -> tensor { // CHECK-NEXT: %1 = stablehlo.divide %0, %cst : tensor // CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } + +// Domain guards (issue #2570): the split rewrites must NOT fire when they would +// narrow the domain and manufacture a NaN on inputs the original handles. + +// log(a * c) with a negative constant must stay as-is: splitting to +// log(a) + log(c) makes log(c) NaN where log(a*c) is a finite real for a < 0. +func.func @main_neg_mul(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<-4.000000e+00> : tensor + %0 = stablehlo.multiply %arg0, %cst : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_neg_mul(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// log(a / c) with a negative constant must stay as-is for the same reason. +func.func @main_neg_div(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<-4.000000e+00> : tensor + %0 = stablehlo.divide %arg0, %cst : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_neg_div(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.divide %arg0, %cst : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// log(pow(x, y)) must stay as-is when x is not provably non-negative: for x < 0, +// log(pow(x, y)) can be a finite real while y * log(x) is NaN. +func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor { + %0 = stablehlo.power %arg0, %arg1 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor { +// CHECK-NEXT: %0 = stablehlo.power %arg0, %arg1 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } From ba6fd786f5997c4ab1b0af2668a80a7935ffa2b3 Mon Sep 17 00:00:00 2001 From: June Kim Date: Mon, 13 Jul 2026 07:43:56 -0700 Subject: [PATCH 03/14] Fix logsimplify lit test: expect non_negative memo attr; clang-format The domain guard for negative-constant log(mul/div) correctly declines to split, but guaranteedNonNegativeResult stamps enzymexla.non_negative on the inspected constant. Update the main_neg_mul/main_neg_div CHECK lines to match the NOTGUARANTEED annotation (repo convention, cf. abspositive.mlir). Also apply clang-format 16 to the AbsOp::create wrapping. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 4 ++-- test/lit_tests/logsimplify.mlir | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index c28ec3c2e5..8d4340a4ab 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27418,8 +27418,8 @@ struct LogSimplify final cast(makeAttr(lhs.getType(), 2))), stablehlo::LogOp::create( rewriter, op.getLoc(), - stablehlo::AbsOp::create(rewriter, op.getLoc(), - lhs.getType(), lhs))); + stablehlo::AbsOp::create(rewriter, op.getLoc(), lhs.getType(), + lhs))); return success(); } diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 2c8db51c3f..c458da6ce8 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -161,7 +161,7 @@ func.func @main_neg_mul(%arg0: tensor) -> tensor { } // CHECK: func.func @main_neg_mul(%arg0: tensor) -> tensor { -// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla]} dense<-4.000000e+00> : tensor // CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: return %1 : tensor @@ -176,7 +176,7 @@ func.func @main_neg_div(%arg0: tensor) -> tensor { } // CHECK: func.func @main_neg_div(%arg0: tensor) -> tensor { -// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor +// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla]} dense<-4.000000e+00> : tensor // CHECK-NEXT: %0 = stablehlo.divide %arg0, %cst : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: return %1 : tensor From 9cc0a7b9d339daa27b42cc0e13953ee71ec82c65 Mon Sep 17 00:00:00 2001 From: June Kim Date: Thu, 16 Jul 2026 11:39:04 -0700 Subject: [PATCH 04/14] Tighten LogSimplify domain guards --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 55 +++++++++++-------- test/lit_tests/logsimplify.mlir | 65 ++++++++++++++++++++++- 2 files changed, 95 insertions(+), 25 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 4ebdc4ff1d..a6d14e6922 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27743,6 +27743,25 @@ struct CommonAssociativeCommutativeOpReorder final } }; +static bool isNonNegativeConstant(Value value) { + DenseElementsAttr attr; + if (!matchPattern(value, m_Constant(&attr)) || + !isa(attr.getElementType())) + return false; + + return llvm::all_of(attr.getValues(), [](const APFloat &element) { + return !element.isNaN() && !element.isNegative(); + }); +} + +static bool isStrictlyPositiveConstant(Value value) { + DenseElementsAttr attr; + return isNonNegativeConstant(value) && + matchPattern(value, m_Constant(&attr)) && + llvm::none_of(attr.getValues(), + [](const APFloat &element) { return element.isZero(); }); +} + struct LogSimplify final : public CheckedOpRewritePattern { using CheckedOpRewritePattern y * log(x) - auto defOp = op.getOperand().getDefiningOp(); - // Only sound when the base x is provably non-negative. For x < 0, - // log(pow(x, y)) can be a finite real (e.g. even integer y, where - // pow(x, y) > 0) while y * log(x) is NaN; likewise y = 0 gives - // log(pow(x, 0)) = 0 but 0 * log(x) = NaN. See issue #2570. - if (defOp && guaranteedNonNegativeResult(defOp.getLhs(), rewriter)) { - rewriter.replaceOpWithNewOp( - op, defOp.getRhs(), - stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs())); - return success(); - } - } - { auto defOp = op.getOperand().getDefiningOp(); if (defOp) { @@ -27796,11 +27801,11 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(mul(a, b)) -> log(a) + // log(b) if a or b is constant - // Only sound when the constant operand is non-negative: a negative - // constant makes log(const) NaN where log(a*b) was a finite real - // (a < 0 makes both sides NaN, so they still agree). Issue #2570. + // Require a strictly positive constant to avoid narrowing the + // domain for negative or zero constants. This does not make the + // rewrite bit-exact under floating-point rounding. See issue #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (guaranteedNonNegativeResult(cst, rewriter)) { + if (isStrictlyPositiveConstant(cst)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); @@ -27855,11 +27860,15 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(div(a, b)) -> log(a) - // log(b) if a or b is constant - // Only sound when the constant operand is non-negative: a negative - // constant makes log(const) NaN where log(a/b) was a finite real. - // Issue #2570. - Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (guaranteedNonNegativeResult(cst, rewriter)) { + // A constant numerator must be strictly positive; a zero numerator + // can turn a finite log(-0) into NaN after splitting. A zero + // denominator is safe here, so it only needs to be non-negative. + // This does not make the rewrite bit-exact under floating-point + // rounding. See issue #2570. + bool constantIsLhs = matchPattern(lhs, m_Constant()); + Value cst = constantIsLhs ? lhs : rhs; + if ((constantIsLhs && isStrictlyPositiveConstant(cst)) || + (!constantIsLhs && isNonNegativeConstant(cst))) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index c458da6ce8..3101b48e4e 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -10,6 +10,67 @@ func.func @main1(%arg0: tensor) -> tensor { // CHECK-NEXT: return %arg0 : tensor // CHECK-NEXT: } +// A zero constant must not split: for a negative argument, the original is +// log(-0) = -inf while log(0) + log(argument) is NaN. +func.func @main_zero_mul(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<0.000000e+00> : tensor + %0 = stablehlo.multiply %cst, %arg0 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_zero_mul(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %cst, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// A zero numerator must not split for the same signed-zero reason. +func.func @main_zero_div(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<0.000000e+00> : tensor + %0 = stablehlo.divide %cst, %arg0 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_zero_div(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.divide %cst, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +// The power rewrite is disabled: zero bases and infinite exponents can change +// finite results into NaNs (e.g. pow(0, 0) and pow(1, inf)). +func.func @main_pow_zero(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<0.000000e+00> : tensor + %0 = stablehlo.power %cst, %arg0 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_pow_zero(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.power %cst, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + +func.func @main_pow_positive(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<1.000000e+00> : tensor + %0 = stablehlo.power %cst, %arg0 : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_pow_positive(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<1.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.power %cst, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + func.func @main2(%arg0: tensor) -> tensor { %0 = stablehlo.multiply %arg0, %arg0 : tensor %1 = stablehlo.multiply %0, %0 : tensor @@ -161,7 +222,7 @@ func.func @main_neg_mul(%arg0: tensor) -> tensor { } // CHECK: func.func @main_neg_mul(%arg0: tensor) -> tensor { -// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla]} dense<-4.000000e+00> : tensor +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor // CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: return %1 : tensor @@ -176,7 +237,7 @@ func.func @main_neg_div(%arg0: tensor) -> tensor { } // CHECK: func.func @main_neg_div(%arg0: tensor) -> tensor { -// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla]} dense<-4.000000e+00> : tensor +// CHECK-NEXT: %cst = stablehlo.constant dense<-4.000000e+00> : tensor // CHECK-NEXT: %0 = stablehlo.divide %arg0, %cst : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: return %1 : tensor From e57a9dafd6073c9f64f3dcaa4f308f9131307683 Mon Sep 17 00:00:00 2001 From: June Kim Date: Thu, 16 Jul 2026 11:43:05 -0700 Subject: [PATCH 05/14] Keep LogSimplify predicate local --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 38 ++++++++++------------- 1 file changed, 16 insertions(+), 22 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index a6d14e6922..8a3a29d496 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27743,25 +27743,6 @@ struct CommonAssociativeCommutativeOpReorder final } }; -static bool isNonNegativeConstant(Value value) { - DenseElementsAttr attr; - if (!matchPattern(value, m_Constant(&attr)) || - !isa(attr.getElementType())) - return false; - - return llvm::all_of(attr.getValues(), [](const APFloat &element) { - return !element.isNaN() && !element.isNegative(); - }); -} - -static bool isStrictlyPositiveConstant(Value value) { - DenseElementsAttr attr; - return isNonNegativeConstant(value) && - matchPattern(value, m_Constant(&attr)) && - llvm::none_of(attr.getValues(), - [](const APFloat &element) { return element.isZero(); }); -} - struct LogSimplify final : public CheckedOpRewritePattern { using CheckedOpRewritePattern(attr.getElementType())) + return false; + + return llvm::all_of(attr.getValues(), + [allowZero](const APFloat &element) { + return !element.isNaN() && + !element.isNegative() && + (allowZero || !element.isZero()); + }); + }; + { // log(exp(x)) -> x auto defOp = op.getOperand().getDefiningOp(); if (defOp) { @@ -27805,7 +27800,7 @@ struct LogSimplify final // domain for negative or zero constants. This does not make the // rewrite bit-exact under floating-point rounding. See issue #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (isStrictlyPositiveConstant(cst)) { + if (isValidLogSplitConstant(cst, /*allowZero=*/false)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); @@ -27867,8 +27862,7 @@ struct LogSimplify final // rounding. See issue #2570. bool constantIsLhs = matchPattern(lhs, m_Constant()); Value cst = constantIsLhs ? lhs : rhs; - if ((constantIsLhs && isStrictlyPositiveConstant(cst)) || - (!constantIsLhs && isNonNegativeConstant(cst))) { + if (isValidLogSplitConstant(cst, /*allowZero=*/!constantIsLhs)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); From 9bab829406af0587ee62f250198fa60c883d45cd Mon Sep 17 00:00:00 2001 From: June Kim Date: Tue, 21 Jul 2026 17:50:07 -0700 Subject: [PATCH 06/14] clang-format: collapse LogSimplify predicate line --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 8a3a29d496..9a43d7857a 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27758,8 +27758,7 @@ struct LogSimplify final return llvm::all_of(attr.getValues(), [allowZero](const APFloat &element) { - return !element.isNaN() && - !element.isNegative() && + return !element.isNaN() && !element.isNegative() && (allowZero || !element.isZero()); }); }; From 6b6f1e5f4bd706052399d5d7b448c5a626b42aa9 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 11:13:32 -0700 Subject: [PATCH 07/14] LogSimplify: gate log(a*a) rewrite on non-negativity instead of abs Per review, replace the abs-based log(mul(a,a)) -> 2*log(abs(a)) rewrite with a guard on guaranteedNonNegativeResult, emitting 2*log(a) only when a is provably non-negative. Avoids inserting an abs (potentially slower than the mul) while staying correct for a < 0. Updates lit tests and adds positive/negative coverage. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 15 +++++----- test/lit_tests/logsimplify.mlir | 34 +++++++++++++++++------ 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 9a43d7857a..47d9e7f7e0 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27776,19 +27776,18 @@ struct LogSimplify final if (defOp) { auto lhs = defOp.getLhs(); auto rhs = defOp.getRhs(); - if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(abs(a)) - // a*a is non-negative for every real a, but 2 * log(a) is NaN for - // a < 0. Take the log of |a| so the rewrite keeps log(a*a)'s domain - // (log(a*a) = 2 * log|a|). See issue #2570. + // log(mul(a, a)) -> 2 * log(a) when a is provably non-negative. + // For a < 0, 2 * log(a) is NaN even though log(a*a) is finite, so we + // only fire this when a is provably non-negative. Gating on the + // non-negative analysis keeps the rewrite correct without inserting an + // abs (which may be slower than the mul). See issue #2570. + if (lhs == rhs && guaranteedNonNegativeResult(lhs, rewriter)) { rewriter.replaceOpWithNewOp( op, stablehlo::ConstantOp::create( rewriter, op.getLoc(), lhs.getType(), cast(makeAttr(lhs.getType(), 2))), - stablehlo::LogOp::create( - rewriter, op.getLoc(), - stablehlo::AbsOp::create(rewriter, op.getLoc(), lhs.getType(), - lhs))); + stablehlo::LogOp::create(rewriter, op.getLoc(), lhs)); return success(); } diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 3101b48e4e..6add6041db 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -78,13 +78,14 @@ func.func @main2(%arg0: tensor) -> tensor { return %2 : tensor } +// arg0 * arg0 is provably non-negative, so log(mul(sq, sq)) folds to +// 2 * log(sq) without inserting an abs. // CHECK: func.func @main2(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor // CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.abs %0 : tensor -// CHECK-NEXT: %2 = stablehlo.log %1 : tensor -// CHECK-NEXT: %3 = stablehlo.multiply %cst, %2 : tensor -// CHECK-NEXT: return %3 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor +// CHECK-NEXT: return %2 : tensor // CHECK-NEXT: } func.func @main3(%arg0: tensor) -> tensor { @@ -134,12 +135,12 @@ func.func @main6(%arg0: tensor) -> tensor { return %1 : tensor } +// arg0 is not provably non-negative, so log(arg0 * arg0) is left untouched +// rather than folded to a NaN-producing 2 * log(arg0). // CHECK: func.func @main6(%arg0: tensor) -> tensor { -// CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor -// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor -// CHECK-NEXT: return %2 : tensor +// CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } func.func @main7(%arg0: tensor) -> tensor { @@ -256,3 +257,20 @@ func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor // CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } + +// abs(arg0) is provably non-negative, so log(mul(a, a)) folds to 2 * log(a) +// directly, with no extra abs inserted. +func.func @main_square_nonneg(%arg0: tensor) -> tensor { + %a = stablehlo.abs %arg0 : tensor + %0 = stablehlo.multiply %a, %a : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_square_nonneg(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor +// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor +// CHECK-NEXT: return %2 : tensor +// CHECK-NEXT: } From aa065977442e5ef5461bd9d3888d457b26a4c914 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 11:44:35 -0700 Subject: [PATCH 08/14] Rename isValidLogSplitConstant to isPositiveFiniteConstant Per review, clearer name for the log-split constant guard. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 47d9e7f7e0..ccfd811319 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27750,7 +27750,7 @@ struct LogSimplify final LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, PatternRewriter &rewriter) const { - auto isValidLogSplitConstant = [](Value value, bool allowZero) { + auto isPositiveFiniteConstant = [](Value value, bool allowZero) { DenseElementsAttr attr; if (!matchPattern(value, m_Constant(&attr)) || !isa(attr.getElementType())) @@ -27798,7 +27798,7 @@ struct LogSimplify final // domain for negative or zero constants. This does not make the // rewrite bit-exact under floating-point rounding. See issue #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (isValidLogSplitConstant(cst, /*allowZero=*/false)) { + if (isPositiveFiniteConstant(cst, /*allowZero=*/false)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); @@ -27860,7 +27860,7 @@ struct LogSimplify final // rounding. See issue #2570. bool constantIsLhs = matchPattern(lhs, m_Constant()); Value cst = constantIsLhs ? lhs : rhs; - if (isValidLogSplitConstant(cst, /*allowZero=*/!constantIsLhs)) { + if (isPositiveFiniteConstant(cst, /*allowZero=*/!constantIsLhs)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); From b76081e1f69b72565c6b4446e75e54dd553cf9b2 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 11:53:09 -0700 Subject: [PATCH 09/14] Rename log-split constant guard to isPositiveNonNanConstant Per review, the predicate accepts +inf, so isPositiveNonNanConstant names it more accurately than isPositiveFiniteConstant. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index ccfd811319..c48d735990 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27750,7 +27750,7 @@ struct LogSimplify final LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, PatternRewriter &rewriter) const { - auto isPositiveFiniteConstant = [](Value value, bool allowZero) { + auto isPositiveNonNanConstant = [](Value value, bool allowZero) { DenseElementsAttr attr; if (!matchPattern(value, m_Constant(&attr)) || !isa(attr.getElementType())) @@ -27798,7 +27798,7 @@ struct LogSimplify final // domain for negative or zero constants. This does not make the // rewrite bit-exact under floating-point rounding. See issue #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (isPositiveFiniteConstant(cst, /*allowZero=*/false)) { + if (isPositiveNonNanConstant(cst, /*allowZero=*/false)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); @@ -27860,7 +27860,7 @@ struct LogSimplify final // rounding. See issue #2570. bool constantIsLhs = matchPattern(lhs, m_Constant()); Value cst = constantIsLhs ? lhs : rhs; - if (isPositiveFiniteConstant(cst, /*allowZero=*/!constantIsLhs)) { + if (isPositiveNonNanConstant(cst, /*allowZero=*/!constantIsLhs)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); From 51ec6b63184fa3d09b74380b4569ff662086c6c8 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 12:17:08 -0700 Subject: [PATCH 10/14] Make isPositiveNonNanConstant a static member of LogSimplify Per review, hoist the log-split constant predicate out of matchAndRewriteImpl into a static member function for easier debugging and faster compile. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 30 +++++++++++++---------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index c48d735990..9df8ba543b 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27748,21 +27748,25 @@ struct LogSimplify final using CheckedOpRewritePattern::CheckedOpRewritePattern; - LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, - PatternRewriter &rewriter) const { - auto isPositiveNonNanConstant = [](Value value, bool allowZero) { - DenseElementsAttr attr; - if (!matchPattern(value, m_Constant(&attr)) || - !isa(attr.getElementType())) - return false; + // True iff `value` is a float constant whose every element is non-NaN and + // non-negative (and, unless allowZero, strictly positive). Keeps the + // log-splitting rewrites below from narrowing the domain on negative or zero + // constants. + static bool isPositiveNonNanConstant(Value value, bool allowZero) { + DenseElementsAttr attr; + if (!matchPattern(value, m_Constant(&attr)) || + !isa(attr.getElementType())) + return false; - return llvm::all_of(attr.getValues(), - [allowZero](const APFloat &element) { - return !element.isNaN() && !element.isNegative() && - (allowZero || !element.isZero()); - }); - }; + return llvm::all_of(attr.getValues(), + [allowZero](const APFloat &element) { + return !element.isNaN() && !element.isNegative() && + (allowZero || !element.isZero()); + }); + } + LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, + PatternRewriter &rewriter) const { { // log(exp(x)) -> x auto defOp = op.getOperand().getDefiningOp(); if (defOp) { From ff1c40f2ed430f1d9d0298184f7cfe5edeb7ffc2 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 12:21:55 -0700 Subject: [PATCH 11/14] Trim LogSimplify guard comments Tighten the domain-guard comments toward the file's terse style: drop the duplicated 'provably non-negative' phrasing and the repeated bit-exactness caveat, keep the non-obvious reasoning. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 9df8ba543b..aa1e227cf2 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27780,11 +27780,9 @@ struct LogSimplify final if (defOp) { auto lhs = defOp.getLhs(); auto rhs = defOp.getRhs(); - // log(mul(a, a)) -> 2 * log(a) when a is provably non-negative. - // For a < 0, 2 * log(a) is NaN even though log(a*a) is finite, so we - // only fire this when a is provably non-negative. Gating on the - // non-negative analysis keeps the rewrite correct without inserting an - // abs (which may be slower than the mul). See issue #2570. + // log(mul(a, a)) -> 2 * log(a). Only when a >= 0: for a < 0, 2*log(a) + // is NaN while log(a*a) is finite. Gating on the analysis avoids an abs + // (possibly slower than the mul). See #2570. if (lhs == rhs && guaranteedNonNegativeResult(lhs, rewriter)) { rewriter.replaceOpWithNewOp( op, @@ -27798,9 +27796,9 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(mul(a, b)) -> log(a) + // log(b) if a or b is constant - // Require a strictly positive constant to avoid narrowing the - // domain for negative or zero constants. This does not make the - // rewrite bit-exact under floating-point rounding. See issue #2570. + // Split only for a strictly positive constant; negative/zero narrows + // the domain, and the split isn't bit-exact under FP rounding. See + // #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; if (isPositiveNonNanConstant(cst, /*allowZero=*/false)) { rewriter.replaceOpWithNewOp( @@ -27857,11 +27855,9 @@ struct LogSimplify final if (anyOperandIsConstant(defOp) && !allOperandsAreConstant(defOp)) { // log(div(a, b)) -> log(a) - // log(b) if a or b is constant - // A constant numerator must be strictly positive; a zero numerator - // can turn a finite log(-0) into NaN after splitting. A zero - // denominator is safe here, so it only needs to be non-negative. - // This does not make the rewrite bit-exact under floating-point - // rounding. See issue #2570. + // Constant numerator must be strictly positive: a zero numerator + // turns a finite log(-0) into NaN after splitting. A zero denominator + // is safe, so it only needs to be non-negative. See #2570. bool constantIsLhs = matchPattern(lhs, m_Constant()); Value cst = constantIsLhs ? lhs : rhs; if (isPositiveNonNanConstant(cst, /*allowZero=*/!constantIsLhs)) { From 48aaf2b1f94b217f16a219441b5f50278063dbf4 Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 22:11:36 -0700 Subject: [PATCH 12/14] Update main2 CHECK for non_negative annotation on squared multiply The guaranteedNonNegativeResult gate stamps enzymexla.non_negative on the inner arg0*arg0 when the log(mul(sq,sq)) rewrite queries it, so the CHECK-NEXT must match the annotated op. main6 is unaffected (arg0 itself is queried, not the multiply). --- test/lit_tests/logsimplify.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 6add6041db..60b0e770dc 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -82,7 +82,7 @@ func.func @main2(%arg0: tensor) -> tensor { // 2 * log(sq) without inserting an abs. // CHECK: func.func @main2(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 {enzymexla.non_negative = [#enzymexla]} : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor // CHECK-NEXT: return %2 : tensor From 5660c46ca7b81e549f91c43a7811154b5ec802a6 Mon Sep 17 00:00:00 2001 From: June Kim Date: Thu, 23 Jul 2026 07:51:39 -0700 Subject: [PATCH 13/14] LogSimplify: exclude infinite constants from log-splits isPositiveNonNanConstant passed +inf (only NaN and negatives were rejected), so log(x * inf) / log(x / inf) could split and flip a finite/NaN result to the other. Gate on isFinite() and rename to isPositiveFiniteConstant (now accurate). Adds an inf-mul guard test. Found by codex review of #2576. --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 16 ++++++++-------- test/lit_tests/logsimplify.mlir | 16 ++++++++++++++++ 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index aa1e227cf2..adee471fdb 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27748,11 +27748,11 @@ struct LogSimplify final using CheckedOpRewritePattern::CheckedOpRewritePattern; - // True iff `value` is a float constant whose every element is non-NaN and - // non-negative (and, unless allowZero, strictly positive). Keeps the - // log-splitting rewrites below from narrowing the domain on negative or zero - // constants. - static bool isPositiveNonNanConstant(Value value, bool allowZero) { + // True iff `value` is a finite float constant whose every element is + // non-negative (and, unless allowZero, strictly positive). Excludes NaN, + // negatives, and infinities, which would let the log-splits below narrow the + // domain. See #2570. + static bool isPositiveFiniteConstant(Value value, bool allowZero) { DenseElementsAttr attr; if (!matchPattern(value, m_Constant(&attr)) || !isa(attr.getElementType())) @@ -27760,7 +27760,7 @@ struct LogSimplify final return llvm::all_of(attr.getValues(), [allowZero](const APFloat &element) { - return !element.isNaN() && !element.isNegative() && + return element.isFinite() && !element.isNegative() && (allowZero || !element.isZero()); }); } @@ -27800,7 +27800,7 @@ struct LogSimplify final // the domain, and the split isn't bit-exact under FP rounding. See // #2570. Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs; - if (isPositiveNonNanConstant(cst, /*allowZero=*/false)) { + if (isPositiveFiniteConstant(cst, /*allowZero=*/false)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); @@ -27860,7 +27860,7 @@ struct LogSimplify final // is safe, so it only needs to be non-negative. See #2570. bool constantIsLhs = matchPattern(lhs, m_Constant()); Value cst = constantIsLhs ? lhs : rhs; - if (isPositiveNonNanConstant(cst, /*allowZero=*/!constantIsLhs)) { + if (isPositiveFiniteConstant(cst, /*allowZero=*/!constantIsLhs)) { rewriter.replaceOpWithNewOp( op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 60b0e770dc..49f0de63c9 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -244,6 +244,22 @@ func.func @main_neg_div(%arg0: tensor) -> tensor { // CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } +// log(a * c) with an infinite constant must stay as-is: log(x * inf) is a +// finite -inf/NaN that splitting to log(x) + log(inf) would not preserve. +func.func @main_inf_mul(%arg0: tensor) -> tensor { + %cst = stablehlo.constant dense<0x7FF0000000000000> : tensor + %0 = stablehlo.multiply %arg0, %cst : tensor + %1 = stablehlo.log %0 : tensor + return %1 : tensor +} + +// CHECK: func.func @main_inf_mul(%arg0: tensor) -> tensor { +// CHECK-NEXT: %cst = stablehlo.constant dense<0x7FF0000000000000> : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor +// CHECK-NEXT: return %1 : tensor +// CHECK-NEXT: } + // log(pow(x, y)) must stay as-is when x is not provably non-negative: for x < 0, // log(pow(x, y)) can be a finite real while y * log(x) is NaN. func.func @main_pow_unknown(%arg0: tensor, %arg1: tensor) -> tensor { From b30b8d7c61a5869db0170f18a145180549aee028 Mon Sep 17 00:00:00 2001 From: June Kim Date: Thu, 23 Jul 2026 10:21:28 -0700 Subject: [PATCH 14/14] Expect non_negative stamp on abs in main_square_nonneg CHECK log(mul(abs,abs)) queries guaranteedNonNegativeResult(abs), which stamps enzymexla.non_negative on the abs op; the CHECK expected a bare abs. Same stamping pattern as main2. This was the last unreached CHECK, so it never surfaced until the earlier lanes went green. --- test/lit_tests/logsimplify.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/lit_tests/logsimplify.mlir b/test/lit_tests/logsimplify.mlir index 49f0de63c9..5446fe8405 100644 --- a/test/lit_tests/logsimplify.mlir +++ b/test/lit_tests/logsimplify.mlir @@ -285,7 +285,7 @@ func.func @main_square_nonneg(%arg0: tensor) -> tensor { // CHECK: func.func @main_square_nonneg(%arg0: tensor) -> tensor { // CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor -// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor +// CHECK-NEXT: %0 = stablehlo.abs %arg0 {enzymexla.non_negative = [#enzymexla]} : tensor // CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor // CHECK-NEXT: return %2 : tensor