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 8 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 @@ -27750,6 +27750,19 @@ struct LogSimplify final

LogicalResult matchAndRewriteImpl(stablehlo::LogOp op,
PatternRewriter &rewriter) const {
auto isValidLogSplitConstant = [](Value value, bool allowZero) {

@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.

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]

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.isNaN() && !element.isNegative() &&
(allowZero || !element.isZero());
});
};

{ // log(exp(x)) -> x
auto defOp = op.getOperand().getDefiningOp<stablehlo::ExpOp>();
if (defOp) {
Expand All @@ -27758,22 +27771,17 @@ 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) 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<stablehlo::MulOp>(
op,
stablehlo::ConstantOp::create(
Expand All @@ -27786,10 +27794,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();
// 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 (isValidLogSplitConstant(cst, /*allowZero=*/false)) {

@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

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 +27853,19 @@ 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();
// 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 (isValidLogSplitConstant(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
139 changes: 134 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 : 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,68 @@ 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(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 : 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