EnzymeAD / EnzymeAD/Enzyme-JAX

NoNanResultAnalysis proves add/sub no-nan for inf minus inf

Open
#2,673 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.