diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 25d2c1cbb1..adee471fdb 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -27748,6 +27748,23 @@ struct LogSimplify final using CheckedOpRewritePattern::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(attr.getElementType())) + return false; + + return llvm::all_of(attr.getValues(), + [allowZero](const APFloat &element) { + return element.isFinite() && !element.isNegative() && + (allowZero || !element.isZero()); + }); + } + LogicalResult matchAndRewriteImpl(stablehlo::LogOp op, PatternRewriter &rewriter) const { { // log(exp(x)) -> x @@ -27758,22 +27775,15 @@ struct LogSimplify final } } - { // log(pow(x, y)) -> y * log(x) - auto defOp = op.getOperand().getDefiningOp(); - if (defOp) { - rewriter.replaceOpWithNewOp( - op, defOp.getRhs(), - stablehlo::LogOp::create(rewriter, op.getLoc(), defOp.getLhs())); - return success(); - } - } - { auto defOp = op.getOperand().getDefiningOp(); 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( op, stablehlo::ConstantOp::create( @@ -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( - 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( + op, stablehlo::LogOp::create(rewriter, op.getLoc(), lhs), + stablehlo::LogOp::create(rewriter, op.getLoc(), rhs)); + return success(); + } } } } @@ -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( - 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( + 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..5446fe8405 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 @@ -17,10 +78,12 @@ 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.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : 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 // CHECK-NEXT: } @@ -72,10 +135,11 @@ 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.log %arg0 : tensor -// CHECK-NEXT: %1 = stablehlo.multiply %cst, %0 : tensor +// CHECK-NEXT: %0 = stablehlo.multiply %arg0, %arg0 : tensor +// CHECK-NEXT: %1 = stablehlo.log %0 : tensor // CHECK-NEXT: return %1 : tensor // CHECK-NEXT: } @@ -145,3 +209,84 @@ 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(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 { + %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: } + +// 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 {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 +// CHECK-NEXT: }