EnzymeAD / EnzymeAD/Enzyme-JAX
NoNanResultAnalysis proves add/sub no-nan for inf minus inf
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
HUMAN: found this by auditing `NoNanResultAnalysis` in the same spirit as the `NonNegativeResultAnalysis` issues (#2648, #2651, #2656). The unsound inference is clear, but I want to be upfront: I could not construct an end-to-end miscompile, and I suspect it may be latent behind constant-folding and the analysis's current op coverage (details in the last section). Filing for visibility with the reasoning and a proposed fix rather than a PR, since I can't gate-verify a repro (no local build). Happy to be told this is already known / not worth guarding.
---
`NoNanResultAnalysis` proves a tensor value is never `NaN`. For `add` and `subtract` the rule is unsound: it can prove a result no-nan when that result is `inf - inf = NaN`.
## The unsound inference
`src/enzyme_ad/jax/Utils.cpp`, `NoNanResultAnalysis::localGuaranteed`, `AddOp`/`SubtractOp` case:
```cpp
} else if (isa(op)) {
// If any one of the operands is a Inf, the result is Inf. If both are Inf,
// the result is NaN.
auto lhsFinite = finiteResultAnalysis->guaranteed(op->getOperand(0), rewriter);
auto rhsFinite = finiteResultAnalysis->guaranteed(op->getOperand(1), rewriter);
if (lhsFinite && rhsFinite) {
return State::GUARANTEED;
}
recursiveCheck = true; // <-- fallback
operandsToCheck.append(op->getOperands().begin(), op->getOperands().end());
}
```
The comment states the hazard exactly: *both operands `Inf` produces `NaN`*. The `lhsFinite && rhsFinite` path is sound (finite operands cannot be `inf`). But when that path fails, control falls through to `recursivelyCheckOperands`, which returns `GUARANTEED` as soon as **both operands are merely no-nan** (`Utils.h` ~742, it only consults the no-nan lattice). `+inf` and `-inf` are no-nan-but-not-finite, so the fallback re-admits precisely the case the finite gate was meant to exclude.
`add`/`sub` is `NaN` iff an operand is `NaN`, or `+inf + (-inf)` / `-inf - (-inf)` etc. Ruling out the first needs both operands no-nan; ruling out the second needs at least one operand finite. The fallback checks only the first.
## Analysis-level witness
```mlir
%inf = stablehlo.constant dense<0x7FF0000000000000> : tensor // +inf: no-nan, not finite
%x = stablehlo.exponential %inf : tensor // exp(+inf) = +inf; no-nan-provable, not finite-provable
%y = stablehlo.subtract %x, %x : tensor // +inf - +inf = NaN
```
`NoNanResultAnalysis` returns `GUARANTEED` for `%y`:
- `finiteResultAnalysis->guaranteed(%x)` is false (`FiniteResultAnalysis::constantFloatCheck` requires `isFinite()`, and `%inf` fails it),
- so the finite gate fails and the fallback runs,
- `guaranteed`-no-nan holds for both `%x` operands (`constantFloatCheck` rejects only `isNaN`; `ExpOp` is no-nan-recursive),
- fallback returns `GUARANTEED`, but `%y` is `NaN`.
The fact feeds real consumers: `canApplyNoNanPattern` (`Utils.cpp` ~526) gates every no-nan rewrite on `guaranteedNoNanResult`, and the float `CompareOp` fold at `EnzymeHLOOpt.cpp` ~13680 folds comparisons assuming IEEE total order.
## Why I'm not sure it's reachable end-to-end (the uncertainty)
I tried to build a lit test that miscompiles and could not, for a structural reason worth stating:
- The only values the analysis proves *no-nan-yet-not-finite* are `±inf` **constants** and `exp`/data-movement of them. Any such constant subgraph (`subtract(+inf, +inf)`) is **constant-folded to a literal `NaN`** first, which the analysis then *correctly* rejects (`constantFloatCheck` sees the `NaN`). Folding masks the bug.
- I could find **no non-constant** value the analysis proves no-nan-yet-actually-infinite: block args aren't no-nan; `convert(int)` / `iota` are finite; `exp(finite)` is already treated as finite (via the separate overflow gap in `FiniteResultAnalysis`), so it takes the sound path; `floor`/`ceil` are no-nan-not-finite but only reach `inf` from an `inf` input that is itself not constructible.
So "reachable" here means *the analysis returns the wrong answer* (demonstrable, above), not *a pipeline emits wrong code* (I couldn't build that). It's a genuine unsound inference, but possibly latent behind folding and current op coverage. It would stop being latent if the analysis gains an op that yields a non-constant `inf` (or if `FiniteResultAnalysis` is tightened for overflow so `exp(finite)` is no longer treated finite).
## Proposed fix
Drop the fallback for `add`/`sub`; require both operands finite (the sound path already present):
```cpp
} else if (isa(op)) {
auto lhsFinite = finiteResultAnalysis->guaranteed(op->getOperand(0), rewriter);
auto rhsFinite = finiteResultAnalysis->guaranteed(op->getOperand(1), rewriter);
// add/sub is NaN-free only when no inf +/- inf can occur; that needs at least
// one finite operand, so a bare no-nan check on both is unsound.
if (lhsFinite && rhsFinite)
return State::GUARANTEED;
return State::NOTGUARANTEED;
}
```
A more precise (still sound) variant keeps the fallback but requires *at least one* operand finite in addition to both being no-nan (`finite +/- no-nan` cannot be `NaN`). The minimal version above is what I'd suggest unless the precision matters for a real rewrite.
Same class as #2648 / #2651 / #2656: a rule that reasons over the extended reals applied to IEEE values.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in src/enzyme_ad/jax/Utils.cpp at NoNanResultAnalysis::localGuaranteed and its AddOp/SubtractOp case; review the related analysis logic in Utils.h. Check the witness involving stablehlo exponential and subtract, then verify that add/sub no longer receives an unsound no-nan guarantee while preserving the finite-operand behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100