Skip to content

Guard LogSimplify rewrites against domain narrowing (#2570) #2576

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wsmoses merged 15 commits into from
Jul 23, 2026
Merged

Guard LogSimplify rewrites against domain narrowing (#2570) #2576

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
86eef28
Guard LogSimplify rewrites against domain narrowing (#2570)
kimjune01 Jun 22, 2026
4f2c333
Guard LogSimplify rewrites against domain narrowing (#2570)
kimjune01 Jun 22, 2026
ba6fd78
Fix logsimplify lit test: expect non_negative memo attr; clang-format
kimjune01 Jul 13, 2026
ab6cee3
Merge fork branch before pushing lit/format fix
kimjune01 Jul 13, 2026
9cc0a7b
Tighten LogSimplify domain guards
kimjune01 Jul 16, 2026
e57a9da
Keep LogSimplify predicate local
kimjune01 Jul 16, 2026
9bab829
clang-format: collapse LogSimplify predicate line
kimjune01 Jul 22, 2026
6b6f1e5
LogSimplify: gate log(a*a) rewrite on non-negativity instead of abs
kimjune01 Jul 22, 2026
aa06597
Rename isValidLogSplitConstant to isPositiveFiniteConstant
kimjune01 Jul 22, 2026
b76081e
Rename log-split constant guard to isPositiveNonNanConstant
kimjune01 Jul 22, 2026
51ec6b6
Make isPositiveNonNanConstant a static member of LogSimplify
kimjune01 Jul 22, 2026
ff1c40f
Trim LogSimplify guard comments
kimjune01 Jul 22, 2026
48aaf2b
Update main2 CHECK for non_negative annotation on squared multiply
kimjune01 Jul 23, 2026
5660c46
LogSimplify: exclude infinite constants from log-splits
kimjune01 Jul 23, 2026
b30b8d7
Expect non_negative stamp on abs in main_square_nonneg CHECK
kimjune01 Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 33 additions & 11 deletions src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27760,7 +27760,11 @@ struct LogSimplify final

{ // log(pow(x, y)) -> y * log(x)
auto defOp = op.getOperand().getDefiningOp<stablehlo::PowOp>();
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)) {

@Antipath1 Antipath1 Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment here claims that this would fix the log(pow(x, 0)) case, but it actually does not. In the specific case where $x = 0$ and $y = 0$:If we start with
$\log(\text{pow}(0, 0)) = \log(1) = 0$,
our pattern rewrites it into $0 \cdot \log(0)$,
which evaluates to NaN (because $\log(0)$ is $-\infty$). That is mathematically incorrect. Soooooooooooooo we still have a correctness bug here, actually. Proving that $x$ or $y$ is never $0$ is kind of hard, so we might genuinely need to put that behind a flag, or just outright remove this rewrite, or try to prove that $x$ is strictly greater than zero (which, I might be wrong, is only provable in very specific scenarios, like the definingOp of $x$ or $y$ being an exponential op which can never be $0$). Nevertheless, I don't think you need to fix this in this PR, since having guaranteedNonNegativeResult is already a massive improvement. So, I think I just need you to update the comment to include this info.

(cc @wsmoses I think this needs your authority on how to deal with this)

rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
op, defOp.getRhs(),
stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs()));
Expand All @@ -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<stablehlo::MulOp>(
op,
stablehlo::ConstantOp::create(
rewriter, op.getLoc(), lhs.getType(),
cast<ElementsAttr>(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<stablehlo::AddOp>(
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.

@wsmoses wsmoses Jul 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you rename this from isValidLogSplitConstant to isPositiveFiniteConstant for clarity

Value cst = matchPattern(lhs, m_Constant()) ? lhs : rhs;
if (guaranteedNonNegativeResult(cst, rewriter)) {
rewriter.replaceOpWithNewOp<stablehlo::AddOp>(
op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs),
stablehlo::LogOp::create(rewriter, op.getLoc(), rhs));
return success();
}

@Antipath1 Antipath1 Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fails when the constant is exactly 0 and the variable is negative. Because guaranteedNonNegativeResult allows the constant to be exactly 0
if cst = 0 and the variable is -2:
Original: $\log(0 \cdot -2) \rightarrow \log(-0.0) \rightarrow \mathbf{-\infty}$
Optimized: $\log(0) + \log(-2) \rightarrow -\infty + \text{NaN} \rightarrow \mathbf{\text{NaN}}$

Similar problem as before, and I still think this should not necessarily block the PR since the status quo is actually a good bit worse.

}
}
}
Expand Down Expand Up @@ -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<stablehlo::SubtractOp>(
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<stablehlo::SubtractOp>(
op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs),
stablehlo::LogOp::create(rewriter, op.getLoc(), rhs));
return success();
}

@Antipath1 Antipath1 Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same problem and still do not think this should be blocking

Fails when the numerator ($a$) is the constant 0, and the denominator ($b$) is negative.The code allows the constant to be either the LHS ($a$) or the RHS ($b$). If the constant happens to be the LHS, and it equals 0, we hit the same trap.
$a = 0$ and $b = -2$:
Original: $\log(0 / -2) \rightarrow \log(-0.0) \rightarrow \mathbf{-\infty}$
Optimized: $\log(0) - \log(-2) \rightarrow -\infty - \text{NaN} \rightarrow \mathbf{\text{NaN}}$(Note: If the denominator $b$ is the constant 0, and $a$ is negative, both sides correctly evaluate to NaN. The bug only exists when the numerator is 0).

}
}
}
Expand Down
64 changes: 57 additions & 7 deletions test/lit_tests/logsimplify.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,11 @@ func.func @main2(%arg0: tensor<f64>) -> tensor<f64> {

// CHECK: func.func @main2(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor<f64>
// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor<f64>
// CHECK-NEXT: return %2 : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.abs %0 : tensor<f64>
// CHECK-NEXT: %2 = stablehlo.log %1 : tensor<f64>
// CHECK-NEXT: %3 = stablehlo.multiply %cst, %2 : tensor<f64>
// CHECK-NEXT: return %3 : tensor<f64>
// CHECK-NEXT: }

func.func @main3(%arg0: tensor<f64>) -> tensor<f64> {
Expand Down Expand Up @@ -74,9 +75,10 @@ func.func @main6(%arg0: tensor<f64>) -> tensor<f64> {

// CHECK: func.func @main6(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.log %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.abs %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: %2 = stablehlo.multiply %cst, %1 : tensor<f64>
// CHECK-NEXT: return %2 : tensor<f64>
// CHECK-NEXT: }

func.func @main7(%arg0: tensor<f64>) -> tensor<f64> {
Expand Down Expand Up @@ -145,3 +147,51 @@ func.func @main11(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %1 = stablehlo.divide %0, %cst : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// 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<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<-4.000000e+00> : tensor<f64>
%0 = stablehlo.multiply %arg0, %cst : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_neg_mul(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} dense<-4.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %cst : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: }

// log(a / c) with a negative constant must stay as-is for the same reason.
func.func @main_neg_div(%arg0: tensor<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<-4.000000e+00> : tensor<f64>
%0 = stablehlo.divide %arg0, %cst : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_neg_div(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant {enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} dense<-4.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.divide %arg0, %cst : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// 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<f64>, %arg1: tensor<f64>) -> tensor<f64> {
%0 = stablehlo.power %arg0, %arg1 : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_pow_unknown(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %0 = stablehlo.power %arg0, %arg1 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: }
Loading