Skip to content

[Fix][CUDA] Preserve NaNs in floating-point min and max - #20387

Open
LngelKyo wants to merge 4 commits into
apache:mainfrom
LngelKyo:fix/cuda-min-max-nan
Open

LngelKyo wants to merge 4 commits into
apache:mainfrom
LngelKyo:fix/cuda-min-max-nan

Conversation

@LngelKyo

@LngelKyo LngelKyo commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

Bug: on main, the CUDA codegen prints bare min()/max() for T.min/T.max on float32/float16/bfloat16, so a NaN operand is dropped

With #20054's test data, the CUDA target gives 4/8 bitwise matches: the NaN lanes are discarded and the ±0 ties disagree with the C host.

Fix

Same semantics as #20054: emit (((a > b) || (a != a)) ? a : b) for max (and the < form for min), so a NaN on either side is preserved and ties take b. Scalar form binds each operand once via SSA and references the temporaries; the vector form expands the same expression per lane. Covers float16/bfloat16/float32/float64. Integer and other non-float min/max keep the existing base-CodeGenC path, including the per-lane vector expansion. Only the CUDA codegen is touched; the host side is #20054's.

Generated source (scalar and per-lane vector), from the branch:

[float32/scalar] C_ptr[((int)threadIdx.x)] = (((v_ > v__1) || (v_ != v_)) ? v_ : v__1);
[float32/vec4]   __1.x = ((v_.x > v__1.x) || (v_.x != v_.x)) ? v_.x : v__1.x;

Tests

Folded into tests/python/codegen/test_target_codegen_cuda.py: 4 dtypes (float32/float64/float16/bfloat16) × 2 ops × {scalar, vec4}, a composite case (C[i] = max(A[i], B[i]) + 1.0 — catches the ternary being parsed as (x + cond) ? va : vb; NaN lanes assert NaN, finite lanes bitwise), and an int32 vec4 regression against numpy. Both compile paths (nvcc, nvrtc) are covered by the autouse fixture.

A further defect surfaced in CI: the scalar path's SSA bindings stayed alive across statements. SSAGetID caches by printed text within the live scope and the base-codegen BufferStore dispatch does not invalidate it, so warp allreduce — which emits red_buf[0] = max(red_buf[0], shuffle_down(...)) repeatedly — read the pre-write value from the second statement on: the first CI run had 92 failures in test_gpu_codegen_allreduce, all finite values with max systematically too small. The scalar bindings are now scoped to the statement (BeginScope/EndScope, the same shape the vector path already had). A chained-statements test (C[v] = max(C[v], A[v]); C[v] = max(C[v], B[v]), C pre-filled) fails on the pre-fix code and passes with the fix.

Verification

RTX A6000, CUDA 13.0.88:

  • positive run: 40/40 passed (all four dtypes, scalar and vec4, both compile paths, plus the chained case);
  • allreduce suite on the fixed branch: 328 passed; full test_target_codegen_cuda.py: 404 passed, 6 skipped;
  • negative control (predicate forced to return false): 34 failed; the 4 int cases and the 2 chained cases stay green, the latter by construction since the control bypasses the new path;
  • restored: 40/40 again.
  • ruff@0.12.3 check/format and clang-format 20.1.8 --dry-run --Werror clean.

Partially addresses #19579 (CUDA side); host side is #20054.

cc @tlopex @yongwww @swjng — CI will likely need approval as before.

Follows apache#20054, which made the C host preserve a NaN operand
in min/max; the CUDA codegen still emitted the bare fmin/fmax-style
min()/max(), which drops a NaN operand under IEEE semantics. The CUDA
path now emits a NaN-preserving ternary for float16/bfloat16/float32/
float64 (scalar and per-lane vector forms): NaN-free pairs compile to
the same min()/max() as before.

- integer and other non-float min/max keep the base CodeGenC path,
  including the vectorized per-lane expansion (a naive scalar print
  would emit vector overloads CUDA does not have);
- scalar form binds each operand once (SSA) before the ternary —
  operands are not re-evaluated per clause;
- float64 included so CUDA matches the C host's semantics.

Tests in test_target_codegen_cuda.py (gpu-marked, nvcc/nvrtc via the
autouse fixture): apache#20054's data with bitwise assertions over all four
dtypes, plus an int32 lanes=4 regression against numpy.
CUDA float add normalizes a NaN result payload, so bitwise comparison
of the NaN lanes after 'max(a, b) + 1.0' is not stable. The composite
case now asserts NaN on lanes 0-2 and bitwise equality with
expected + 1.0 on the finite lanes; the purpose of the test (catching a
ternary parsed as (x + cond) ? va : vb) is unchanged.
The scalar NaN-preserving path bound its operands with SSAGetID but left
the bindings alive after the statement. SSAGetID caches by expression
text within the live scope, and CodeGenC::Dispatch_(BufferStoreNode) does
not invalidate that cache — so a block that emits
red_buf[0] = max(red_buf[0], shuffle_down(...)) repeatedly (warp
allreduce) hit the cache from the second statement on and read the
pre-write value, silently dropping most of the reduction. The GPU
unittest shard caught it: 92 failures in test_gpu_codegen_allreduce, all
finite values, max systematically too small.

Bind inside BeginScope/EndScope, same shape as the vector path already
uses. A chained-statements test (C[v] = max(A[v], B[v]); C[v] =
max(C[v], D[v])) added to pin it.
@LngelKyo
LngelKyo marked this pull request as draft September 18, 2026 18:08
The previous form (max(A, B) then max(C, D)) never read the same
expression text twice with a write in between, so it passed on the
pre-fix code. Rewrite as C[v] = max(C[v], A[v]); C[v] = max(C[v], B[v])
— both statements print the same read text for C — and pre-fill C with
a known c0 (zeros) so a cache hit reads a known wrong value; with
c0 = 0 and a > max(c0, b) on every lane, the pre-fix code computes
max(c0, b) on at least one lane and fails.
@LngelKyo
LngelKyo marked this pull request as ready for review September 18, 2026 20:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant