wuyii8941 opened a new issue, #19518:
URL: https://github.com/apache/tvm/issues/19518
## Expected behavior
`R.max()` and `R.min()` with NaN inputs should propagate NaN in the output,
consistent with IEEE 754 semantics: any comparison involving NaN returns NaN.
This is the behavior on CPU (`llvm` target).
## Actual behavior
On CUDA, `R.max()` and `R.min()` silently ignore NaN values and return the
max/min of only the non-NaN elements. CPU correctly propagates NaN.
`R.sum()` and `R.mean()` handle NaN correctly on both CPU and CUDA
(propagate NaN), so this is specific to `reduce_max` and `reduce_min`.
## Reproduction
```python
import numpy as np
import tvm
from tvm import relax
import tvm.relax.op as R
from tvm.s_tir import dlight
def build_reduce(shape, axis, reduce_fn):
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorStructInfo(shape, "float32"))
with bb.function("main", [x]):
with bb.dataflow():
out = bb.emit(reduce_fn(x, axis=axis, keepdims=False))
gv = bb.emit_output(out)
bb.emit_func_output(gv)
return bb.get()
def run_on_target(mod, x_np, target):
if target == "llvm":
pipeline =
tvm.ir.transform.Sequential([relax.transform.LegalizeOps()])
mod_l = pipeline(mod)
exe = relax.build(mod_l, target="llvm")
dev = tvm.cpu()
else:
pipeline = tvm.ir.transform.Sequential([
relax.transform.LegalizeOps(),
dlight.ApplyDefaultSchedule(dlight.gpu.Fallback()),
])
with tvm.target.Target("cuda"):
mod_l = pipeline(mod)
exe = relax.build(mod_l, target="cuda")
dev = tvm.cuda()
vm = relax.VirtualMachine(exe, device=dev)
return vm["main"](tvm.runtime.tensor(x_np, device=dev)).numpy()
np.random.seed(42)
shape = (4, 8)
x = np.random.randn(*shape).astype("float32")
x[1, 3] = float("nan") # inject NaN
# reduce_max
mod = build_reduce(shape, axis=1, reduce_fn=R.max)
cpu_out = run_on_target(mod, x, "llvm")
cuda_out = run_on_target(mod, x, "cuda")
print(f"Input row 1: {x[1]}")
print(f"CPU max row 1: {cpu_out[1]} (nan = correct, IEEE 754)")
print(f"CUDA max row 1: {cuda_out[1]} (finite value = WRONG)")
# reduce_min
mod_min = build_reduce(shape, axis=1, reduce_fn=R.min)
cpu_min = run_on_target(mod_min, x, "llvm")
cuda_min = run_on_target(mod_min, x, "cuda")
print(f"CPU min row 1: {cpu_min[1]} (nan = correct)")
print(f"CUDA min row 1: {cuda_min[1]} (finite value = WRONG)")
# reduce_sum — NOT affected (both propagate NaN)
mod_sum = build_reduce(shape, axis=1, reduce_fn=R.sum)
cpu_sum = run_on_target(mod_sum, x, "llvm")
cuda_sum = run_on_target(mod_sum, x, "cuda")
print(f"CPU sum row 1: {cpu_sum[1]} (nan = correct)")
print(f"CUDA sum row 1: {cuda_sum[1]} (nan = correct)")
```
## Expected output
```
CPU max row 1: nan
CUDA max row 1: nan <-- should match CPU
```
## Actual output
```
CPU max row 1: nan
CUDA max row 1: 1.579 <-- NaN silently dropped, returns max of non-NaN
values
```
## Root cause
The CUDA codegen for `reduce_max` / `reduce_min` uses a tree reduction with
`tir.max()` / `tir.min()`, which on CUDA maps to `fmaxf()` / `fminf()`. Per the
CUDA math API, `fmaxf(x, NaN) = x` — NaN is treated as missing rather than
propagated. The CPU backend uses a comparison-based reduction that naturally
propagates NaN per IEEE 754.
A correct NaN-propagating max would be: `result = (isnan(a) || isnan(b)) ?
NaN : max(a, b)`, or equivalently using the `__hmax_nan` / raw comparison
approach.
## Impact
Any model using `reduce_max` or `reduce_min` on data that may contain NaN
(e.g., attention masks with `-inf`/`NaN`, loss values, data with missing
entries) will **silently produce wrong results** on CUDA while appearing
correct on CPU. This is a correctness bug, not a precision issue.
## Environment
- TVM: main branch
- Target: `cuda` (tested on Tesla T4, sm_75)
- Python: 3.11
- OS: Ubuntu Linux
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
