[spichardo]

Use the MLX GPU binary for device="gpu" / is_gpu_simulation=True on Apple Silicon

What

macOS has no CUDA. With this PR, GPU runs of the C++ backend on Apple Silicon use kspaceFirstOrder-MLX, the k-Wave
CUDA solver ported to MLX/Metal. It is released by kspacefirstorder-unified as kspaceFirstOrder-MLX-darwin plus
mlx.metallib (the PR adding it to kspacefirstorder-unified is waltsims/kspacefirstorder-unified#31).

• kwave/__init__.py:
  • New mlx architecture, darwin only, which downloads kspaceFirstOrder-MLX-darwin (saved as
    kspaceFirstOrder-MLX) and mlx.metallib into BINARY_PATH. Intel Macs keep getting no C++ binaries, as before.
  • MLX is an optional architecture: a failed download (e.g. a pinned release without the asset) logs a warning
    instead of aborting import kwave, so the OMP backend stays usable.
• SimulationExecutionOptions.binary_name, CppSimulation._resolve_binary_path:
  • GPU on darwin resolves to kspaceFirstOrder-MLX.
  • "MLX" is accepted as a backend value and maps to cpp/gpu in compat.
• Messages:
  • The "GPU not supported on macOS" errors now say that the MLX binary (Apple Silicon) is needed.
  • The CUDA compute-capability warning is skipped for the MLX binary.
• Docs: docstrings mention MLX on macOS.
• Tests:
  • Updated the two macOS-GPU error tests.
  • Added tests for MLX selection, darwin URLs and local filenames, and the non-fatal optional download.

Before merging

BINARY_VERSION must point to the first kspacefirstorder-unified release that contains kspaceFirstOrder-MLX-darwin
and mlx.metallib. This PR keeps v1.4.2, which has neither: macOS users then just get a warning and no GPU binary.

Testing

• pytest tests/test__init__.py tests/test_executor.py tests/test_cpp_simulation.py tests/test_checks.py tests/test_compute_capability_warning.py: 43 passed, 2 skipped (macOS).
• Manual, on an Apple M3 Max with a locally built release binary in kwave/bin/darwin:
  • kspaceFirstOrder(..., backend="cpp", device="gpu") matches device="cpu" (v1.4.2 OMP binary) to 1e-6 to 1.2e-5
    on ivp_3D_simulation, ivp_heterogeneous_medium and tvsp_homogeneous_medium_monopole.
  • The legacy kspace_first_order_2d_gpu (at_array_as_source) runs on the MLX binary.

Confidence Score: 5/5

The latest test fix appears safe to merge, with no remaining review findings.

Summary

Adds MLX binary downloads and GPU selection for Apple Silicon. The latest change fixes the macOS-specific test on Windows.

• The macOS installer looks for the MLX binary and its Metal library.
• C++ GPU simulations use MLX on Apple Silicon Macs.

_{Reviews (4) · Last reviewed commit: "Mock macOS platform in MLX binary test" · Reviewed by Greptile}

[spichardo]

@greptileai review

[spichardo]

@greptileai review