From df447fdf10b53e275a5def5ef7044e274ceaa0dd Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 17:10:27 -0700 Subject: [PATCH 1/2] domain-guard triage: debug remark at each domain-fact consumer Add DOMAIN_GUARD_TRUST, an LLVM_DEBUG one-liner emitted at each rewrite that trusts a domain-narrowing analysis fact (non_negative / no_nan). Instruments the five consuming sites on main: ReduceMaxMinMulPositiveScalar, the two CompareSimplify non_negative folds, the CompareSimplify no_nan fold, and AbsPositiveSimplify. The predicate and trusted value are already on the IR via setGuaranteedInIR; the remark adds the only missing bit -- the consuming rewrite's name -- plus a grep anchor ([domain-guard]) and a time-ordered trace, so a later domain-narrowing NaN is one grep from the rewrite that produced it. Zero always-on cost: LLVM_DEBUG compiles out under NDEBUG; no lit test uses -debug-only, so FileCheck is unaffected. LogSimplify's log(a*a)/log(a*b) folds consume no fact on main (that is #2570); they become remark sites once their guard lands -- the shared macro makes that a one-liner. No local build available (heavy LLVM/MLIR/XLA); this is source-reasoned, CI is the verifier. See scope/domain-guard-triage/SCOPING.md (gitignored). --- src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 26 +++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp index 25d2c1cbb1..c21dc68280 100644 --- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp +++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp @@ -73,6 +73,24 @@ #include #define DEBUG_TYPE "enzymehloopt" +// Domain-guard triage (branch scope/domain-guard-triage). +// +// Emit one debug line at each rewrite that trusts a domain-narrowing analysis +// fact (non_negative / no_nan / finite). The predicate and the trusted value are +// already recoverable from the IR (setGuaranteedInIR writes e.g. +// `enzymexla.non_negative = GUARANTEED` onto the value and it survives to the +// pass output). What this adds is the only bit that is NOT on the IR: the name +// of the *consuming rewrite*, plus a time-ordered trace. A later +// domain-narrowing NaN is then one grep ("[domain-guard]") from the rewrite that +// produced it. +// +// Zero always-on cost: LLVM_DEBUG compiles out entirely under NDEBUG and is a +// single DebugFlag branch otherwise. No lit test uses -debug-only, so this never +// affects FileCheck. Enable with `-debug-only=enzymehloopt`. +#define DOMAIN_GUARD_TRUST(REWRITE, PREDICATE, VALUE) \ + LLVM_DEBUG(llvm::dbgs() << "[domain-guard] " REWRITE " trusted " PREDICATE \ + << " on " << (VALUE) << "\n") + namespace mlir { namespace enzyme { #define GEN_PASS_DEF_ENZYMEHLOOPTPASS @@ -10541,6 +10559,7 @@ struct ReduceMaxMinMulPositiveScalar // For now, we accept non-negative which includes zero // This is still correct for max/min but could potentially // multiply by zero which doesn't break correctness + DOMAIN_GUARD_TRUST("ReduceMaxMinMulPositiveScalar", "non_negative", v); return true; } return false; @@ -13678,6 +13697,10 @@ struct CompareOpCanon final if ((compType && *compType == ComparisonType::SIGNED) || (isFloat && (!compType || *compType == ComparisonType::FLOAT) && guaranteedNoNanResult(otherOperand, rewriter))) { + // isFloat is true only on the no_nan-trusting arm; the SIGNED arm is + // integer (isFloat == false) and consults no analysis fact. + if (isFloat) + DOMAIN_GUARD_TRUST("CompareSimplify", "no_nan", otherOperand); if (isNegative) { switch (direction) { case ComparisonDirection::EQ: @@ -13721,11 +13744,13 @@ struct CompareOpCanon final }; if (rhsAttr && guaranteedNonNegativeResult(lhs, rewriter)) { + DOMAIN_GUARD_TRUST("CompareSimplify", "non_negative", lhs); if (succeeded(simplifyNonNegative(rhsAttr, direction, lhs))) { return success(); } } if (lhsAttr && guaranteedNonNegativeResult(rhs, rewriter)) { + DOMAIN_GUARD_TRUST("CompareSimplify", "non_negative", rhs); if (succeeded( simplifyNonNegative(lhsAttr, invertDirection(direction), rhs))) { return success(); @@ -22239,6 +22264,7 @@ struct AbsPositiveSimplify return failure(); if (guaranteedNonNegativeResult(operand.getDefiningOp(), rewriter)) { + DOMAIN_GUARD_TRUST("AbsPositiveSimplify", "non_negative", operand); rewriter.replaceOp(op, op.getOperand()); return success(); } From bb3cf38de5c884b7cad0eb8c3ab6103086d086cb Mon Sep 17 00:00:00 2001 From: June Kim Date: Wed, 22 Jul 2026 17:12:49 -0700 Subject: [PATCH 2/2] domain-guard triage: guard regression for CompareOpCanon domain folds The abs (abspositive.mlir test5/test6) and reduce (reduce_max_min_mul_positive_ scalar.mlir reduce_max_mul_no_guarantee) consumers already have negative/ unknown-sign edge coverage. The CompareOpCanon non_negative/no_nan folds did not. Add compare_domain_guard.mlir: an unknown-sign, possibly-NaN operand must not let compare(x, const) collapse to a boolean constant. Guard-only by design -- it asserts a stablehlo.compare survives, not an exact folded form -- so it is robust to unrelated operand/direction canonicalizations and forbids only the fold-to-constant that the #2648/#2651 failure shape would produce. Passes on current main (the guard already holds there); it is a regression tripwire, not a fix. No local build (heavy LLVM/MLIR/XLA); CHECK lines reasoned, not run. CI verifies. --- test/lit_tests/compare_domain_guard.mlir | 37 ++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 test/lit_tests/compare_domain_guard.mlir diff --git a/test/lit_tests/compare_domain_guard.mlir b/test/lit_tests/compare_domain_guard.mlir new file mode 100644 index 0000000000..2a14bc4254 --- /dev/null +++ b/test/lit_tests/compare_domain_guard.mlir @@ -0,0 +1,37 @@ +// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s + +// Domain-guard regression for the CompareOpCanon folds that trust an analysis +// fact (non_negative / no_nan). These lock the *guard*: an operand of unknown +// sign that is not provably no-NaN must NOT let `compare(x, const)` collapse to +// a boolean constant. If NonNegativeResultAnalysis / NoNanResultAnalysis ever +// wrongly proves one of these operands (the #2648 / #2651 failure shape), the +// fold fires and the compare disappears -- these CHECKs then fail. +// +// The assertion is deliberately robust: it only requires that a `stablehlo.compare` +// survive (other canonicalizations may rewrite its operands, which is fine). +// The single thing it forbids is the fold to a constant. +// +// No local build available (heavy LLVM/MLIR/XLA); CHECK lines are reasoned +// against the pass logic, not run. CI is the verifier. + +// A float operand of unknown sign and unknown NaN-status: `x - y` is neither +// guaranteed non-negative nor guaranteed no-NaN, so `(x - y) < 0` must be kept. +func.func @compare_unknown_float_kept(%arg0: tensor<4xf64>, %arg1: tensor<4xf64>) -> tensor<4xi1> { + %0 = stablehlo.subtract %arg0, %arg1 : tensor<4xf64> + %c = stablehlo.constant dense<0.000000e+00> : tensor<4xf64> + %1 = stablehlo.compare LT, %0, %c : (tensor<4xf64>, tensor<4xf64>) -> tensor<4xi1> + return %1 : tensor<4xi1> +} + +// CHECK-LABEL: func.func @compare_unknown_float_kept +// CHECK: stablehlo.compare + +// A signed integer operand of unknown sign: `x < -1` must be kept. +func.func @compare_unknown_signed_kept(%arg0: tensor<4xi32>) -> tensor<4xi1> { + %c = stablehlo.constant dense<-1> : tensor<4xi32> + %0 = stablehlo.compare LT, %arg0, %c, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1> + return %0 : tensor<4xi1> +} + +// CHECK-LABEL: func.func @compare_unknown_signed_kept +// CHECK: stablehlo.compare