Skip to content

Use cast_to for fused AsType in compiled Metal kernels - #4351

Open
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:fix/compiled-subnormal-bool-cast
Open

Use cast_to for fused AsType in compiled Metal kernels#4351
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:fix/compiled-subnormal-bool-cast

Conversation

@katlun-lgtm

Copy link
Copy Markdown
Contributor

Follow-up to #4224 (#4205). The compiled/fused path still flushes float32 and bfloat16 subnormals to False on .astype(mx.bool_): build_kernel in mlx/backend/metal/compiled.cpp emits a raw static_cast for is_static_cast primitives, so a fused AsType never reaches cast_to. This swaps it for cast_to, which is already in every generated kernel's preamble via metal::utils().

Repro — needs at least two fusable ops; a single-op mx.compile(lambda x: x.astype(mx.bool_)) doesn't fuse and falls back to the already-correct copy kernel:

import numpy as np, mlx.core as mx

sub = mx.array(np.array([0x00000001] * 4, dtype=np.uint32)).view(mx.float32)
print(np.array(sub.astype(mx.bool_)))   # [True True True True] -- fixed by #4224

f = mx.compile(lambda x: mx.broadcast_to(x, (2, 4)).astype(mx.bool_))
print(np.array(f(sub)))                 # all False on main     -- this PR

broadcast_to is deliberate: the fused broadcast preserves the input bits (0x00000001 in and out), so what fails is the cast, not FTZ arithmetic. Same failure on bfloat16. float16 is unaffected — Metal's static_cast<bool> on half is already correct (probed directly with mx.fast.metal_kernel).

Out of scope, and not claimed fixed: !=, >, == on subnormals still diverge from the CPU backend after this. Eager x + 0.0 zeroes the bits too, so that's ordinary Metal FTZ arithmetic, a different class from the cast. I also checked the other static_cast in this file (constant-input loads around line 121): it casts the printed constant to its own dtype, so there is nothing to fix there.

Tested on an M3 Max: the new test in test_compile.py fails on current main and passes with the change; full python/tests suite is green (801 passed, 51 skipped, 11293 subtests).

Investigation and test runs were AI-assisted; I've reviewed and understand the change.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

…als to bool false

Follow-up to ml-explore#4224: the fused path in compiled.cpp still emitted a raw
static_cast for AsType, so .astype(mx.bool_) inside a compiled graph of
two or more fusable ops flushed float32/bfloat16 subnormals to False.
cast_to is already in the generated preamble via metal::utils().
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants