From 27b4bd69e61a85228979d01e57a8ae96bcba97ee Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:12:40 +0000 Subject: [PATCH 1/2] enable consensus with gpu by adding tiled matrix calculation feature --- src/cnmf/cnmf.py | 1 + src/cnmf/gpunmf/solver_cd.py | 337 +++++++++++++++++++++------- src/cnmf/gpunmf/solver_cd_triton.py | 21 +- src/cnmf/gpunmf/utils.py | 92 +++++++- 4 files changed, 356 insertions(+), 95 deletions(-) diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index 58f01bd..2d56896 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -1271,6 +1271,7 @@ def main(): parser.add_argument("--gpu-check-every", type=int, help="[factorize,consensus,gpu] Eager-mode convergence check interval") parser.add_argument("--gpu-compile-block", type=int, help="[factorize,consensus,gpu] Number of MU iterations per compiled block") parser.add_argument("--gpu-batch", type=int, help="[factorize] Replicates run per GPU solver launch; 1 = single-replicate") + parser.add_argument("--gpu-row-tiling-ratio", type=float, help="[consensus,gpu] Force fixed-H CD row tiling with a ratio in (0, 1]; omitted uses automatic VRAM sizing") args = parser.parse_args() diff --git a/src/cnmf/gpunmf/solver_cd.py b/src/cnmf/gpunmf/solver_cd.py index 7fd234e..7164d8f 100644 --- a/src/cnmf/gpunmf/solver_cd.py +++ b/src/cnmf/gpunmf/solver_cd.py @@ -4,6 +4,7 @@ from . import utils + # --------------------------------------------------------------------- # CD solver (sklearn-compatible Fast-HALS; batch-aware) # --------------------------------------------------------------------- @@ -71,10 +72,14 @@ def _validate_cd_runtime(torch, rc, nmf_kwargs): raise ValueError("shuffle must be a boolean for solver='cd'") -def _hals_sweep_torch(factor, gram, cross, permutation, active): +def _hals_sweep_torch( + factor, gram, cross, permutation, active, return_row_violation=False +): """Apply one literal torch port of sklearn's serial CD coordinate sweep.""" replicates, components, rows = factor.shape - violation = factor.new_zeros(replicates) + violation = factor.new_zeros( + (replicates, rows) if return_row_violation else replicates + ) zero = factor.new_zeros(()) cyclic = permutation is None @@ -105,7 +110,16 @@ def _hals_sweep_torch(factor, gram, cross, permutation, active): projected_gradient = gradient.where( old_value != 0, zero.minimum(gradient) ) - violation = violation + projected_gradient.abs().sum(dim=1) * active + if return_row_violation: + violation = ( + violation + + projected_gradient.abs() * active[:, None] + ) + else: + violation = ( + violation + + projected_gradient.abs().sum(dim=1) * active + ) if cyclic: hessian = gram[:, component, component, None] @@ -141,7 +155,9 @@ def _get_hals_cuda_backend(): return _HALS_CUDA_BACKEND -def _hals_sweep(factor, gram, cross, permutation, active): +def _hals_sweep( + factor, gram, cross, permutation, active, return_row_violation=False +): """Use the fused CUDA sweep when available, otherwise use torch.""" if factor.is_cuda: hals_sweep_cuda = _get_hals_cuda_backend() @@ -153,8 +169,22 @@ def _hals_sweep(factor, gram, cross, permutation, active): (factor.shape[0], 1), ) ).long() - return hals_sweep_cuda(factor, gram, cross, permutation, active) - return _hals_sweep_torch(factor, gram, cross, permutation, active) + return hals_sweep_cuda( + factor, + gram, + cross, + permutation, + active, + return_row_violation=return_row_violation, + ) + return _hals_sweep_torch( + factor, + gram, + cross, + permutation, + active, + return_row_violation=return_row_violation, + ) def _regularize_cd_products(gram, cross, l1_reg, l2_reg): @@ -187,21 +217,16 @@ def _batch_invariant_cd_products(factor, data): return gram, cross -def _fit_cd( - torch, - Xg, - Wt, - H, - max_iter, - tol, - update_h, - regularization, - shuffle, - seeds, - tf32, - device, -): +def _fit_cd(torch, rc, Xg, Wt, H, update_h, seeds, nmf_kwargs): """Run batched sklearn-compatible Fast-HALS to projected-gradient convergence.""" + max_iter = rc.max_iter + tol = rc.tol + regularization = _cd_regularization( + nmf_kwargs, Xg.shape[0], Xg.shape[1] + ) + shuffle = bool(nmf_kwargs.get("shuffle", False)) + tf32 = utils._want_tf32(torch, rc) + device = rc.device if Xg.dtype != Wt.dtype or Xg.dtype != H.dtype: raise RuntimeError( "NMF runtime tensors must share dtype; got " @@ -274,85 +299,227 @@ def next_permutation(): return Wt, H, n_iter -def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): - """Run full or fixed-H sklearn-compatible CD replicates.""" - torch = utils._loud_import_torch() - seeds = utils._normalize_seeds(seeds) - rc = utils._gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) - _validate_cd_runtime(torch, rc, nmf_kwargs) +def _fit_cd_fixed_h_tiled( + torch, rc, Xcompute, H, row_batch, seeds, nmf_kwargs +): + """Update W in row tiles from one shared set of fixed-H products.""" + max_iter = rc.max_iter + tol = rc.tol + regularization = _cd_regularization( + nmf_kwargs, Xcompute.shape[0], Xcompute.shape[1] + ) + shuffle = bool(nmf_kwargs.get("shuffle", False)) + tf32 = utils._want_tf32(torch, rc) + device = rc.device + replicates, components, _ = H.shape + n_rows = Xcompute.shape[0] + Wt_host = np.zeros( + (replicates, components, n_rows), dtype=Xcompute.dtype + ) + active = torch.ones(replicates, dtype=torch.bool, device=device) + n_iter = torch.zeros(replicates, dtype=torch.int64, device=device) + violation_init = H.new_zeros(replicates) + l1_w, l2_w, _l1_h, _l2_h = regularization + + # H is fixed, so H @ H.T and H @ X.T are invariant across iterations. + # Build them once on the host for both tiled and untiled execution. This + # makes the arithmetic inputs independent of the requested row partition + # and avoids transferring/reducing the multi-gigabyte X every iteration. + Hcompute = H.cpu().numpy() + gram_host = np.empty( + (replicates, components, components), dtype=Xcompute.dtype + ) + cross_host = np.empty( + (replicates, components, n_rows), dtype=Xcompute.dtype + ) + for replicate in range(replicates): + gram_host[replicate] = Hcompute[replicate] @ Hcompute[replicate].T + cross_host[replicate] = Hcompute[replicate] @ Xcompute.T + if l1_w != 0.0: + cross_host -= l1_w + gram = torch.as_tensor(gram_host, dtype=H.dtype, device=device) + if l2_w != 0.0: + gram.diagonal(dim1=-2, dim2=-1).add_(l2_w) + + if shuffle: + try: + from sklearn.utils import check_random_state + except ModuleNotFoundError as exc: + raise RuntimeError("scikit-learn is required for shuffled CD") from exc + rngs = [check_random_state(seed) for seed in seeds] + cyclic_permutation = None + else: + rngs = None + cyclic_permutation = ( + torch.arange(components, dtype=torch.int64, device=device) + .expand(replicates, -1) + .contiguous() + if H.is_cuda and _get_hals_cuda_backend() is not None + else None + ) + + def next_permutation(): + if rngs is None: + return cyclic_permutation + values = np.stack( + [rng.permutation(components) for rng in rngs], axis=0 + ).astype(np.int64, copy=False) + return torch.as_tensor(values, dtype=torch.int64, device=device) + + with torch.no_grad(), utils._cuda_tf32(torch, tf32, device): + for iteration in range(1, max_iter + 1): + permutation = next_permutation() + row_violation_host = np.empty( + (replicates, n_rows), dtype=Xcompute.dtype + ) + + for start in range(0, n_rows, row_batch): + stop = min(start + row_batch, n_rows) + Wt = torch.as_tensor( + np.ascontiguousarray(Wt_host[:, :, start:stop]), + dtype=H.dtype, + device=device, + ) + cross = torch.as_tensor( + np.ascontiguousarray(cross_host[:, :, start:stop]), + dtype=H.dtype, + device=device, + ) + row_violation = _hals_sweep( + Wt, + gram, + cross, + permutation, + active, + return_row_violation=True, + ) + row_violation_host[:, start:stop] = ( + row_violation.cpu().numpy() + ) + Wt_host[:, :, start:stop] = Wt.cpu().numpy() + + # Reduce once in global row order so the convergence decision is + # independent of where the requested row tiles begin and end. + violation = torch.as_tensor( + row_violation_host, dtype=H.dtype, device=device + ).sum(dim=1) + + n_iter = n_iter.new_full((), iteration).where(active, n_iter) + if iteration == 1: + violation_init = violation.clone() + + zero_init = violation_init == 0 + denominator = violation_init.where( + ~zero_init, violation_init.new_ones(()) + ) + converged = active & ( + zero_init | ((violation / denominator) <= tol) + ) + active = active & ~converged + if not bool(active.any()): + break + + return Wt_host, n_iter + + +# --------------------------------------------------------------------- +# CD solver entry points +# --------------------------------------------------------------------- + + +def _nmf_gpu_cd_update_h(torch, rc, seeds, nmf_kwargs): + """Run CD with both W and H updated.""" host_dtype = _numpy_staging_dtype(torch, rc.dtype) Xcompute = np.ascontiguousarray(rc.Xnp, dtype=host_dtype) Xg = torch.as_tensor(Xcompute, dtype=rc.dtype, device=rc.device) replicates = len(seeds) - update_h = nmf_kwargs.get("update_H", True) is not False - - if update_h: - init = nmf_kwargs.get("init") - if init == "custom": - W0 = _to_checked_custom_factor( - nmf_kwargs.get("W"), - (Xcompute.shape[0], rc.k), - "W", - host_dtype, - ) - H0 = _to_checked_custom_factor( - nmf_kwargs.get("H"), - (rc.k, Xcompute.shape[1]), - "H", - host_dtype, - ) - Wt0 = np.repeat(W0.T[None, :, :], replicates, axis=0) - Hs0 = np.repeat(H0[None, :, :], replicates, axis=0) - else: - Wt0 = np.empty( - (replicates, rc.k, Xcompute.shape[0]), dtype=host_dtype - ) - Hs0 = np.empty( - (replicates, rc.k, Xcompute.shape[1]), dtype=host_dtype - ) - for replicate, seed in enumerate(seeds): - W0, H0 = utils._init_wh(Xcompute, rc.k, seed, init) - Wt0[replicate] = W0.T - Hs0[replicate] = H0 - Wt = torch.as_tensor(Wt0, dtype=rc.dtype, device=rc.device) - H = torch.as_tensor(Hs0, dtype=rc.dtype, device=rc.device) - else: - H0 = utils._to_checked_fixed_h( - nmf_kwargs.get("H"), rc.k, Xcompute.shape[1] + + init = nmf_kwargs.get("init") + if init == "custom": + W0 = _to_checked_custom_factor( + nmf_kwargs.get("W"), + (Xcompute.shape[0], rc.k), + "W", + host_dtype, ) - H0 = np.ascontiguousarray(H0, dtype=host_dtype) - # sklearn CD ignores supplied W and initializes fixed-H usages to zero. - Wt = torch.zeros( - (replicates, rc.k, Xcompute.shape[0]), - dtype=rc.dtype, - device=rc.device, + H0 = _to_checked_custom_factor( + nmf_kwargs.get("H"), + (rc.k, Xcompute.shape[1]), + "H", + host_dtype, ) - H = ( - torch.as_tensor(H0, dtype=rc.dtype, device=rc.device) - .unsqueeze(0) - .expand(replicates, -1, -1) - .contiguous() + Wt0 = np.repeat(W0.T[None, :, :], replicates, axis=0) + Hs0 = np.repeat(H0[None, :, :], replicates, axis=0) + else: + Wt0 = np.empty( + (replicates, rc.k, Xcompute.shape[0]), dtype=host_dtype + ) + Hs0 = np.empty( + (replicates, rc.k, Xcompute.shape[1]), dtype=host_dtype ) + for replicate, seed in enumerate(seeds): + W0, H0 = utils._init_wh(Xcompute, rc.k, seed, init) + Wt0[replicate] = W0.T + Hs0[replicate] = H0 - regularization = _cd_regularization( - nmf_kwargs, Xcompute.shape[0], Xcompute.shape[1] + Wt = torch.as_tensor(Wt0, dtype=rc.dtype, device=rc.device) + H = torch.as_tensor(Hs0, dtype=rc.dtype, device=rc.device) + Wt, H, _ = _fit_cd(torch, rc, Xg, Wt, H, True, seeds, nmf_kwargs) + + Hc = H.cpu().double().numpy() + Wc = Wt.transpose(-2, -1).cpu().double().numpy() + return [(Hc[r], Wc[r]) for r in range(replicates)] + + +def _nmf_gpu_cd_fixed_h(torch, rc, seeds, nmf_kwargs): + """Run CD with H fixed and only W updated.""" + host_dtype = _numpy_staging_dtype(torch, rc.dtype) + Xcompute = np.ascontiguousarray(rc.Xnp, dtype=host_dtype) + replicates = len(seeds) + + H0 = utils._to_checked_fixed_h( + nmf_kwargs.get("H"), rc.k, Xcompute.shape[1] + ) + H0 = np.ascontiguousarray(H0, dtype=host_dtype) + H = ( + torch.as_tensor(H0, dtype=rc.dtype, device=rc.device) + .unsqueeze(0) + .expand(replicates, -1, -1) + .contiguous() ) - Wt, H, _ = _fit_cd( + + row_batch, tiling_ratio = utils.resolve_row_batch( + torch, + rc.device, + Xcompute, + replicates, + rc.k, + configured_ratio=rc.opt["row_tiling_ratio"], + ) + + Wt_host, _ = _fit_cd_fixed_h_tiled( torch, - Xg, - Wt, + rc, + Xcompute, H, - rc.max_iter, - rc.tol, - update_h, - regularization, - bool(nmf_kwargs.get("shuffle", False)), + row_batch, seeds, - utils._want_tf32(torch, rc), - rc.device, + nmf_kwargs, ) + Wc = Wt_host.transpose(0, 2, 1).astype(np.float64, copy=False) Hc = H.cpu().double().numpy() - Wc = Wt.transpose(-2, -1).cpu().double().numpy() return [(Hc[r], Wc[r]) for r in range(replicates)] + + +def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): + """Dispatch full or fixed-H sklearn-compatible CD replicates.""" + torch = utils._loud_import_torch() + seeds = utils._normalize_seeds(seeds) + rc = utils._gpu_setup(torch, X, nmf_kwargs, gpu_kwargs) + _validate_cd_runtime(torch, rc, nmf_kwargs) + + if nmf_kwargs.get("update_H", True) is not False: + return _nmf_gpu_cd_update_h(torch, rc, seeds, nmf_kwargs) + return _nmf_gpu_cd_fixed_h(torch, rc, seeds, nmf_kwargs) diff --git a/src/cnmf/gpunmf/solver_cd_triton.py b/src/cnmf/gpunmf/solver_cd_triton.py index d946113..10b159f 100644 --- a/src/cnmf/gpunmf/solver_cd_triton.py +++ b/src/cnmf/gpunmf/solver_cd_triton.py @@ -18,7 +18,7 @@ def _hals_sweep_kernel( cross, permutation, active, - block_violation, + row_violation, M: tl.constexpr, K: tl.constexpr, N_BLOCKS: tl.constexpr, @@ -87,14 +87,15 @@ def _hals_sweep_kernel( ) tl.store( - block_violation - + replicate_offset * N_BLOCKS - + row_block.to(tl.int64), - tl.sum(violation, axis=0), + row_violation + replicate_offset * M + rows_offset, + violation, + mask=row_mask, ) -def hals_sweep_cuda(factor, gram, cross, permutation, active): +def hals_sweep_cuda( + factor, gram, cross, permutation, active, return_row_violation=False +): """Update one factor in place and return violation per replicate.""" if not factor.is_cuda: raise ValueError("hals_sweep_cuda requires CUDA tensors") @@ -135,7 +136,7 @@ def hals_sweep_cuda(factor, gram, cross, permutation, active): block = 128 n_blocks = triton.cdiv(rows, block) - block_violation = factor.new_empty((replicates, n_blocks)) + row_violation = factor.new_empty((replicates, rows)) grid = (replicates * n_blocks,) _hals_sweep_kernel[grid]( factor, @@ -143,7 +144,7 @@ def hals_sweep_cuda(factor, gram, cross, permutation, active): cross, permutation, active, - block_violation, + row_violation, M=rows, K=components, N_BLOCKS=n_blocks, @@ -151,4 +152,6 @@ def hals_sweep_cuda(factor, gram, cross, permutation, active): IS_FP64=factor.element_size() == 8, num_warps=4, ) - return block_violation.sum(dim=1) + if return_row_violation: + return row_violation + return row_violation.sum(dim=1) diff --git a/src/cnmf/gpunmf/utils.py b/src/cnmf/gpunmf/utils.py index 319d336..09fc515 100644 --- a/src/cnmf/gpunmf/utils.py +++ b/src/cnmf/gpunmf/utils.py @@ -12,7 +12,7 @@ "max_iter": 1000, "tol": 1e-4, "init": "random", - "solver": "mu", + "solver": "cd", } @@ -25,6 +25,7 @@ "check_every": 10, "compile_block": 1, "batch": 1, + "row_tiling_ratio": None, } @@ -38,6 +39,7 @@ "gpu_check_every", "gpu_compile_block", "gpu_batch", + "gpu_row_tiling_ratio", ) @@ -92,6 +94,7 @@ def gpu_kwargs_from_args(args): "check_every": args.gpu_check_every, "compile_block": args.gpu_compile_block, "batch": args.gpu_batch, + "row_tiling_ratio": args.gpu_row_tiling_ratio, } if args.engine != "gpu": if any(value is not None for value in raw.values()): @@ -115,6 +118,9 @@ def parse_typed(value, default, cast, normalize=None): def parse_positive_int(value, default): return max(1, parse_typed(value, default, int)) + def parse_optional_tiling_ratio(value): + return None if value is None else _validate_row_tiling_ratio(value) + return dict( device = parse_typed(raw.get("device"), DEFAULT_GPU["device"], str, str.lower), dtype = parse_typed(raw.get("dtype"), DEFAULT_GPU["dtype"], str, str.lower), @@ -124,9 +130,93 @@ def parse_positive_int(value, default): check_every = parse_positive_int(raw.get("check_every"), DEFAULT_GPU["check_every"]), compile_block = parse_positive_int(raw.get("compile_block"), DEFAULT_GPU["compile_block"]), batch = parse_positive_int(raw.get("batch"), DEFAULT_GPU["batch"]), + row_tiling_ratio = parse_optional_tiling_ratio(raw.get("row_tiling_ratio")), ) +# --------------------------------------------------------------------- +# GPU row tiling +# --------------------------------------------------------------------- + + +_MIB = 1 << 20 + + +def _validate_row_tiling_ratio(value): + """Return a finite configured row ratio in the interval (0, 1].""" + try: + ratio = float(value) + except (TypeError, ValueError): + raise ValueError( + "gpu row tiling ratio must be a finite number in (0, 1]" + ) from None + if not np.isfinite(ratio) or not 0 < ratio <= 1: + raise ValueError( + "gpu row tiling ratio must be a finite number in (0, 1]" + ) + return ratio + + +def _check_row_batch(row_batch, n_rows): + """Validate and cap a resolved row batch.""" + if row_batch < 1 or n_rows < 1: + raise ValueError("gpu row batch and row count must be at least 1") + return min(n_rows, row_batch) + + +def resolve_row_batch( + torch, + device, + Xcompute, + replicates, + n_components, + *, + configured_ratio=None, + reserve_bytes=512 * _MIB, + reserve_fraction=0.10, +): + """Resolve a configured row ratio or automatically size one from VRAM.""" + n_rows, _n_features = Xcompute.shape + if n_rows < 1: + raise ValueError("fixed-H input must contain at least one row") + + if configured_ratio is not None: + tiling_ratio = _validate_row_tiling_ratio(configured_ratio) + row_batch = max(1, int(n_rows * tiling_ratio)) + return _check_row_batch(row_batch, n_rows), tiling_ratio + + tiling_ratio = 1.0 + if str(device).split(":", 1)[0] == "cuda": + if reserve_bytes < 0 or not 0 <= reserve_fraction < 1: + raise ValueError("invalid VRAM reserve") + free_vram_bytes, _ = torch.cuda.mem_get_info(device) + factor_bytes = ( + replicates + * n_components + * n_rows + * Xcompute.dtype.itemsize + ) + required_bytes = Xcompute.nbytes + 2 * factor_bytes + reserve = max( + reserve_bytes, + int(free_vram_bytes * reserve_fraction), + ) + usable_vram_bytes = max(0, free_vram_bytes - reserve) + if required_bytes > usable_vram_bytes: + bytes_per_row = (required_bytes + n_rows - 1) // n_rows + if usable_vram_bytes < bytes_per_row: + raise MemoryError( + "GPU fixed-H CD does not have enough VRAM for one row" + ) + portions = ( + required_bytes + usable_vram_bytes - 1 + ) // usable_vram_bytes + tiling_ratio = 1 / portions + + row_batch = max(1, int(n_rows * tiling_ratio)) + return _check_row_batch(row_batch, n_rows), tiling_ratio + + # --------------------------------------------------------------------- # Runtime backend selection (device / dtype / TF32) # --------------------------------------------------------------------- From 93a852dff6f4ac9f1cda989f1a95a41fc9c984fe Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:13:22 +0000 Subject: [PATCH 2/2] add tiled refit usage with cd solver parity tests --- tests/test_nmf_gpu.py | 126 +++++++++++++++++++++++++++++++++++++++--- tests/test_prepare.py | 28 ++++++++++ tests/test_tiling.py | 88 +++++++++++++++++++++++++++++ 3 files changed, 233 insertions(+), 9 deletions(-) create mode 100644 tests/test_tiling.py diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 72761c2..231c3f6 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -679,8 +679,9 @@ def test_cd_batched_seeds_match_independent_runs( @pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +@pytest.mark.parametrize("row_tiling_ratio", [None, 0.5]) def test_sklearn_cd_fixed_h_matches_batched_gpu_refit( - kernel, dtype_name, np_dtype, rtol, atol + kernel, dtype_name, np_dtype, rtol, atol, row_tiling_ratio ): """Fixed-H CD must keep H and use sklearn's exact-zero W start.""" require_nmf_runtime() @@ -700,6 +701,7 @@ def test_sklearn_cd_fixed_h_matches_batched_gpu_refit( "dtype": dtype_name, "allow_tf32": False, "compile": False, + "row_tiling_ratio": row_tiling_ratio, } actual = kernel.solver_cd._nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs) @@ -715,6 +717,88 @@ def test_sklearn_cd_fixed_h_matches_batched_gpu_refit( ) +@pytest.mark.parametrize("dtype_name,np_dtype,rtol,atol", CD_PARITY_CASES) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("max_iter", [1, 9]) +@pytest.mark.parametrize("row_tiling_ratio", [0.25, 0.5]) +def test_cd_fixed_h_row_tiling_matches_untiled( + kernel, dtype_name, np_dtype, rtol, atol, device, max_iter, + row_tiling_ratio, monkeypatch, +): + """Tiled and untiled fixed-H refits must be bitwise identical.""" + torch = require_nmf_runtime() + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + rng = np.random.default_rng(137) + X = rng.random((257, 11), dtype=np_dtype) + np_dtype(0.1) + fixed_H = rng.random((4, 11), dtype=np_dtype) + np_dtype(0.1) + seeds = [7, 43] + nmf_kwargs = _cd_nmf_kwargs( + 4, + seed=0, + max_iter=max_iter, + update_H=False, + H=fixed_H, + ) + gpu_kwargs = { + "device": device, + "dtype": dtype_name, + "allow_tf32": False, + "compile": False, + } + + convergence_iterations = [] + sweep_rows = [] + original_fit = kernel.solver_cd._fit_cd_fixed_h_tiled + original_sweep = kernel.solver_cd._hals_sweep + + def capture_iterations(*args, **kwargs): + result = original_fit(*args, **kwargs) + convergence_iterations.append(result[1].cpu().numpy().copy()) + return result + + monkeypatch.setattr( + kernel.solver_cd, + "_fit_cd_fixed_h_tiled", + capture_iterations, + ) + + def capture_sweep(factor, *args, **kwargs): + sweep_rows.append(factor.shape[-1]) + return original_sweep(factor, *args, **kwargs) + + monkeypatch.setattr(kernel.solver_cd, "_hals_sweep", capture_sweep) + + untiled = kernel.solver_cd._nmf_gpu_cd( + X, + seeds, + nmf_kwargs, + dict(gpu_kwargs, row_tiling_ratio=1.0), + ) + untiled_sweep_count = len(sweep_rows) + tiled = kernel.solver_cd._nmf_gpu_cd( + X, + seeds, + nmf_kwargs, + dict(gpu_kwargs, row_tiling_ratio=row_tiling_ratio), + ) + + expected_row_batch = int(X.shape[0] * row_tiling_ratio) + tiled_sweep_rows = sweep_rows[untiled_sweep_count:] + assert tiled_sweep_rows[0] == expected_row_batch + assert max(tiled_sweep_rows) <= expected_row_batch + + for (untiled_H, untiled_W), (tiled_H, tiled_W) in zip( + untiled, tiled + ): + np.testing.assert_array_equal(tiled_H, untiled_H) + np.testing.assert_array_equal(tiled_W, untiled_W) + np.testing.assert_array_equal( + convergence_iterations[1], convergence_iterations[0] + ) + + + def test_sklearn_cd_custom_init_matches_gpu_kernel(kernel): """Custom W/H should be consumed exactly as sklearn consumes them.""" require_nmf_runtime() @@ -1087,6 +1171,7 @@ def test_resolve_gpu_opts_dict_values_override_defaults(kernel): "check_every": 7, "compile_block": 9, "batch": 1, # not overridden here -> default + "row_tiling_ratio": None, } @@ -1460,24 +1545,24 @@ def test_nmf_gpu_batch_rejects_unknown_solver(kernel): ) -def test_nmf_gpu_batch_defaults_to_mu(kernel, monkeypatch): - """Missing solver configuration should route to MU.""" +def test_nmf_gpu_batch_defaults_to_cd(kernel, monkeypatch): + """Missing solver configuration should route to CD.""" calls = [] - def fake_mu(X, seeds, nmf_kwargs, gpu_kwargs=None): + def fake_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): calls.append((X, seeds, nmf_kwargs, gpu_kwargs)) - return ["mu-result"] + return ["cd-result"] - monkeypatch.setitem(kernel._GPU_SOLVERS, "mu", fake_mu) + monkeypatch.setitem(kernel._GPU_SOLVERS, "cd", fake_cd) X = np.ones((3, 2)) seeds = [7] nmf_kwargs = {"n_components": 1} gpu_kwargs = {"device": "cpu"} - assert kernel.utils.DEFAULT_NMF["solver"] == "mu" + assert kernel.utils.DEFAULT_NMF["solver"] == "cd" assert kernel._nmf_gpu_batch( X, seeds, nmf_kwargs, gpu_kwargs - ) == ["mu-result"] + ) == ["cd-result"] assert len(calls) == 1 actual_X, actual_seeds, actual_kwargs, actual_gpu_kwargs = calls[0] assert actual_X is X @@ -1693,6 +1778,7 @@ def _engine_args(**overrides): "gpu_check_every": None, "gpu_compile_block": None, "gpu_batch": None, + "gpu_row_tiling_ratio": None, } values.update(overrides) return SimpleNamespace(**values) @@ -1713,6 +1799,7 @@ def run_nmf_gpu(kernel, X, nmf_kwargs, gpu_kwargs=None): gpu_check_every=gpu_kwargs.get("check_every"), gpu_compile_block=gpu_kwargs.get("compile_block"), gpu_batch=gpu_kwargs.get("batch"), + gpu_row_tiling_ratio=gpu_kwargs.get("row_tiling_ratio"), ) return kernel._nmf_gpu(args, X, nmf_kwargs) @@ -1740,6 +1827,25 @@ def test_gpu_kwargs_from_args_carries_batch_and_fills_its_default(kernel): assert kernel.utils.gpu_kwargs_from_args(default_args)["batch"] == 1 + + + +@pytest.mark.parametrize("ratio", ["0.5", 1.0]) +def test_gpu_kwargs_from_args_accepts_valid_row_tiling_ratio(kernel, ratio): + actual = kernel.utils.gpu_kwargs_from_args( + _engine_args(engine="gpu", gpu_row_tiling_ratio=ratio) + ) + assert actual["row_tiling_ratio"] == float(ratio) + + +@pytest.mark.parametrize("ratio", [0, -0.1, 1.1, np.nan, np.inf, "bad"]) +def test_gpu_kwargs_from_args_rejects_invalid_row_tiling_ratio(kernel, ratio): + with pytest.raises(ValueError, match="row tiling ratio"): + kernel.utils.gpu_kwargs_from_args( + _engine_args(engine="gpu", gpu_row_tiling_ratio=ratio) + ) + + # --------------------------------------------------------------------- # CLI parsing, engine wiring, and fixed-H consensus refit (integration) # --------------------------------------------------------------------- @@ -1749,7 +1855,8 @@ def test_engine_args_default_to_none_until_user_selects_an_engine(kernel): assert args.engine is None for name in ("gpu_device", "gpu_dtype", "gpu_allow_tf32", "gpu_compile", - "gpu_eps", "gpu_check_every", "gpu_compile_block", "gpu_batch"): + "gpu_eps", "gpu_check_every", "gpu_compile_block", "gpu_batch", + "gpu_row_tiling_ratio"): assert getattr(args, name) is None, f"{name} should default to None" @@ -1792,6 +1899,7 @@ def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel): "check_every": 5, "compile_block": 100, "batch": 1, # default added; not set on the CLI here + "row_tiling_ratio": None, } diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 2d6eb32..1240966 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -41,6 +41,7 @@ def configure_gpu(cnmf_obj, *, command="factorize", solver="mu", gpu_check_every=gpu_kwargs.get("check_every"), gpu_compile_block=gpu_kwargs.get("compile_block"), gpu_batch=gpu_kwargs.get("batch"), + gpu_row_tiling_ratio=gpu_kwargs.get("row_tiling_ratio"), ) configured = gpunmf.configure_nmf_engine( lambda output_dir, name: cnmf_obj, @@ -509,6 +510,28 @@ def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): assert spectra.shape == (2, 3) +@pytest.mark.parametrize("row_tiling_ratio", [None, 0.5]) +def test_gpu_consensus_keeps_original_method_and_binds_args_to_nmf( + mock_cnmf, row_tiling_ratio +): + """Consensus stays original while `_nmf` retains its parsed GPU arguments.""" + original_consensus = mock_cnmf.consensus + mock_cnmf, args = configure_gpu( + mock_cnmf, + command="consensus", + solver="cd", + gpu_kwargs={ + "device": "cpu", + "row_tiling_ratio": row_tiling_ratio, + }, + ) + + assert mock_cnmf.consensus == original_consensus + assert mock_cnmf._nmf.func is gpunmf._nmf_gpu + assert mock_cnmf._nmf.args == (args,) + assert gpunmf.utils.gpu_kwargs_from_args(args)["row_tiling_ratio"] == row_tiling_ratio + + def test_consensus_gpu_engine_smoke_writes_expected_outputs(mock_cnmf, monkeypatch, tmp_path): """A tiny CPU-backed GPU-engine consensus run should write the expected consensus outputs.""" import cnmf.gpunmf as gpu_mod @@ -535,7 +558,10 @@ def test_consensus_gpu_engine_smoke_writes_expected_outputs(mock_cnmf, monkeypat ) save_df_to_npz(merged, mock_cnmf.paths["merged_spectra"] % 2) + refit_rows = [] + def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): + refit_rows.append(X_arg.shape[0]) return fake_gpu_nmf_output(X_arg, nmf_kwargs) monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) @@ -548,6 +574,8 @@ def fake_nmf_gpu(args, X_arg, nmf_kwargs, gpu_kwargs=None): mock_cnmf.consensus(k=2, density_threshold=2.0, local_neighborhood_size=0.5, show_clustering=False, refit_usage=False) + assert refit_rows + density = "2_0" expected_files = [ mock_cnmf.paths["consensus_spectra"] % (2, density), diff --git a/tests/test_tiling.py b/tests/test_tiling.py new file mode 100644 index 0000000..850a638 --- /dev/null +++ b/tests/test_tiling.py @@ -0,0 +1,88 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +from cnmf.gpunmf import utils + + +_GIB = 1 << 30 + + +def fake_torch(*, free_bytes=6 * _GIB, total_bytes=8 * _GIB): + calls = [] + cuda = SimpleNamespace( + mem_get_info=lambda device: calls.append(device) + or (free_bytes, total_bytes), + ) + return SimpleNamespace(cuda=cuda), calls + + +def test_resolve_row_batch_uses_all_rows_when_cpu_is_automatic(): + torch, calls = fake_torch() + X = np.zeros((100, 20), dtype=np.float32) + + row_batch, ratio = utils.resolve_row_batch(torch, "cpu", X, 2, 5) + + assert (row_batch, ratio) == (100, 1.0) + assert calls == [] + + +def test_resolve_row_batch_calculates_one_fitting_ratio_from_cuda_vram(): + torch, calls = fake_torch(free_bytes=5_000, total_bytes=20_000) + X = np.zeros((100, 20), dtype=np.float32) + + row_batch, ratio = utils.resolve_row_batch( + torch, + "cuda:1", + X, + 2, + 5, + reserve_bytes=0, + reserve_fraction=0, + ) + + assert calls == ["cuda:1"] + assert ratio == 0.25 + assert row_batch == 25 + + +def test_resolve_row_batch_uses_configured_ratio_without_vram_lookup(): + torch, calls = fake_torch() + X = np.zeros((101, 20), dtype=np.float64) + + row_batch, ratio = utils.resolve_row_batch( + torch, + "cuda", + X, + 2, + 5, + configured_ratio=0.5, + ) + + assert (row_batch, ratio) == (50, 0.5) + assert calls == [] + + +@pytest.mark.parametrize("ratio", [0, -0.1, 1.1, np.nan, np.inf, "bad"]) +def test_resolve_row_batch_rejects_invalid_configured_ratio(ratio): + torch, _calls = fake_torch() + X = np.zeros((10, 4), dtype=np.float32) + + with pytest.raises(ValueError, match="row tiling ratio"): + utils.resolve_row_batch( + torch, + "cpu", + X, + 1, + 2, + configured_ratio=ratio, + ) + + +def test_resolve_row_batch_rejects_input_without_rows(): + torch, _calls = fake_torch() + X = np.zeros((0, 4), dtype=np.float32) + + with pytest.raises(ValueError, match="at least one row"): + utils.resolve_row_batch(torch, "cpu", X, 1, 2)