diff --git a/kwave/__init__.py b/kwave/__init__.py index 0534de526..67b22602f 100644 --- a/kwave/__init__.py +++ b/kwave/__init__.py @@ -7,6 +7,7 @@ import warnings from pathlib import Path from typing import List +from urllib.error import URLError from urllib.request import urlretrieve # Test installation with: @@ -78,8 +79,15 @@ "vcruntime140_threads.dll", ] +# Metal kernels of the MLX GPU binary (Apple Silicon). The darwin MLX binary links MLX statically and loads +# this file from its own directory, so it is downloaded next to it. +MLX_METALLIB = "mlx.metallib" + EXECUTABLE_PREFIX = "kspaceFirstOrder-" -ARCHITECTURES = ["omp", "cuda"] +# GPU backend per platform: CUDA on Linux/Windows, MLX (Metal) on macOS/Apple Silicon. +ARCHITECTURES = ["omp", "cuda", "mlx"] +# Backends whose binaries are optional: a failed download only disables that backend instead of aborting `import kwave`. +OPTIONAL_ARCHITECTURES = ["mlx"] def _platform_binary_url(platform_name: str, architecture: str) -> list: @@ -92,11 +100,15 @@ def _platform_binary_url(platform_name: str, architecture: str) -> list: """ if architecture == "cuda" and platform_name == "darwin": return [] + if architecture == "mlx" and platform_name != "darwin": + return [] if platform_name == "darwin" and platform_name == PLATFORM and _darwin_unsupported: return [] - backend = "CUDA" if architecture == "cuda" else "OMP" + backend = {"cuda": "CUDA", "mlx": "MLX"}.get(architecture, "OMP") ext = ".exe" if platform_name == "windows" else "" exe_url = f"{_UNIFIED_RELEASE_URL}{EXECUTABLE_PREFIX}{backend}-{platform_name}{ext}" + if architecture == "mlx": + return [exe_url, _UNIFIED_RELEASE_URL + MLX_METALLIB] # Shared runtime DLLs ship with the OMP entry only — both backends find # them at the same BINARY_PATH, so attaching to both would download the # full 21-DLL bundle twice on first install. @@ -238,6 +250,19 @@ def _record_binary_metadata(binary_version: str, binary_filepath: str, binary_ur json.dump(metadata, outfile, indent=4) +def _remove_installed_binaries(system_os: str, bin_type: str) -> None: + """Remove every local file (and its metadata) of a backend, e.g. after an incomplete download of a bundle.""" + for url in URL_DICT[system_os][bin_type]: + filename = _local_filename(url.split("/")[-1]) + for path in (BINARY_PATH / filename, BINARY_PATH / f"{filename}_metadata.json"): + try: + os.remove(path) + except FileNotFoundError: + pass + except OSError: # pragma: no cover - defensive + logging.warning(f"Could not remove {path}; please delete it manually.") + + def download_binaries(system_os: str, bin_type: str): """ Download binary from release url @@ -255,6 +280,10 @@ def download_binaries(system_os: str, bin_type: str): binary_version, asset_name = url.split("/")[-2:] filename = _local_filename(asset_name) + # Keep valid cached files: a missing optional backend must not force (network) re-downloads of the others. + if system_os == PLATFORM and _is_binary_present(filename, bin_type): + continue + logging.log(logging.INFO, f"Downloading {asset_name} -> {filename} in {BINARY_PATH}...") # Create the directory if it does not yet exist @@ -267,6 +296,15 @@ def download_binaries(system_os: str, bin_type: str): _ensure_executable(binary_filepath) _record_binary_metadata(binary_version=binary_version, binary_filepath=binary_filepath, binary_url=url, filename=filename) + except URLError as e: + if bin_type not in OPTIONAL_ARCHITECTURES: + raise + # e.g. the pinned release does not ship this backend yet: keep the other backends usable. A partial + # bundle (e.g. the MLX executable without mlx.metallib) cannot run, so remove what is already there. + _remove_installed_binaries(system_os, bin_type) + logging.warning(f"Could not download {asset_name} ({e}); the {bin_type.upper()} backend will not be available.") + return + except TimeoutError: logging.log( logging.WARN, diff --git a/kwave/compat.py b/kwave/compat.py index 70f393914..5a4809bfd 100644 --- a/kwave/compat.py +++ b/kwave/compat.py @@ -54,7 +54,7 @@ def options_to_kwargs(simulation_options=None, execution_options=None): if opts.backend == "python": kwargs["backend"] = "python" kwargs["device"] = "gpu" if opts.is_gpu_simulation else "cpu" - elif opts.backend == "CUDA": + elif opts.backend in ("CUDA", "MLX"): kwargs["backend"] = "cpp" kwargs["device"] = "gpu" elif opts.backend == "OMP": diff --git a/kwave/executor.py b/kwave/executor.py index fc986228d..6a70be369 100644 --- a/kwave/executor.py +++ b/kwave/executor.py @@ -26,7 +26,8 @@ def _make_binary_executable(self): if not binary_path.exists(): if kwave.PLATFORM == "darwin" and self.execution_options.is_gpu_simulation: raise ValueError( - "GPU simulations are currently not supported on MacOS. " + f"GPU binary not found at {binary_path}. GPU simulations on macOS require Apple Silicon " + "and the MLX binary (kspaceFirstOrder-MLX). " "Try running the simulation on CPU by setting is_gpu_simulation=False." ) raise FileNotFoundError(f"Binary not found at {binary_path}") diff --git a/kwave/kspaceFirstOrder.py b/kwave/kspaceFirstOrder.py index 94352d2f1..d2d982b93 100644 --- a/kwave/kspaceFirstOrder.py +++ b/kwave/kspaceFirstOrder.py @@ -177,7 +177,8 @@ def kspaceFirstOrder( the compiled C++ binary. Default ``"python"``. device: ``"cpu"`` or ``"gpu"``. For ``backend="python"`` this selects NumPy (cpu) vs CuPy (gpu). For ``backend="cpp"`` it - selects the OMP vs CUDA binary. Default ``"cpu"``. + selects the OMP vs GPU binary (CUDA on Linux/Windows, + MLX on macOS with Apple Silicon). Default ``"cpu"``. dtype: Numerical precision for state arrays in the Python backend. Accepts dtype-like input — a numpy dtype (``np.float32``, ``np.float64``; cupy aliases like ``cp.float32`` work since cupy @@ -207,7 +208,7 @@ def kspaceFirstOrder( debug: Print detailed diagnostic output. Default ``False``. num_threads: Thread count for the C++ OMP binary. ``None`` uses all available cores. Default ``None``. - device_num: GPU device index for CUDA execution. Default ``None``. + device_num: GPU device index for CUDA/MLX execution. Default ``None``. binary_path: Path to a custom C++ binary. When ``None`` (default), the binary bundled with ``k-wave-data`` is used. Only applies when ``backend="cpp"``. diff --git a/kwave/kspaceFirstOrder2D.py b/kwave/kspaceFirstOrder2D.py index 5dd5aeb6b..ae843a0e9 100644 --- a/kwave/kspaceFirstOrder2D.py +++ b/kwave/kspaceFirstOrder2D.py @@ -33,7 +33,8 @@ def kspace_first_order_2d_gpu( 2D time-domain simulation of wave propagation on a GPU using C++ CUDA code. This function provides a blind interface to the C++/CUDA version of kspaceFirstOrder2D - (called kspaceFirstOrder-CUDA). The function works by saving the input files to disk, + (called kspaceFirstOrder-CUDA; kspaceFirstOrder-MLX on macOS with Apple Silicon). + The function works by saving the input files to disk, running the C++/CUDA simulation, and loading the output files back into Python. Parameters: diff --git a/kwave/kspaceFirstOrder3D.py b/kwave/kspaceFirstOrder3D.py index b946ec803..76bf6ab61 100644 --- a/kwave/kspaceFirstOrder3D.py +++ b/kwave/kspaceFirstOrder3D.py @@ -34,6 +34,8 @@ def kspaceFirstOrder3DG( This function provides a blind interface to the C++/CUDA version of kspaceFirstOrder3D (called kspaceFirstOrder-CUDA). The function works by saving the input files to disk, running the C++/CUDA simulation, and loading the output files back into Python. + On macOS with Apple Silicon, the MLX port of the same code (kspaceFirstOrder-MLX) is used + instead, running on the Apple GPU through Metal. Parameters: ----------- diff --git a/kwave/options/simulation_execution_options.py b/kwave/options/simulation_execution_options.py index cea2b9e88..2859d5068 100644 --- a/kwave/options/simulation_execution_options.py +++ b/kwave/options/simulation_execution_options.py @@ -17,7 +17,8 @@ class SimulationExecutionOptions: Attributes: backend: Execution backend to use. Options: - "OMP": C++ OpenMP binary (default for CPU) - - "CUDA": C++ CUDA binary (default for GPU) + - "CUDA": C++ CUDA binary (default for GPU on Linux/Windows) + - "MLX": C++ MLX binary for the Apple Silicon GPU (default for GPU on macOS) - "python": Pure Python/CuPy solver (no external binaries required) """ @@ -39,7 +40,7 @@ def __init__( checkpoint_interval: Optional[int] = None, # [seconds] checkpoint_timesteps: Optional[int] = None, # [timestep integer] checkpoint_file: Optional[Path | str] = None, # [path to hdf5 file] - backend: Optional[str] = None, # "OMP", "CUDA", or "python" + backend: Optional[str] = None, # "OMP", "CUDA", "MLX", or "python" ): self._num_threads_explicit = num_threads is not None self.backend = backend @@ -113,7 +114,7 @@ def backend(self) -> Optional[str]: @backend.setter def backend(self, value: Optional[str]): - valid_backends = [None, "OMP", "CUDA", "python"] + valid_backends = [None, "OMP", "CUDA", "MLX", "python"] if value not in valid_backends: raise ValueError(f"Backend must be one of {valid_backends}, got '{value}'") self._backend = value @@ -136,14 +137,18 @@ def is_gpu_simulation(self, value: Optional[bool]): @property def binary_name(self) -> str: - valid_binary_names = ["kspaceFirstOrder-OMP", "kspaceFirstOrder-CUDA"] + valid_binary_names = ["kspaceFirstOrder-OMP", "kspaceFirstOrder-CUDA", "kspaceFirstOrder-MLX"] if self._binary_name is None: # set default binary name based on GPU simulation value if self.is_gpu_simulation is None: raise ValueError("`is_gpu_simulation` must be set to either True or False before determining the binary name.") - if self.is_gpu_simulation: - self._binary_name = "kspaceFirstOrder-CUDA" + if self.backend == "MLX": + # explicit choice of the Apple Silicon GPU binary + self._binary_name = "kspaceFirstOrder-MLX" + elif self.is_gpu_simulation: + # macOS has no CUDA: the GPU binary runs on the Apple Silicon GPU through MLX (Metal) + self._binary_name = "kspaceFirstOrder-MLX" if PLATFORM == "darwin" else "kspaceFirstOrder-CUDA" else: self._binary_name = "kspaceFirstOrder-OMP" diff --git a/kwave/solvers/cpp_simulation.py b/kwave/solvers/cpp_simulation.py index 65a95b74b..a55e125f9 100644 --- a/kwave/solvers/cpp_simulation.py +++ b/kwave/solvers/cpp_simulation.py @@ -339,7 +339,8 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path: When *binary_path* is provided the caller's custom binary is used and verified to exist. Otherwise the bundled binary is selected based on - *device* (``"gpu"`` → CUDA, anything else → OMP) and verified to exist. + *device* (``"gpu"`` → CUDA, or MLX on macOS; anything else → OMP) and + verified to exist. Args: device: ``"cpu"`` or ``"gpu"``. @@ -352,8 +353,8 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path: Raises: FileNotFoundError: When the binary does not exist at the resolved path. - ValueError: When ``device="gpu"`` is requested on macOS where no - CUDA binary is available. + ValueError: When ``device="gpu"`` is requested on macOS and the MLX + binary is not available (e.g. on an Intel Mac). """ if binary_path is not None: resolved = Path(binary_path) @@ -363,11 +364,18 @@ def _resolve_binary_path(device: str, binary_path: str | None = None) -> Path: import kwave - binary_name = "kspaceFirstOrder-CUDA" if device == "gpu" else "kspaceFirstOrder-OMP" + if device == "gpu": + # macOS has no CUDA: use the MLX binary, which runs on the Apple Silicon GPU (Metal) + binary_name = "kspaceFirstOrder-MLX" if kwave.PLATFORM == "darwin" else "kspaceFirstOrder-CUDA" + else: + binary_name = "kspaceFirstOrder-OMP" resolved = kwave.BINARY_PATH / binary_name if not resolved.exists(): if kwave.PLATFORM == "darwin" and device == "gpu": - raise ValueError("GPU simulations are not supported on macOS. Use device='cpu'.") + raise ValueError( + f"GPU binary {binary_name} not found at {resolved}. GPU simulations on macOS require Apple Silicon " + "and the MLX binary from the k-wave-python binary release. Use device='cpu' otherwise." + ) raise FileNotFoundError(f"C++ binary not found at {resolved}. Install with: pip install k-wave-data") return resolved @@ -402,7 +410,7 @@ def _execute(self, input_file, output_file, *, device, num_threads, device_num, resolved = self._resolve_binary_path(device, binary_path) resolved.chmod(resolved.stat().st_mode | stat.S_IEXEC) - if device == "gpu": + if device == "gpu" and not resolved.name.startswith("kspaceFirstOrder-MLX"): self._warn_if_unsupported_gpu() # Build command-line options diff --git a/tests/test__init__.py b/tests/test__init__.py index 0f428ad3b..4a6ede629 100644 --- a/tests/test__init__.py +++ b/tests/test__init__.py @@ -102,5 +102,101 @@ def test_existing_non_executable_binary_is_healed(tmp_path, monkeypatch): assert post_mode & stat.S_IXOTH, "other exec bit not healed on cache hit" +def test_darwin_gpu_urls_point_to_mlx_binary_and_metallib(): + """macOS has no CUDA: its GPU backend is the MLX binary plus the Metal kernel library it loads at run time.""" + import kwave + + mlx_urls = kwave._platform_binary_url("darwin", "mlx") + if kwave.PLATFORM == "darwin" and kwave._darwin_unsupported: + assert mlx_urls == [] + return + assert [url.rsplit("/", 1)[-1] for url in mlx_urls] == ["kspaceFirstOrder-MLX-darwin", "mlx.metallib"] + assert all(url.startswith(kwave._UNIFIED_RELEASE_URL) for url in mlx_urls) + assert kwave._platform_binary_url("darwin", "cuda") == [] + assert kwave._platform_binary_url("linux", "mlx") == [] + assert kwave._platform_binary_url("windows", "mlx") == [] + + +def test_local_filename_strips_platform_suffix_for_mlx(): + import kwave + + assert kwave._local_filename("kspaceFirstOrder-MLX-darwin") == "kspaceFirstOrder-MLX" + assert kwave._local_filename("mlx.metallib") == "mlx.metallib" + + +def test_failed_optional_mlx_download_does_not_abort(tmp_path, monkeypatch): + """A release without the MLX assets must not break `import kwave` (the OMP backend stays usable).""" + from urllib.error import HTTPError + + import kwave + + monkeypatch.setattr(kwave, "BINARY_PATH", tmp_path) + + def fake_urlretrieve(url, dest): + raise HTTPError(url, 404, "Not Found", None, None) + + monkeypatch.setattr(kwave, "urlretrieve", fake_urlretrieve) + monkeypatch.setattr(kwave, "URL_DICT", {kwave.PLATFORM: {"mlx": ["https://example.invalid/kspaceFirstOrder-MLX-darwin"]}}) + kwave.download_binaries(kwave.PLATFORM, "mlx") # must not raise + + monkeypatch.setattr(kwave, "URL_DICT", {kwave.PLATFORM: {"omp": ["https://example.invalid/kspaceFirstOrder-OMP-darwin"]}}) + with pytest.raises(HTTPError): + kwave.download_binaries(kwave.PLATFORM, "omp") + + +def _fake_release(tmp_path, monkeypatch, fail_assets): + """Point kwave at a fake release with an OMP binary and a two-file MLX bundle; assets in fail_assets 404.""" + from urllib.error import HTTPError + + import kwave + + base = "https://example.invalid/releases/download/v9.9.9/" + url_dict = { + kwave.PLATFORM: { + "omp": [base + "kspaceFirstOrder-OMP-darwin"], + "mlx": [base + "kspaceFirstOrder-MLX-darwin", base + "mlx.metallib"], + } + } + downloads = [] + + def fake_urlretrieve(url, dest): + downloads.append(url.rsplit("/", 1)[-1]) + if url.rsplit("/", 1)[-1] in fail_assets: + raise HTTPError(url, 404, "Not Found", None, None) + Path(dest).write_bytes(b"payload-" + url.encode()) + + monkeypatch.setattr(kwave, "BINARY_PATH", tmp_path) + monkeypatch.setattr(kwave, "URL_DICT", url_dict) + monkeypatch.setattr(kwave, "ARCHITECTURES", ["omp", "mlx"]) + monkeypatch.setattr(kwave, "urlretrieve", fake_urlretrieve) + return downloads + + +def test_cached_omp_is_not_downloaded_again_when_mlx_is_missing(tmp_path, monkeypatch): + """A missing optional MLX bundle must not re-download (or require the network for) a valid cached OMP binary.""" + import kwave + + downloads = _fake_release(tmp_path, monkeypatch, fail_assets={"kspaceFirstOrder-MLX-darwin", "mlx.metallib"}) + kwave.install_binaries() + assert (tmp_path / "kspaceFirstOrder-OMP").exists() + assert not kwave.binaries_present() # MLX still missing + + downloads.clear() + kwave.install_binaries() # second import: OMP is cached and valid + assert "kspaceFirstOrder-OMP-darwin" not in downloads + assert downloads == ["kspaceFirstOrder-MLX-darwin"] + + +def test_incomplete_mlx_bundle_is_removed(tmp_path, monkeypatch): + """If mlx.metallib fails after the MLX executable was downloaded, the unusable executable must not stay installed.""" + import kwave + + _fake_release(tmp_path, monkeypatch, fail_assets={"mlx.metallib"}) + kwave.download_binaries(kwave.PLATFORM, "mlx") + assert not (tmp_path / "kspaceFirstOrder-MLX").exists() + assert not (tmp_path / "kspaceFirstOrder-MLX_metadata.json").exists() + assert not (tmp_path / "mlx.metallib").exists() + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_cpp_simulation.py b/tests/test_cpp_simulation.py index b99a0062c..decb8bd75 100644 --- a/tests/test_cpp_simulation.py +++ b/tests/test_cpp_simulation.py @@ -70,10 +70,17 @@ def test_default_missing_cuda_on_linux_raises_file_not_found(self, tmp_path): with pytest.raises(FileNotFoundError, match="pip install k-wave-data"): CppSimulation._resolve_binary_path("gpu") - def test_default_gpu_on_macos_raises_value_error(self, tmp_path): - # macOS + gpu must raise before the binary-existence check. + def test_default_gpu_on_macos_selects_mlx_binary(self, tmp_path): + mlx_binary = tmp_path / "kspaceFirstOrder-MLX" + mlx_binary.write_text("#!/bin/sh\n") with patch("kwave.BINARY_PATH", tmp_path), patch("kwave.PLATFORM", "darwin"): - with pytest.raises(ValueError, match="not supported on macOS"): + resolved = CppSimulation._resolve_binary_path("gpu") + assert resolved.name == "kspaceFirstOrder-MLX" + + def test_default_gpu_on_macos_missing_mlx_raises_value_error(self, tmp_path): + # macOS + gpu without the MLX binary (e.g. Intel Mac) must raise a ValueError pointing to device='cpu'. + with patch("kwave.BINARY_PATH", tmp_path), patch("kwave.PLATFORM", "darwin"): + with pytest.raises(ValueError, match="kspaceFirstOrder-MLX not found"): CppSimulation._resolve_binary_path("gpu") @pytest.mark.skipif(sys.platform == "win32", reason="Windows does not honor Unix executable bits") diff --git a/tests/test_executor.py b/tests/test_executor.py index afb4aa810..2d08ee9c1 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -255,9 +255,7 @@ def test_cpu_environment_variable(self): del os.environ["KWAVE_FORCE_CPU"] def test_executor_gpu_cuda_failure_darwin(self): - expected_error_msg = ( - "GPU simulations are currently not supported on MacOS. Try running the simulation on CPU by setting is_gpu_simulation=False." - ) + expected_error_msg = "GPU simulations on macOS require Apple Silicon and the MLX binary" # Configure the mock path object mock_binary_path = MagicMock(spec=Path) mock_binary_path.chmod.side_effect = FileNotFoundError diff --git a/tests/test_simulation_execution_options.py b/tests/test_simulation_execution_options.py index 0af0c330a..f461de46e 100644 --- a/tests/test_simulation_execution_options.py +++ b/tests/test_simulation_execution_options.py @@ -10,6 +10,8 @@ OMP_BINARY_NAME = "kspaceFirstOrder-OMP{}".format(".exe" if PLATFORM == "windows" else "") CUDA_BINARY_NAME = "kspaceFirstOrder-CUDA{}".format(".exe" if PLATFORM == "windows" else "") +# macOS has no CUDA: the default GPU binary there is the MLX (Apple Silicon) one +GPU_BINARY_NAME = "kspaceFirstOrder-MLX" if PLATFORM == "darwin" else CUDA_BINARY_NAME class TestSimulationExecutionOptions(unittest.TestCase): @@ -73,12 +75,19 @@ def test_is_gpu_simulation_setter(self): options = self.default_options options.is_gpu_simulation = True self.assertTrue(options.is_gpu_simulation) - self.assertEqual(options.binary_name, CUDA_BINARY_NAME) + self.assertEqual(options.binary_name, GPU_BINARY_NAME) options.is_gpu_simulation = False self.assertFalse(options.is_gpu_simulation) self.assertEqual(options.binary_name, OMP_BINARY_NAME) + def test_mlx_backend_selects_mlx_binary(self): + """backend="MLX" alone (without is_gpu_simulation=True) must select the MLX binary (a macOS-only backend).""" + with patch("kwave.options.simulation_execution_options.PLATFORM", "darwin"): + options = SimulationExecutionOptions(backend="MLX") + self.assertFalse(options.is_gpu_simulation) + self.assertEqual(options.binary_name, "kspaceFirstOrder-MLX") + def test_device_num_setter_invalid(self): """Test setting an invalid device number.""" options = self.default_options