diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index 0a15f17..fe11a6c 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -1234,6 +1234,11 @@ def main(): """ import sys, argparse + try: + from cnmf.nmf_gpu import configure_nmf_engine, gpu_kwargs_from_args, parse_gpu_args, validate_engine_args_for_command + except ImportError: + from nmf_gpu import configure_nmf_engine, gpu_kwargs_from_args, parse_gpu_args, validate_engine_args_for_command + parser = argparse.ArgumentParser() parser.add_argument('command', type=str, choices=['prepare', 'factorize', 'combine', 'consensus', 'k_selection_plot']) @@ -1257,11 +1262,18 @@ def main(): parser.add_argument('--local-neighborhood-size', type=float, help='[consensus] Fraction of the number of replicates to use as nearest neighbors for local density filtering', default=0.30) parser.add_argument('--show-clustering', dest='show_clustering', help='[consensus] Produce a clustergram figure summarizing the spectra clustering', action='store_true') parser.add_argument('--build-reference', dest='build_reference', help='[consensus] Generates a reference spectra for use in starCAT', action='store_true', default=True) + parse_gpu_args(parser) args = parser.parse_args() + try: + engine_commands = ('factorize',) + validate_engine_args_for_command(args, engine_commands) + except ValueError as e: + parser.error(str(e)) cnmf_obj = cNMF(output_dir=args.output_dir, name=args.name) + cnmf_obj = configure_nmf_engine(cnmf_obj, engine=args.engine or 'cpu', gpu_kwargs=gpu_kwargs_from_args(args)) if args.command == 'prepare': cnmf_obj.prepare(args.counts, components=args.components, n_iter=args.n_iter, densify=args.densify, diff --git a/src/cnmf/nmf_gpu.py b/src/cnmf/nmf_gpu.py index d1045fd..4487860 100644 --- a/src/cnmf/nmf_gpu.py +++ b/src/cnmf/nmf_gpu.py @@ -69,6 +69,76 @@ } +def parse_gpu_args(parser): + """Register cNMF CLI flags for the optional PyTorch GPU NMF engine.""" + group = parser.add_argument_group("NMF engine options") + group.add_argument("--engine", type=str.lower, choices=["cpu", "gpu"], help="[factorize] NMF engine to use (default cpu)") + group.add_argument("--gpu-device", type=str, help="[factorize,gpu] Device for GPU NMF: auto, cpu, cuda, cuda:N, or mps") + group.add_argument("--gpu-dtype", type=str.lower, choices=["auto", "fp32", "fp64", "bf16"], help="[factorize,gpu] Storage and matmul dtype for GPU NMF (default auto)") + group.add_argument("--gpu-allow-tf32", action="store_const", const=True, help="[factorize,gpu] Allow TF32 for CUDA fp32 matrix multiplication") + group.add_argument("--gpu-compile", action="store_const", const=True, help="[factorize,gpu] Enable torch.compile for the GPU NMF update step") + group.add_argument("--gpu-eps", type=float, help="[factorize,gpu] Multiplicative-update denominator guard") + group.add_argument("--gpu-check-every", type=int, help="[factorize,gpu] Eager-mode convergence check interval") + group.add_argument("--gpu-compile-block", type=int, help="[factorize,gpu] Number of MU iterations per compiled block") + return parser + + +def gpu_kwargs_from_args(args): + """Collect parsed cNMF CLI GPU flags into a kernel gpu_kwargs dict.""" + raw = { + "device": args.gpu_device, + "dtype": args.gpu_dtype, + "allow_tf32": args.gpu_allow_tf32, + "compile": args.gpu_compile, + "eps": args.gpu_eps, + "check_every": args.gpu_check_every, + "compile_block": args.gpu_compile_block, + } + if args.engine != "gpu": + if any(value is not None for value in raw.values()): + raise ValueError("GPU options require --engine gpu") + return None + return _resolve_gpu_opts(raw) + + +def validate_engine_args_for_command(args, available_commands): + """Engine/GPU CLI options are only valid for commands that support the selected engine.""" + available_commands = tuple(available_commands) + if args.command in available_commands: + return + + gpu_arg_names = [ + "engine", + "gpu_device", + "gpu_dtype", + "gpu_allow_tf32", + "gpu_compile", + "gpu_eps", + "gpu_check_every", + "gpu_compile_block", + ] + if any(getattr(args, name) is not None for name in gpu_arg_names): + commands = ", ".join(available_commands) + raise ValueError(f"NMF engine/GPU options are only valid with: {commands}") + + +def configure_nmf_engine(cnmf_obj, engine="cpu", gpu_kwargs=None): + """Configure a cNMF instance with an optional GPU NMF adapter.""" + if engine not in ("cpu", "gpu"): + raise ValueError("engine must be 'cpu' or 'gpu'") + if engine == "cpu": + return cnmf_obj + + def _gpu_nmf(X, nmf_kwargs): + nmf_kwargs = dict(nmf_kwargs) + nmf_kwargs["engine"] = "gpu" + nmf_kwargs["gpu"] = gpu_kwargs or {} + return _nmf_gpu(cnmf_obj, X, nmf_kwargs) + + cnmf_obj._nmf = _gpu_nmf + return cnmf_obj + + def _resolve_gpu_opts(gpu_kwargs): """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict.""" raw = dict(gpu_kwargs or {}) @@ -359,6 +429,9 @@ def factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs=None): return _to_nmf_output(H, W) -def _nmf_gpu(self, X, nmf_kwargs, gpu_kwargs=None): +def _nmf_gpu(self, X, nmf_kwargs): """cNMF adapter: same core NMF kernel, with `self` ignored for monkeypatch compatibility.""" + nmf_kwargs = dict(nmf_kwargs) + gpu_kwargs = nmf_kwargs.pop("gpu", None) + nmf_kwargs.pop("engine", None) return factorize_nmf_gpu(X, nmf_kwargs, gpu_kwargs) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 664b4bd..0c1b9d4 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -133,7 +133,7 @@ def test_factorize_nmf_gpu_fp32_compute_still_returns_float64_numpy_outputs(kern def test_nmf_gpu_adapter_ignores_self_and_delegates_to_factorize_nmf_gpu(kernel, monkeypatch): - """Verify the cNMF monkeypatch adapter forwards args and ignores its bound `self`.""" + """Verify the cNMF adapter extracts embedded GPU args and ignores its bound `self`.""" calls = [] sentinel = (object(), object()) @@ -143,13 +143,14 @@ def fake_factorize(X, nmf_kwargs, gpu_kwargs=None): monkeypatch.setattr(kernel, "factorize_nmf_gpu", fake_factorize) X = np.ones((3, 2)) - nmf_kwargs = {"n_components": 1} gpu_kwargs = {"device": "cpu"} + nmf_kwargs = {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1} - result = kernel._nmf_gpu(object(), X, nmf_kwargs, gpu_kwargs) + result = kernel._nmf_gpu(object(), X, nmf_kwargs) assert result is sentinel - assert calls == [(X, nmf_kwargs, gpu_kwargs)] + assert calls == [(X, {"n_components": 1}, gpu_kwargs)] + assert nmf_kwargs == {"engine": "gpu", "gpu": gpu_kwargs, "n_components": 1} # ---------------------------------------------------------------------