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 all 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
61 changes: 42 additions & 19 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 @@ -27748,6 +27748,23 @@ struct LogSimplify final
using CheckedOpRewritePattern<stablehlo::LogOp,
LogSimplify>::CheckedOpRewritePattern;

// 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<FloatType>(attr.getElementType()))
return false;

return llvm::all_of(attr.getValues<APFloat>(),
[allowZero](const APFloat &element) {
return element.isFinite() && !element.isNegative() &&
(allowZero || !element.isZero());
});
}

LogicalResult matchAndRewriteImpl(stablehlo::LogOp op,
PatternRewriter &rewriter) const {
{ // log(exp(x)) -> x
Expand All @@ -27758,22 +27775,15 @@ struct LogSimplify final
}
}

{ // log(pow(x, y)) -> y * log(x)
auto defOp = op.getOperand().getDefiningOp<stablehlo::PowOp>();
if (defOp) {
rewriter.replaceOpWithNewOp<stablehlo::MulOp>(
op, defOp.getRhs(),
stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs()));
return success();
}
}

{
auto defOp = op.getOperand().getDefiningOp<stablehlo::MulOp>();
if (defOp) {
auto lhs = defOp.getLhs();
auto rhs = defOp.getRhs();
if (lhs == rhs) { // log(mul(a, a)) -> 2 * log(a)
// 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<stablehlo::MulOp>(
op,
stablehlo::ConstantOp::create(
Expand All @@ -27786,10 +27796,16 @@ struct LogSimplify final
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();
// 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 (isPositiveFiniteConstant(cst, /*allowZero=*/false)) {
rewriter.replaceOpWithNewOp<stablehlo::AddOp>(
op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs),
stablehlo::LogOp::create(rewriter, op.getLoc(), rhs));
return success();
}
}
}
}
Expand Down Expand Up @@ -27839,10 +27855,17 @@ 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();
// 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 (isPositiveFiniteConstant(cst, /*allowZero=*/!constantIsLhs)) {
rewriter.replaceOpWithNewOp<stablehlo::SubtractOp>(
op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs),
stablehlo::LogOp::create(rewriter, op.getLoc(), rhs));
return success();
}
}
}
}
Expand Down
155 changes: 150 additions & 5 deletions test/lit_tests/logsimplify.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,80 @@ func.func @main1(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: return %arg0 : tensor<f64>
// 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<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
%0 = stablehlo.multiply %cst, %arg0 : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

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

// A zero numerator must not split for the same signed-zero reason.
func.func @main_zero_div(%arg0: tensor<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
%0 = stablehlo.divide %cst, %arg0 : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_zero_div(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.divide %cst, %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// 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<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
%0 = stablehlo.power %cst, %arg0 : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_pow_zero(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.power %cst, %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: }

func.func @main_pow_positive(%arg0: tensor<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<f64>
%0 = stablehlo.power %cst, %arg0 : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_pow_positive(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<1.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.power %cst, %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: }

func.func @main2(%arg0: tensor<f64>) -> tensor<f64> {
%0 = stablehlo.multiply %arg0, %arg0 : tensor<f64>
%1 = stablehlo.multiply %0, %0 : tensor<f64>
%2 = stablehlo.log %1 : tensor<f64>
return %2 : tensor<f64>
}

// 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<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: %0 = stablehlo.multiply %arg0, %arg0 {enzymexla.non_negative = [#enzymexla<guaranteed GUARANTEED>]} : 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: }
Expand Down Expand Up @@ -72,10 +135,11 @@ func.func @main6(%arg0: tensor<f64>) -> tensor<f64> {
return %1 : tensor<f64>
}

// 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<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: %0 = stablehlo.multiply %arg0, %arg0 : tensor<f64>
// CHECK-NEXT: %1 = stablehlo.log %0 : tensor<f64>
// CHECK-NEXT: return %1 : tensor<f64>
// CHECK-NEXT: }

Expand Down Expand Up @@ -145,3 +209,84 @@ 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 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 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(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<f64>) -> tensor<f64> {
%cst = stablehlo.constant dense<0x7FF0000000000000> : tensor<f64>
%0 = stablehlo.multiply %arg0, %cst : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_inf_mul(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<0x7FF0000000000000> : 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(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: }

// 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<f64>) -> tensor<f64> {
%a = stablehlo.abs %arg0 : tensor<f64>
%0 = stablehlo.multiply %a, %a : tensor<f64>
%1 = stablehlo.log %0 : tensor<f64>
return %1 : tensor<f64>
}

// CHECK: func.func @main_square_nonneg(%arg0: tensor<f64>) -> tensor<f64> {
// CHECK-NEXT: %cst = stablehlo.constant dense<2.000000e+00> : tensor<f64>
// CHECK-NEXT: %0 = stablehlo.abs %arg0 {enzymexla.non_negative = [#enzymexla<guaranteed GUARANTEED>]} : 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: }
Loading