Skip to content

Emit valid kernel source for non-finite float constants - #3648

Merged
angeloskath merged 2 commits into
ml-explore:mainfrom
tillahoffmann:fix-nan-constant-codegen
Jun 11, 2026
Merged

angeloskath merged 2 commits into
ml-explore:mainfrom
tillahoffmann:fix-nan-constant-codegen

Conversation

@tillahoffmann

Copy link
Copy Markdown
Contributor

Summary

Compiling a function that bakes a NaN (or infinity) scalar constant into a fused kernel fails to compile. print_float_constant streams the value with operator<<, producing a bare token like nan, which is interpolated verbatim into the generated kernel source and is not a valid identifier:

ternary_ops.h:23:35: error: use of undeclared identifier 'nan'
  auto tmp_D = static_cast<float>(nan);

Eager evaluation is fine; only the compiled path is affected.

Repro

import mlx.core as mx
f = mx.compile(lambda x: mx.where(x > 0, x, float("nan")))
f(mx.array([1.0, -1.0]))   # before: Metal compile error; after: array([1, nan])

Fix

Emit the NAN / INFINITY / -INFINITY macros for non-finite values. Both kernel toolchains provide them — Metal via <metal_stdlib> (#define NAN __builtin_nanf(""), #define INFINITY __builtin_inff()) and the CPU C++ codegen via <cmath> — so the same emitted text is valid for both backends. The finite check widens to double first so it is exact for every float type.

Tests

  • tests/compile_tests.cpp: test compile with non-finite constants.
  • python/tests/test_compile.py: test_compile_nonfinite_constants (float32/float16/bfloat16, NaN and -inf).

Verified red→green on both: the unfixed build reproduces the nan error, the fixed build computes the expected values. Full C++ compile suite and test_compile.py pass.

🤖 Generated with Claude Code

Compiling a function that bakes a NaN (or infinity) scalar constant into a
fused kernel failed because print_float_constant streamed a bare token such
as `nan`, which is not a valid identifier in the generated kernel source
(e.g. Metal: "use of undeclared identifier 'nan'").

Emit the NAN / INFINITY / -INFINITY macros instead; both the Metal
(<metal_stdlib>) and CPU (<cmath>) kernel toolchains provide them. The
finite check widens to double first so it is exact for every float type.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The prebuilt CPU preamble is produced by running the preprocessor over
compiled_preamble.h, which consumes the NAN / INFINITY macro definitions
(macros vanish during preprocessing, leaving only the expanded <cmath>
declarations). print_float_constant emits those macros for non-finite
constants, so the generated kernel failed to compile with "'NAN' was not
declared in this scope" on the Linux/Windows prebuilt-preamble path.

Define the macros in the generated source when they are missing, backed
by std::numeric_limits (available via <limits>). The #ifndef guards make
this a no-op on paths where <cmath>'s macros are already in scope (e.g.
the include-based preamble path), avoiding redefinition clashes.

Also drop the now-dead `#define inf` from the CUDA backends: non-finite
constants are emitted as the standard INFINITY/NAN macros (which CUDA
provides), so the bare `inf` token is no longer generated.

Verified on a Linux + g++ build (the failing path): the regression test
fails with the original "'NAN' was not declared" error without this
change and passes with it.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
@tillahoffmann
tillahoffmann force-pushed the fix-nan-constant-codegen branch from 2061983 to 9db52bc Compare June 10, 2026 12:54
@angeloskath
angeloskath merged commit d6385cf into ml-explore:main Jun 11, 2026
16 checks passed
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.

3 participants