Skip to content

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

Draft
LngelKyo wants to merge 1 commit into
apache:mainfrom
LngelKyo:fix/cuda-min-max-nan
Draft

LngelKyo wants to merge 1 commit 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.

@LngelKyo
LngelKyo marked this pull request as draft September 18, 2026 18:08
@LngelKyo
LngelKyo marked this pull request as ready for review September 18, 2026 20:52
@LngelKyo
LngelKyo force-pushed the fix/cuda-min-max-nan branch from c87055f to 4232dd4 Compare September 21, 2026 22:53
@LngelKyo
LngelKyo marked this pull request as draft September 21, 2026 22:53
@LngelKyo
LngelKyo force-pushed the fix/cuda-min-max-nan branch from cbe6a06 to 5d7a3a0 Compare September 21, 2026 23:52
The CUDA codegen prints bare min()/max() for T.min/T.max, so a NaN
operand is dropped: with apache#20054's data only 4 of 8 lanes match bitwise,
and the +-0 ties disagree with the C host.

Emit ((a > b) || (a != a)) ? a : b for max (and the < form for min),
matching apache#20054's host-side semantics: a NaN on either side survives
and ties take b. Covers float16/bfloat16/float32/float64; integer and
other non-float min/max keep the base CodeGenC path, including the
per-lane vector expansion. Only the CUDA codegen is touched; the host
side is apache#20054's.

Operands are bound once via SSAGetID inside their own scope. The scope
matters: SSAGetID caches by printed text within the live scope and the
BufferStore dispatch does not invalidate it, so without it warp
allreduce -- which emits red_buf[0] = max(red_buf[0], shuffle_down(...))
repeatedly -- reads the pre-write value from the second statement on.

Constant operands take the simplified forms raised in the apache#20054
review: a constant NaN lhs is the result, a constant NaN rhs needs only
the NaN test, and a constant non-NaN operand drops the NaN clause,
using a reversed non-strict compare when the constant is on the right.
Vector constants keep the general per-lane path.

Tests go in tests/python/codegen/test_target_codegen_cuda.py, gpu-marked
and covering both compile paths through the autouse fixture: apache#20054's
data over four dtypes x {scalar, vec4}, the constant-operand matrix,
nested uses that would expose a missing parenthesis, a chained-statement
case that pins the SSA scoping, and an int32 vec4 regression against
numpy.
@LngelKyo
LngelKyo force-pushed the fix/cuda-min-max-nan branch from b629c8e to 6810d51 Compare September 22, 2026 00:08

This branch has not been deployed

No deployments
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