EnzymeAD / EnzymeAD/Enzyme-JAX
DivZeroPad incorrectly replaces NaN with 0
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
The DivZeroPad optimization pass:
```cpp
template
static LogicalResult getDefiningZeroPadding(OpTy op, PatternRewriter &rewriter,
stablehlo::PadOp &pad,
Value &otherArg,
bool &isOtherArgLHS) {
pad = op.getLhs().template getDefiningOp();
otherArg = op.getRhs();
isOtherArgLHS = false;
if (!pad) {
pad = op.getRhs().template getDefiningOp();
otherArg = op.getLhs();
isOtherArgLHS = true;
}
if (!pad)
return rewriter.notifyMatchFailure(op, "operands not produced by pad");
// if (!llvm::hasSingleElement(pad->getUsers()))
// return rewriter.notifyMatchFailure(op, "pad has multiple users");
if (!matchPattern(pad.getPaddingValue(), m_AnyZeroFloat()))
return rewriter.notifyMatchFailure(op, "padding value not zero");
return success();
}
struct DivZeroPad
: public CheckedOpRewritePattern {
using CheckedOpRewritePattern::CheckedOpRewritePattern;
LogicalResult matchAndRewriteImpl(stablehlo::DivOp op,
PatternRewriter &rewriter) const {
stablehlo::PadOp pad;
Value otherArg;
bool otherIsLHS;
if (failed(getDefiningZeroPadding(op, rewriter, pad, otherArg, otherIsLHS)))
return failure();
if (anyPadSizesNegative(pad))
return failure();
if (otherIsLHS)
return failure();
auto otherArgType = cast(otherArg.getType());
SmallVector limitDims = llvm::to_vector(otherArgType.getShape());
for (auto &&[limit, pad] : llvm::zip(limitDims, pad.getEdgePaddingHigh())) {
limit -= pad;
}
SmallVector interior = llvm::to_vector(pad.getInteriorPadding());
for (int64_t &value : interior) {
value += 1;
}
auto slice = stablehlo::SliceOp::create(rewriter, pad.getLoc(), otherArg,
pad.getEdgePaddingLow(), limitDims,
interior);
auto mul = stablehlo::DivOp::create(
rewriter, op.getLoc(),
otherIsLHS ? slice.getResult() : pad.getOperand(),
otherIsLHS ? pad.getOperand() : slice.getResult());
auto newPad = stablehlo::PadOp::create(
rewriter, pad.getLoc(), mul.getResult(), pad.getPaddingValue(),
pad.getEdgePaddingLowAttr(), pad.getEdgePaddingHighAttr(),
pad.getInteriorPaddingAttr());
rewriter.replaceOp(op, newPad);
return success();
}
};
```
The DivZeroPad optimization pattern contains a mathematical flaw that can silently mask numerical instability.
The pass looks for a division where the numerator has been padded with zeros (e.g., pad(A, 0.0) / B). To save compute, it optimizes this by slicing the denominator, performing the division only on the unpadded region, and then padding the final result with 0.0s (pad(A / slice(B), 0.0)).
The underlying assumption is that because the numerator's padding is 0.0, the result of the division in that region will always be 0.0.
However, in IEEE 754 floating-point math:
0.0 / 0.0 = NaN
0.0 / NaN = NaN
If the denominator (B) happens to contain 0.0 or NaN in that padded region, the unoptimized code correctly propagates a NaN. The optimized code skips the math entirely and outputs a 0.0, illegally hiding the NaN / divide-by-zero error.
### How to Reproduce
This was discovered using a custom differential fuzzer on test/lit_tests/divpad.mlir. When the denominator block argument is injected with tensors containing 0.0 or NaN, the optimized IR returns 0.000000e+00 while the unoptimized IR correctly returns NaN. (@wsmoses we talked about this in the meeting a few weeks back. I should have a PR for this ready in a few days but it currently produces a lot of false positives I believe)
### Root Cause
Looking at struct DivZeroPad, it appears to be a copy-paste of MulZeroPad (the variable is even still named mul inside the pass). For multiplication, 0.0 * X = 0.0 is safe. For division, the swap requires strict guards, which are currently missing.
### Proposed Fix
The pass needs to guarantee that the denominator in the padded region is safe before applying the rewrite.
This optimization should only fire if:
The denominator is proven safe: The pass checks that the denominator cannot contain zeros (via an analysis pass) AND cannot contain NaNs (via NoNanResultAnalysis).
OR the pass is restricted to fast-math: The pattern is moved so it only registers and runs when a global no_nan=true flag is enabled by the user.
(Disclaimer: The text of this issue was drafted with the assistance of an LLM to clearly summarize my debugging and fuzzer findings due to me being too lazy to write the issue myself)
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the DivZeroPad pattern and test/lit_tests/divpad.mlir, then compare optimized and unoptimized results when the denominator contains 0.0 or NaN. Trace the existing rewrite and relevant safety analyses, and add a regression test showing that padded-region NaNs are not replaced with zero; done means the optimization is applied only under the stated safety conditions.
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
- Mostly clear
- Newbie friendliness
- 48/100