Guard LogSimplify rewrites against domain narrowing (#2570)#2576
Conversation
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.
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.
86eef28 to
4f2c333
Compare
|
//test/lit_tests:logsimplify.mlir.test is failing in CI |
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.
# Conflicts: # src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp # test/lit_tests/logsimplify.mlir
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2576 +/- ##
==========================================
+ Coverage 31.07% 31.70% +0.63%
==========================================
Files 225 226 +1
Lines 44509 44892 +383
==========================================
+ Hits 13830 14234 +404
+ Misses 30679 30658 -21 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Lit test fixed in |
Antipath1
left a comment
Antipath1
left a comment
There was a problem hiding this comment.
I think this improves the status quo quite a lot thank you so I think that this should actually get merged even in this state but I think the comments should be updated to reflect the remaining problems.
| // 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)) { |
There was a problem hiding this comment.
The comment here claims that this would fix the log(pow(x, 0)) case, but it actually does not. In the specific case where
our pattern rewrites it into
which evaluates to NaN (because
(cc @wsmoses I think this needs your authority on how to deal with this)
| // 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<stablehlo::AddOp>( | ||
| op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), | ||
| stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); | ||
| return success(); | ||
| } |
There was a problem hiding this comment.
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:
Optimized:
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.
| // 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(); | ||
| } |
There was a problem hiding this comment.
Same problem and still do not think this should be blocking
Fails when the numerator (
Original:
Optimized:
|
Good catches, thanks. I tightened the patch:
But even with positive constants, the multiply/divide identities aren’t fully IEEE-equivalent because rounding, overflow, and underflow can differ. It may be out of scope. If required, those rewrites may need separate gating or removal. |
|
@Antipath1 and I spoke about this during the meeting today. @kimjune01 rather than folding to an abs (which may be slower than the mul), what if you gated the transformation on the result of the is value positive analysis |
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.
|
Done -- dropped the While wiring this up I hit a latent soundness bug in that analysis's square rule -- it compares defining ops, so |
| // 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)) { |
There was a problem hiding this comment.
can you rename this from isValidLogSplitConstant to isPositiveFiniteConstant for clarity
wsmoses
left a comment
wsmoses
left a comment
There was a problem hiding this comment.
lgtm modulo the fn name for simplicity
Per review, clearer name for the log-split constant guard.
|
Renamed to One nuance on the name so it doesn't mislead later: the predicate currently accepts |
|
maybe the isPositiveNonNanConstant then? |
Per review, the predicate accepts +inf, so isPositiveNonNanConstant names it more accurately than isPositiveFiniteConstant.
|
Good call -- |
|
|
||
| LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, | ||
| PatternRewriter &rewriter) const { | ||
| auto isPositiveNonNanConstant = [](Value value, bool allowZero) { |
There was a problem hiding this comment.
sorry one final minor comment, can you make this a static member function rather than a lambda?
it'll be easier to debug in the future [and faster to compile/run]
Per review, hoist the log-split constant predicate out of matchAndRewriteImpl into a static member function for easier debugging and faster compile.
|
Done in |
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.
|
ci fails: |
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).
|
Fixed in |
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 EnzymeAD#2576.
|
Also excluded infinite constants from the splits. |
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.
3 participants
Fixes #2570 (maintainer-approved: "makes sense, PR welcome!").
Problem
Four
LogSimplifysub-cases (EnzymeHLOOpt.cpp) apply a real-analysis identity valid only for positive bases, with no domain guard, so the optimized program returns NaN where the unoptimized one returns a finite real:log(a*a) -> 2*log(a)a < 0log(4)=1.386becomes2*log(-2)=NaNlog(pow(x,y)) -> y*log(x)x < 0(even y), ory = 0log(1)=0becomes0*log(-2)=NaNlog(a*b) -> log(a)+log(b)(one const)< 0log(6)=1.79becomesNaNlog(a/b) -> log(a)-log(b)(one const)< 0log(2/3)=-0.405becomesNaNThis is not the same as the opt-in float policy behind
log(sqrt(x)) -> 0.5*log(x). That rewrite is an exact identity in the reals and only perturbs float rounding. These four are false in the reals themselves: as partial functions the two sides have different domains (log(a*a)is real fora != 0,2*log(a)only fora > 0), so the divergence appears before floating point is involved. Same finite-to-NaN domain narrowing as theCbrtOpderivative in #2571 / #2575.Fix
Reuse the existing
guaranteedNonNegativeResultanalysis to gate each case on a provably-non-negative base:log(a*a) -> 2*log(abs(a)): unconditionally domain-sound, sincea*aandabs(a)share a domain. No guard needed; just take the log of|a|.log(pow(x,y)): fire only whenxis 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.Test
test/lit_tests/logsimplify.mlir's@main2/@main6(thelog(a*a)cases) are updated to the newabsform, and three cases are added asserting the guarded rewrites no longer fire (negative-constant mul, negative-constant div, unknown-signpowbase). The pipeline enables onlylog_simplify;log_const_prop, so the positive-constant cases (@main3/@main4/@main7/@main8) still fire unchanged.Found mechanically with a value-and-gradient soundness gate over the rewrite library; happy to share the process if useful.