From 8edaeba1d10b07039f3645dad2aa01d43cffebd1 Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:36:20 +0000 Subject: [PATCH] Add row tiling to the GPU CD solver Tiling off keeps the pre-tiling full-matrix computation. A row tiling ratio, or 0 for automatic VRAM sizing, streams X and W in row windows and computes each window's products at window size. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/cnmf/cnmf.py | 1 + src/cnmf/gpunmf/__init__.py | 3 +- src/cnmf/gpunmf/solver_cd.py | 368 +++++++++++++++++++---------- src/cnmf/gpunmf/tiling.py | 112 +++++++++ src/cnmf/gpunmf/utils.py | 37 ++- tests/test_cd_tiling_runtime.py | 126 ++++++++++ tests/test_nmf_gpu.py | 154 +++++++++++- tests/test_prepare.py | 1 + tests/test_synthetic_regression.py | 3 + tests/test_tiling.py | 161 +++++++++++++ 10 files changed, 830 insertions(+), 136 deletions(-) create mode 100644 src/cnmf/gpunmf/tiling.py create mode 100644 tests/test_cd_tiling_runtime.py create mode 100644 tests/test_tiling.py diff --git a/src/cnmf/cnmf.py b/src/cnmf/cnmf.py index b400b32..b7b4347 100755 --- a/src/cnmf/cnmf.py +++ b/src/cnmf/cnmf.py @@ -1291,6 +1291,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="[factorize,consensus,gpu] CD row tile ratio: omitted or 1 = no tiling; 0 = automatic VRAM sizing; (0, 1) = explicit ratio") args = parser.parse_args() diff --git a/src/cnmf/gpunmf/__init__.py b/src/cnmf/gpunmf/__init__.py index 48e8260..07d73b0 100644 --- a/src/cnmf/gpunmf/__init__.py +++ b/src/cnmf/gpunmf/__init__.py @@ -9,7 +9,7 @@ import numpy as np -from . import solver_cd, solver_mu, utils +from . import solver_cd, solver_mu, tiling, utils __all__ = [ @@ -18,6 +18,7 @@ "prepare_gpu", "solver_cd", "solver_mu", + "tiling", "utils", ] diff --git a/src/cnmf/gpunmf/solver_cd.py b/src/cnmf/gpunmf/solver_cd.py index 7fd234e..01e83fa 100644 --- a/src/cnmf/gpunmf/solver_cd.py +++ b/src/cnmf/gpunmf/solver_cd.py @@ -2,7 +2,7 @@ import numpy as np -from . import utils +from . import tiling, utils # --------------------------------------------------------------------- # CD solver (sklearn-compatible Fast-HALS; batch-aware) @@ -166,112 +166,260 @@ def _regularize_cd_products(gram, cross, l1_reg, l2_reg): return gram, cross -def _batch_invariant_cd_products(factor, data): - """Build CD products with one identical 2D GEMM path per replicate. +def _coordinate_orders(torch, seeds, components, shuffle, device): + """Return next_order(), giving each sweep's [replicate, component] order. - A single batched matmul may select a different fp32 CUDA reduction path as - the replicate count changes. Fast-HALS amplifies those small differences, - which can change projected-gradient stopping decisions. Keep the factor - state and HALS sweep batched, but compute each replicate's products through - the same two-dimensional matmul shape used by a batch of one. + Shuffled orders come from one sklearn RandomState per seed, one draw per + sweep. The cyclic order is None for the torch sweep's fast path and a + fixed index tensor for the CUDA sweep. """ - replicates, components, _ = factor.shape - gram = factor.new_empty((replicates, components, components)) - cross = factor.new_empty((replicates, components, data.shape[-1])) - for replicate in range(replicates): - replicate_factor = factor[replicate] - gram[replicate].copy_( - replicate_factor @ replicate_factor.transpose(-2, -1) - ) - cross[replicate].copy_(replicate_factor @ data) - return gram, cross + if shuffle: + from sklearn.utils import check_random_state + rngs = [check_random_state(seed) for seed in seeds] -def _fit_cd( - torch, - Xg, - Wt, - H, - max_iter, - tol, - update_h, - regularization, - shuffle, - seeds, - tf32, - device, -): - """Run batched sklearn-compatible Fast-HALS to projected-gradient convergence.""" - if Xg.dtype != Wt.dtype or Xg.dtype != H.dtype: - raise RuntimeError( - "NMF runtime tensors must share dtype; got " - f"{sorted(map(str, {Xg.dtype, Wt.dtype, H.dtype}))}" - ) + def next_order(): + orders = np.stack([rng.permutation(components) for rng in rngs]) + return torch.as_tensor(orders, dtype=torch.int64, device=device) - replicates, components, _ = Wt.shape - active = torch.ones(replicates, dtype=torch.bool, device=device) - n_iter = torch.zeros(replicates, dtype=torch.int64, device=device) - violation_init = Wt.new_zeros(replicates) - l1_w, l2_w, l1_h, l2_h = regularization + return next_order - 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 = ( + cyclic = None + if device.startswith("cuda") and _get_hals_cuda_backend() is not None: + cyclic = ( torch.arange(components, dtype=torch.int64, device=device) - .expand(replicates, -1) + .expand(len(seeds), -1) .contiguous() - if Xg.is_cuda and _get_hals_cuda_backend() is not None - else None ) + return lambda: cyclic - 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): - gram, cross = _batch_invariant_cd_products( - H, Xg.transpose(-2, -1) - ) - gram, cross = _regularize_cd_products(gram, cross, l1_w, l2_w) - violation = _hals_sweep( - Wt, gram, cross, next_permutation(), active - ) - if update_h: - gram, cross = _batch_invariant_cd_products(Wt, Xg) - gram, cross = _regularize_cd_products(gram, cross, l1_h, l2_h) - violation = violation + _hals_sweep( - H, gram, cross, next_permutation(), active +class _StopRule: + """sklearn's per-replicate CD stop: violation / first violation <= tol.""" + + def __init__(self, torch, replicates, tol, device): + self.tol = tol + self.active = torch.ones(replicates, dtype=torch.bool, device=device) + self.n_iter = torch.zeros(replicates, dtype=torch.int64, device=device) + self.violation_init = None + + def update(self, iteration, violation): + """Record one iteration; return True once every replicate has stopped.""" + self.n_iter = self.n_iter.new_full((), iteration).where( + self.active, self.n_iter + ) + if self.violation_init is None: + self.violation_init = violation.clone() + + zero_init = self.violation_init == 0 + denominator = self.violation_init.where( + ~zero_init, self.violation_init.new_ones(()) + ) + converged = self.active & ( + zero_init | ((violation / denominator) <= self.tol) + ) + self.active = self.active & ~converged + return not bool(self.active.any()) + + +def _window_budget( + update_h, replicates, components, n_features, itemsize, torch_sweep +): + """Return (fixed, per-row) device bytes for one fit; H is already allocated. + + A window of w rows needs fixed + w * per_row bytes. The resident full + matrix is the single window w = n_rows. + """ + R, k, m, s = replicates, components, n_features, itemsize + # Coordinate orders (int64), n_iter, active flags and violation totals. + fixed = R * (8 * k + 8 + 1 + 3 * s) + # The W-sweep Gram, plus one replicate's GEMM output before it is copied. + fixed += (R + 1) * k * k * s + # W and cross-product windows, plus one replicate's H·Xᵀ output. + per_row = (2 * R * k + k) * s + # Without Triton, the torch sweep holds about ten [replicate, row] + # temporaries plus an int64 gather index for shuffled orders. + sweep = 10 * R * s + 8 * R if torch_sweep else 0 + per_row += sweep + if update_h: + # The X window, the H-update Gram and cross accumulators, one + # replicate's W·X output, and the H sweep over n_features columns. + per_row += m * s + fixed += R * k * k * s + (R + 1) * k * m * s + sweep * m + else: + # The one-off cross-product pass holds an X window, every replicate's + # cross window and one replicate's GEMM output. + per_row = max(per_row, (m + R * k + k) * s) + return fixed, per_row + + +def _fixed_h_cross(torch, H, Xcompute, window_rows, l1_w): + """Return H·Xᵀ minus the L1 shift for every row, one window at a time.""" + replicates, components, _ = H.shape + n_rows = Xcompute.shape[0] + cross_host = np.empty( + (replicates, components, n_rows), dtype=Xcompute.dtype + ) + windows = tiling._RowWindows( + torch, H, n_rows, window_rows, {"X": (Xcompute, 0)} + ) + for start, stop, window in windows: + cross = H.new_empty((replicates, components, stop - start)) + for replicate in range(replicates): + cross[replicate].copy_(H[replicate] @ window["X"].T) + if l1_w != 0.0: + cross.sub_(l1_w) + cross_host[:, :, start:stop] = cross.cpu().numpy() + del cross + return cross_host + + +def _fit_cd_fixed_h( + torch, rc, Xcompute, Wt_host, H, window_rows, seeds, nmf_kwargs +): + """Update W against a fixed H, one row window at a time. + + H·Hᵀ and H·Xᵀ never change, so they are computed once, on the device, and + then travel with W's rows. One window keeps everything resident, which is + the pre-tiling arithmetic. + """ + replicates, components, _ = H.shape + n_rows = Xcompute.shape[0] + l1_w, l2_w, _, _ = _cd_regularization( + nmf_kwargs, n_rows, Xcompute.shape[1] + ) + next_order = _coordinate_orders( + torch, seeds, components, nmf_kwargs.get("shuffle", False), rc.device + ) + stop_rule = _StopRule(torch, replicates, rc.tol, rc.device) + gram = H.new_empty((replicates, components, components)) + tf32 = utils._want_tf32(torch, rc) + + with torch.no_grad(), utils._cuda_tf32(torch, tf32, rc.device): + for replicate in range(replicates): + gram[replicate].copy_(H[replicate] @ H[replicate].T) + if l2_w != 0.0: + gram.diagonal(dim1=-2, dim2=-1).add_(l2_w) + cross_host = _fixed_h_cross(torch, H, Xcompute, window_rows, l1_w) + windows = tiling._RowWindows( + torch, H, n_rows, window_rows, + {"Wt": (Wt_host, 2), "cross": (cross_host, 2)}, writable=("Wt",), + ) + + for iteration in range(1, rc.max_iter + 1): + order = next_order() + violation = None + for start, stop, window in windows: + window_violation = _hals_sweep( + window["Wt"], gram, window["cross"], order, stop_rule.active + ) + violation = ( + window_violation if violation is None + else violation + window_violation ) + if stop_rule.update(iteration, violation): + break - n_iter = n_iter.new_full((), iteration).where(active, n_iter) - if iteration == 1: - violation_init = violation.clone() + return windows.host("Wt"), H, stop_rule.n_iter - zero_init = violation_init == 0 - denominator = violation_init.where( - ~zero_init, violation_init.new_ones(()) - ) - converged = active & ( - zero_init | ((violation / denominator) <= tol) + +def _fit_cd_factorize( + torch, rc, Xcompute, Wt_host, H, window_rows, seeds, nmf_kwargs +): + """Update W and H, streaming X and W in row windows when they do not fit. + + Products are computed at window size and summed over windows for the H + update. One window keeps X and W resident and is exactly the pre-tiling + full-matrix arithmetic. Smaller windows split the sums over cells, so + their results agree numerically rather than bit for bit. + """ + replicates, components, n_features = H.shape + n_rows = Xcompute.shape[0] + l1_w, l2_w, l1_h, l2_h = _cd_regularization( + nmf_kwargs, n_rows, n_features + ) + next_order = _coordinate_orders( + torch, seeds, components, nmf_kwargs.get("shuffle", False), rc.device + ) + stop_rule = _StopRule(torch, replicates, rc.tol, rc.device) + gram_w = H.new_empty((replicates, components, components)) + gram_h = H.new_empty((replicates, components, components)) + cross_h = H.new_empty(H.shape) + tf32 = utils._want_tf32(torch, rc) + windows = tiling._RowWindows( + torch, H, n_rows, window_rows, + {"X": (Xcompute, 0), "Wt": (Wt_host, 2)}, writable=("Wt",), + ) + + with torch.no_grad(), utils._cuda_tf32(torch, tf32, rc.device): + for iteration in range(1, rc.max_iter + 1): + # One 2-D GEMM per replicate: a batched matmul may pick another + # reduction path as gpu-batch changes, which Fast-HALS amplifies. + for replicate in range(replicates): + gram_w[replicate].copy_(H[replicate] @ H[replicate].T) + if l2_w != 0.0: + gram_w.diagonal(dim1=-2, dim2=-1).add_(l2_w) + # One order for the whole W sweep, not a fresh one per window. + order_w = next_order() + violation = None + + for start, stop, window in windows: + Xw, Ww = window["X"], window["Wt"] + cross_w = H.new_empty((replicates, components, stop - start)) + for replicate in range(replicates): + cross_w[replicate].copy_(H[replicate] @ Xw.T) + if l1_w != 0.0: + cross_w.sub_(l1_w) + window_violation = _hals_sweep( + Ww, gram_w, cross_w, order_w, stop_rule.active + ) + violation = ( + window_violation if violation is None + else violation + window_violation + ) + # Use UPDATED W. Raw products accumulate across all cells; + # apply H regularization only once, with full sample scaling. + for replicate in range(replicates): + gram = Ww[replicate] @ Ww[replicate].T + cross = Ww[replicate] @ Xw + if start == 0: + gram_h[replicate].copy_(gram) + cross_h[replicate].copy_(cross) + else: + gram_h[replicate].add_(gram) + cross_h[replicate].add_(cross) + # Drop this window before the next one is uploaded. + del Xw, Ww, cross_w + + _regularize_cd_products(gram_h, cross_h, l1_h, l2_h) + violation = violation + _hals_sweep( + H, gram_h, cross_h, next_order(), stop_rule.active ) - active = active & ~converged - if not bool(active.any()): + if stop_rule.update(iteration, violation): break - return Wt, H, n_iter + return windows.host("Wt"), H, stop_rule.n_iter + + +def _fit_cd(torch, rc, Xcompute, Wt_host, H, seeds, nmf_kwargs, update_h): + """Run batched sklearn-compatible Fast-HALS to projected-gradient convergence. + + Wt_host is in host memory and H is already on the device. Without row + tiling this is the pre-tiling full-matrix computation; with tiling, the + products are computed one row window at a time. Returns the updated + (Wt_host, H, n_iter). + """ + n_rows, n_features = Xcompute.shape + window_rows = tiling._resolve_window_rows( + torch, rc.opt["row_tiling_ratio"], rc.device, n_rows, + budget=lambda: _window_budget( + update_h, len(seeds), rc.k, n_features, Xcompute.dtype.itemsize, + torch_sweep=_get_hals_cuda_backend() is None, + ), + solver="factorization CD" if update_h else "fixed-H CD", + ) + fit = _fit_cd_factorize if update_h else _fit_cd_fixed_h + return fit(torch, rc, Xcompute, Wt_host, H, window_rows, seeds, nmf_kwargs) def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): @@ -283,7 +431,6 @@ def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): 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 @@ -315,44 +462,21 @@ def _nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs=None): 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] ) 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, - ) - H = ( - torch.as_tensor(H0, dtype=rc.dtype, device=rc.device) - .unsqueeze(0) - .expand(replicates, -1, -1) - .contiguous() + # sklearn CD ignores supplied W and starts fixed-H usages at zero. + Wt0 = np.zeros( + (replicates, rc.k, Xcompute.shape[0]), dtype=host_dtype ) + Hs0 = np.repeat(H0[None, :, :], replicates, axis=0) - regularization = _cd_regularization( - nmf_kwargs, Xcompute.shape[0], Xcompute.shape[1] - ) - Wt, H, _ = _fit_cd( - torch, - Xg, - Wt, - H, - rc.max_iter, - rc.tol, - update_h, - regularization, - bool(nmf_kwargs.get("shuffle", False)), - seeds, - utils._want_tf32(torch, rc), - rc.device, + H = torch.as_tensor(Hs0, dtype=rc.dtype, device=rc.device) + Wt_host, H, _ = _fit_cd( + torch, rc, Xcompute, Wt0, H, seeds, nmf_kwargs, update_h ) - Hc = H.cpu().double().numpy() - Wc = Wt.transpose(-2, -1).cpu().double().numpy() + Wc = Wt_host.transpose(0, 2, 1).astype(np.float64, copy=False) return [(Hc[r], Wc[r]) for r in range(replicates)] diff --git a/src/cnmf/gpunmf/tiling.py b/src/cnmf/gpunmf/tiling.py new file mode 100644 index 0000000..34b4571 --- /dev/null +++ b/src/cnmf/gpunmf/tiling.py @@ -0,0 +1,112 @@ +"""Row tiling shared by the GPU NMF solvers. + +A solver works on the rows (cells) of X and W either all at once, resident on +the device, or in windows streamed from host memory. This module holds the +solver-neutral parts: turning the row tiling ratio into a window size, and +moving row windows between host and device. Each solver supplies its own +memory model and arithmetic. +""" + +import numpy as np + +# Automatic sizing leaves max(512 MiB, 10% of free VRAM) unallocated for the +# CUDA context, cuBLAS workspaces and allocator fragmentation. +_VRAM_RESERVE_BYTES = 512 * (1 << 20) +_VRAM_RESERVE_FRACTION = 0.10 + + +def _resolve_window_rows(torch, ratio, device, n_rows, budget, solver): + """Return the rows per window for a row tiling ratio. + + None or 1 keeps every row resident. A ratio in (0, 1) takes + floor(n_rows * ratio) rows, at least one. 0 takes the most rows that fit + in free CUDA memory, so the full matrix stays resident whenever it fits; + budget() gives the solver's (fixed_bytes, bytes_per_row). + """ + if ratio is None: + return n_rows + if ratio != 0: + return max(1, int(n_rows * ratio)) + if not device.startswith("cuda"): + return n_rows + + fixed_bytes, bytes_per_row = budget() + free_bytes, _ = torch.cuda.mem_get_info(device) + reserve = max( + _VRAM_RESERVE_BYTES, int(free_bytes * _VRAM_RESERVE_FRACTION) + ) + available = free_bytes - reserve - fixed_bytes + if available < bytes_per_row: + raise MemoryError( + f"GPU {solver} does not have enough VRAM for one row; " + "reduce gpu batch or n_components" + ) + return min(n_rows, available // bytes_per_row) + + +def _to_device(torch, array, like): + """Copy a host array to like's device and dtype as one contiguous tensor.""" + return torch.as_tensor( + np.ascontiguousarray(array), dtype=like.dtype, device=like.device + ) + + +class _RowWindows: + """Host arrays whose rows are processed on the device one window at a time. + + arrays maps a name to (host array, row axis). With one window covering + every row, the arrays are uploaded once and stay resident. Otherwise each + window is uploaded when visited, and the writable arrays are copied back + to the host when the caller moves on. + """ + + def __init__(self, torch, like, n_rows, window_rows, arrays, writable=()): + self._torch = torch + self._like = like + self._arrays = arrays + self._writable = tuple(writable) + self.resident = window_rows >= n_rows + self._bounds = [ + (start, min(start + window_rows, n_rows)) + for start in range(0, n_rows, window_rows) + ] + self._resident = ( + {name: self._upload(name, 0, n_rows) for name in arrays} + if self.resident else None + ) + + def _rows(self, name, start, stop): + array, axis = self._arrays[name] + index = [slice(None)] * array.ndim + index[axis] = slice(start, stop) + return array, tuple(index) + + def _upload(self, name, start, stop): + array, index = self._rows(name, start, stop) + return _to_device(self._torch, array[index], self._like) + + def __iter__(self): + """Yield (start, stop, tensors), with tensors[name] on the device. + + A streamed window is released when the caller moves on. Callers that + keep their own references to its tensors must drop them first, or two + windows are alive during the next upload. + """ + for start, stop in self._bounds: + if self.resident: + yield start, stop, self._resident + continue + tensors = {name: self._upload(name, start, stop) for name in self._arrays} + try: + yield start, stop, tensors + finally: + for name in self._writable: + array, index = self._rows(name, start, stop) + array[index] = tensors[name].cpu().numpy() + tensors.clear() + + def host(self, name): + """Return an array's current values in host memory.""" + if self.resident: + return self._resident[name].cpu().numpy() + return self._arrays[name][0] diff --git a/src/cnmf/gpunmf/utils.py b/src/cnmf/gpunmf/utils.py index 319d336..cb65c07 100644 --- a/src/cnmf/gpunmf/utils.py +++ b/src/cnmf/gpunmf/utils.py @@ -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", ) @@ -49,6 +51,7 @@ def _validate_engine_args(args, available_solvers): # sequencial validation of engine args for the given command _validate_engine_args_for_command(args) _validate_nmf_solver(args.solver, args.beta_loss, available_solvers) + _validate_row_tiling_ratio(args.gpu_row_tiling_ratio) def _validate_engine_args_for_command(args, available_commands=("prepare", "factorize", "consensus")): @@ -92,6 +95,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()): @@ -101,6 +105,19 @@ def gpu_kwargs_from_args(args): +def _validate_row_tiling_ratio(value): + """None or 1 disables tiling; 0 selects automatic sizing; (0, 1) sets a ratio.""" + if value is None: + return None + try: + ratio = float(value) + except (TypeError, ValueError): + ratio = None + if ratio is None or not 0 <= ratio <= 1: + raise ValueError("gpu row tiling ratio must be None or a finite number in [0, 1]") + return ratio + + def _resolve_gpu_opts(gpu_kwargs): """Merge Nextflow-provided gpu_kwargs over defaults into a typed opts dict.""" raw = dict(gpu_kwargs or {}) @@ -115,15 +132,19 @@ def parse_typed(value, default, cast, normalize=None): def parse_positive_int(value, default): return max(1, parse_typed(value, default, int)) + def parse_row_tiling_ratio(value, default): + return parse_typed(value, default, _validate_row_tiling_ratio) + 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), - allow_tf32 = parse_bool(raw.get("allow_tf32"), DEFAULT_GPU["allow_tf32"]), - compile = parse_bool(raw.get("compile"), DEFAULT_GPU["compile"]), - eps = parse_typed(raw.get("eps"), DEFAULT_GPU["eps"], float), - 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"]), + device = parse_typed(raw.get("device"), DEFAULT_GPU["device"], str, str.lower), + dtype = parse_typed(raw.get("dtype"), DEFAULT_GPU["dtype"], str, str.lower), + allow_tf32 = parse_bool(raw.get("allow_tf32"), DEFAULT_GPU["allow_tf32"]), + compile = parse_bool(raw.get("compile"), DEFAULT_GPU["compile"]), + eps = parse_typed(raw.get("eps"), DEFAULT_GPU["eps"], float), + 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_row_tiling_ratio(raw.get("row_tiling_ratio"), DEFAULT_GPU["row_tiling_ratio"]), ) diff --git a/tests/test_cd_tiling_runtime.py b/tests/test_cd_tiling_runtime.py new file mode 100644 index 0000000..34da7ad --- /dev/null +++ b/tests/test_cd_tiling_runtime.py @@ -0,0 +1,126 @@ +"""CPU-only checks for CD's tiling memory model, window wiring and stop rule. + +Solver-neutral window sizing and row windows are covered in test_tiling.py. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +from utils import kernel, require_nmf_runtime + + +# Hand-computed budgets for G=13, K=3, B=2; H is already allocated. +# Per row: W and cross windows (2BK) plus one H·Xᵀ output (K); factorization +# adds the X window (G); fixed-H takes the larger of that and its one-off +# product pass (G + BK + K). The torch sweep adds 10B values + 8B index bytes. +# Fixed: orders and state B(8K + 9) bytes + 3B values, plus (B + 1)K² Gram +# values; factorization adds BK² + (B + 1)KG values and, with the torch +# sweep, G times its per-row sweep bytes. +@pytest.mark.parametrize("update_h,dtype,torch_sweep,fixed,per_row", [ + (False, np.float32, False, 198, 88), + (False, np.float64, False, 330, 176), + (True, np.float32, False, 738, 112), + (True, np.float64, False, 1410, 224), + (False, np.float32, True, 198, 156), + (True, np.float32, True, 1986, 208), +]) +def test_cd_window_budget(kernel, update_h, dtype, torch_sweep, fixed, per_row): + itemsize = np.dtype(dtype).itemsize + assert kernel.solver_cd._window_budget( + update_h, 2, 3, 13, itemsize, torch_sweep + ) == (fixed, per_row) + + +@pytest.mark.parametrize("update_h,torch_sweep", [(False, False), (True, False), (True, True)]) +def test_fit_cd_sizes_auto_windows_from_the_cd_budget(kernel, monkeypatch, update_h, torch_sweep): + cd = kernel.solver_cd + monkeypatch.setattr(cd, "_HALS_CUDA_BACKEND", None if torch_sweep else object()) + window_rows = [] + loop = "_fit_cd_factorize" if update_h else "_fit_cd_fixed_h" + monkeypatch.setattr(cd, loop, lambda *args: window_rows.append(args[5])) + fixed, per_row = cd._window_budget(update_h, 2, 3, 13, 4, torch_sweep) + reserve = kernel.tiling._VRAM_RESERVE_BYTES + free_bytes = reserve + fixed + 8 * per_row - 1 + backend = SimpleNamespace(cuda=SimpleNamespace( + mem_get_info=lambda device: (free_bytes, 16 * (1 << 30)) + )) + rc = SimpleNamespace(k=3, device="cuda:7", opt={"row_tiling_ratio": 0}) + X = np.empty((2065, 13), dtype=np.float32) + + def fit(): + cd._fit_cd(backend, rc, X, None, None, [0, 1], {}, update_h) + + fit() + assert window_rows == [7] + free_bytes = reserve + fixed + per_row - 1 + mode = "factorization" if update_h else "fixed-H" + with pytest.raises(MemoryError, match=f"GPU {mode} CD .* one row"): + fit() + + +@pytest.mark.parametrize("tol", [0, 1e-4]) +def test_stop_rule_zero_initial_violation_stops_immediately(kernel, tol): + torch = require_nmf_runtime() + stop = kernel.solver_cd._StopRule(torch, 2, tol, "cpu") + assert stop.update(1, torch.zeros(2)) + assert stop.n_iter.tolist() == [1, 1] + assert stop.active.tolist() == [False, False] + + +def test_stop_rule_tracks_and_freezes_each_replicate(kernel): + torch = require_nmf_runtime() + stop = kernel.solver_cd._StopRule(torch, 3, 0.1, "cpu") + initial = torch.tensor([10.0, 20.0, 0.0], dtype=torch.float64) + assert not stop.update(1, initial) + initial.zero_() # The stored baseline must own its data. + assert stop.violation_init.tolist() == [10.0, 20.0, 0.0] + assert not stop.update(2, torch.tensor([1.0, 4.0, 0.0])) + assert stop.active.tolist() == [False, True, False] + assert stop.n_iter.tolist() == [2, 2, 1] + assert stop.update(3, torch.tensor([999.0, 1.0, 999.0])) + assert stop.active.tolist() == [False, False, False] + assert stop.n_iter.tolist() == [2, 3, 1] + + +def test_stop_rule_tol_zero_still_stops_at_exact_zero(kernel): + torch = require_nmf_runtime() + stop = kernel.solver_cd._StopRule(torch, 1, 0, "cpu") + assert not stop.update(1, torch.tensor([1.0])) + assert not stop.update(2, torch.tensor([1e-20])) + assert stop.update(3, torch.tensor([0.0])) + assert stop.n_iter.item() == 3 + + +@pytest.mark.parametrize("update_h", [False, True], ids=["fixed-h", "factorize"]) +def test_row_windows_never_upload_the_full_matrix(kernel, monkeypatch, update_h): + """Below ratio 1, X, W and cross-products move one window at a time.""" + torch = require_nmf_runtime() + rng = np.random.default_rng(5) + n_rows = 101 # Distinct from every other dimension below. + X = rng.random((n_rows, 7)) + 0.1 + nmf_kwargs = dict( + solver="cd", n_components=3, init="random", max_iter=4, tol=0.0, + update_H=update_h, + ) + if not update_h: + nmf_kwargs["H"] = rng.random((3, 7)) + 0.1 + shapes = [] + as_tensor = torch.as_tensor + + def recording_as_tensor(value, *args, **kwargs): + shapes.append(np.shape(value)) + return as_tensor(value, *args, **kwargs) + + monkeypatch.setattr(torch, "as_tensor", recording_as_tensor) + outputs = kernel.solver_cd._nmf_gpu_cd( + X, [0, 1], nmf_kwargs, + dict(device="cpu", dtype="fp64", row_tiling_ratio=0.25), + ) + + window = int(n_rows * 0.25) + row_counts = [shape[0] if len(shape) == 2 else shape[-1] for shape in shapes if len(shape) >= 2] + assert n_rows not in row_counts + assert max(row_counts) == max(window, 7) # 7 is H's feature axis. + assert all(np.isfinite(W).all() and (W >= 0).all() for _, W in outputs) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 72761c2..053e728 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -49,7 +49,8 @@ sklearn CD parity Pins sklearn's serial W-then-H Fast-HALS updates, regularization scaling, shuffled coordinate stream, fixed-H zero initialization, and per-replicate - projected-gradient stopping for fp64 and fp32 batches. + projected-gradient stopping for fp64 and fp32 batches. Tiling must preserve + W/H bits and stopping across storage windows; sklearn parity uses tolerances. """ import builtins @@ -924,15 +925,18 @@ def test_sklearn_cd_matches_batched_cuda_when_available( "dtype_name,np_dtype", [("fp32", np.float32), ("fp64", np.float64)], ) +@pytest.mark.parametrize("row_tiling_ratio", [None, 0.5, 1 / 3]) +@pytest.mark.parametrize("update_h", [False, True], ids=["fixed-h", "factorize"]) def test_cd_cuda_results_are_invariant_to_batch_width( - kernel, dtype_name, np_dtype + kernel, dtype_name, np_dtype, row_tiling_ratio, update_h ): """Changing the replicate batch width must not change a CD trajectory.""" torch = require_nmf_runtime() if not torch.cuda.is_available(): pytest.skip("CUDA is not available") - X = np.random.default_rng(117).random((41, 23), dtype=np_dtype) + rng = np.random.default_rng(117) + X = rng.random((2065, 23), dtype=np_dtype) X += np_dtype(0.1) seeds = [3, 37, 83] nmf_kwargs = _cd_nmf_kwargs( @@ -943,12 +947,16 @@ def test_cd_cuda_results_are_invariant_to_batch_width( alpha_W=0.0, alpha_H=0.0, l1_ratio=0.0, + update_H=update_h, ) + if not update_h: + nmf_kwargs["H"] = rng.random((5, 23), dtype=np_dtype) + np_dtype(0.1) gpu_kwargs = { "device": "cuda", "dtype": dtype_name, "allow_tf32": False, "compile": False, + "row_tiling_ratio": row_tiling_ratio, } batched = kernel.solver_cd._nmf_gpu_cd(X, seeds, nmf_kwargs, gpu_kwargs) @@ -956,8 +964,140 @@ def test_cd_cuda_results_are_invariant_to_batch_width( (single_H, single_W), = kernel.solver_cd._nmf_gpu_cd( X, [seed], nmf_kwargs, gpu_kwargs ) - np.testing.assert_array_equal(batched_H, single_H) - np.testing.assert_array_equal(batched_W, single_W) + np.testing.assert_array_equal(batched_H.view(np.uint64), single_H.view(np.uint64)) + np.testing.assert_array_equal(batched_W.view(np.uint64), single_W.view(np.uint64)) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("dtype_name,np_dtype", [("fp32", np.float32), ("fp64", np.float64)]) +@pytest.mark.parametrize("update_h", [False, True], ids=["fixed-h", "factorize"]) +@pytest.mark.parametrize("shuffle", [False, True], ids=["cyclic", "shuffled"]) +def test_cd_row_tiling_agrees_with_untiled( + kernel, monkeypatch, device, dtype_name, np_dtype, update_h, shuffle +): + """Row windows change only rounding; untiled runs keep their exact bits. + + Seed 213 generates synthetic positive input, not a downloaded fixture. + Ratio 1 and automatic sizing keep all 2065 rows resident, so they match + the untiled bits. Half and third windows split the sums over cells, so + factors, reconstruction error and stopping must agree numerically. + """ + torch = require_nmf_runtime() + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + rng = np.random.default_rng(213) + X = (rng.random((2065, 17)) + 0.1).astype(np_dtype) + nmf_kwargs = _cd_nmf_kwargs( + 3, seed=0, max_iter=40, tol=0.03, update_H=update_h, shuffle=shuffle + ) + if not update_h: + nmf_kwargs["H"] = (rng.random((3, 17)) + 0.1).astype(np_dtype) + options = dict(device=device, dtype=dtype_name, allow_tf32=False, compile=False) + captured = [] + original_fit = kernel.solver_cd._fit_cd + + def capture_iterations(*args, **kwargs): + result = original_fit(*args, **kwargs) + captured.append(result[2].cpu().numpy().copy()) + return result + + monkeypatch.setattr(kernel.solver_cd, "_fit_cd", capture_iterations) + expected = kernel.solver_cd._nmf_gpu_cd(X, [0, 7], nmf_kwargs, options) + expected_n_iter = captured[-1] + assert (expected_n_iter < nmf_kwargs["max_iter"]).all() + # CPU/sklearn and CUDA use different arithmetic: require numerical parity. + rtol, atol = (2e-4, 2e-5) if dtype_name == "fp32" else (2e-9, 2e-10) + sse_rtol = 1e-5 if dtype_name == "fp32" else 1e-12 + for replicate, seed in enumerate((0, 7)): + reference_H, reference_W, reference_n_iter = _sklearn_cd_reference( + X, dict(nmf_kwargs, random_state=seed) + ) + np.testing.assert_allclose(expected[replicate][0], reference_H, rtol=rtol, atol=atol) + np.testing.assert_allclose(expected[replicate][1], reference_W, rtol=rtol, atol=atol) + assert expected_n_iter[replicate] == reference_n_iter + + def sse(H, W): + return float(np.square(X.astype(np.float64) - W @ H).sum()) + + for ratio in (0, 1, 0.5, 1 / 3): + actual = kernel.solver_cd._nmf_gpu_cd( + X, [0, 7], nmf_kwargs, dict(options, row_tiling_ratio=ratio) + ) + n_iter = captured[-1] + for replicate, ((actual_H, actual_W), (expected_H, expected_W)) in enumerate( + zip(actual, expected) + ): + assert_valid_nmf_output(X, actual_H, actual_W, 3) + if not update_h: + np.testing.assert_array_equal(actual_H, nmf_kwargs["H"]) + if ratio in (0, 1): + # Every row stays resident: the untiled path, bit for bit. + np.testing.assert_array_equal(n_iter, expected_n_iter) + np.testing.assert_array_equal(actual_H.view(np.uint64), expected_H.view(np.uint64)) + np.testing.assert_array_equal(actual_W.view(np.uint64), expected_W.view(np.uint64)) + continue + assert abs(int(n_iter[replicate]) - int(expected_n_iter[replicate])) <= 1 + expected_sse = sse(expected_H, expected_W) + assert abs(sse(actual_H, actual_W) - expected_sse) <= sse_rtol * expected_sse + if n_iter[replicate] == expected_n_iter[replicate]: + np.testing.assert_allclose(actual_H, expected_H, rtol=rtol, atol=atol) + np.testing.assert_allclose(actual_W, expected_W, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("ratio", [None, 0, 1, 0.5, 1 / 3]) +def test_gpu_kwargs_from_args_forwards_cd_row_tiling_ratio(kernel, ratio): + """The optional CLI value must survive argument-to-runtime resolution.""" + args = _engine_args(engine="gpu", solver="cd", gpu_row_tiling_ratio=ratio) + assert kernel.utils.gpu_kwargs_from_args(args)["row_tiling_ratio"] == ratio + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("dtype_name,np_dtype", [("fp32", np.float32), ("fp64", np.float64)]) +def test_cd_fixed_h_products_are_computed_once_on_selected_device( + kernel, monkeypatch, device, dtype_name, np_dtype +): + """Fixed-H products are computed once on the device, one window at a time. + + GEMMs follow the row window: the whole matrix when untiled, otherwise + floor(n * ratio) rows plus the remainder. Counting Tensor matmuls also + catches accidentally restoring NumPy products. + """ + torch = require_nmf_runtime() + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + rng = np.random.default_rng(215) + X = (rng.random((2065, 17)) + 0.1).astype(np_dtype) + fixed_H = (rng.random((3, 17)) + 0.1).astype(np_dtype) + nmf_kwargs = _cd_nmf_kwargs(3, seed=0, max_iter=4, update_H=False, H=fixed_H) + calls = [] + matmul = torch.Tensor.__matmul__ + + def record_matmul(left, right): + assert left.device.type == right.device.type == device + assert not torch.is_grad_enabled() + if device == "cuda": + assert not torch.backends.cuda.matmul.allow_tf32 + calls.append((tuple(left.shape), tuple(right.shape))) + return matmul(left, right) + + monkeypatch.setattr(torch.Tensor, "__matmul__", record_matmul) + for ratio in (None, 0.5, 1 / 3): + calls.clear() + outputs = kernel.solver_cd._nmf_gpu_cd( + X, [0, 7], nmf_kwargs, + dict(device=device, dtype=dtype_name, row_tiling_ratio=ratio, allow_tf32=False), + ) + rows = len(X) if ratio is None else max(1, int(len(X) * ratio)) + windows = [min(rows, len(X) - start) for start in range(0, len(X), rows)] + # Two replicate Grams, then each window's cross-products for each + # replicate; nothing is recomputed during the four W iterations. + assert calls == [((3, 17), (17, 3))] * 2 + [ + ((3, 17), (17, window)) for window in windows for _ in range(2) + ] + for actual_H, actual_W in outputs: + assert_valid_nmf_output(X, actual_H, actual_W, 3) + np.testing.assert_array_equal(actual_H, fixed_H) # --------------------------------------------------------------------- @@ -1087,6 +1227,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, # not overridden here -> default } @@ -1693,6 +1834,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 +1855,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) @@ -1792,6 +1935,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, # default added; not set on the CLI here } diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 2d6eb32..b05be16 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, diff --git a/tests/test_synthetic_regression.py b/tests/test_synthetic_regression.py index 6d73700..2c5c63c 100644 --- a/tests/test_synthetic_regression.py +++ b/tests/test_synthetic_regression.py @@ -163,9 +163,11 @@ def test_prepare_matches_dylkot_reference(generated_fixture, tmp_path): assert actual_parameters == expected_parameters +@pytest.mark.parametrize("row_tiling_ratio", [None, 0.5, 1 / 3], ids=["untiled", "half", "third"]) def test_gpu_cd_factorize_matches_dylkot_cpu_reference( generated_fixture, tmp_path, + row_tiling_ratio, ): """GPU-engine FP64 CD on CPU must reproduce canonical CPU-CD spectra.""" pytest.importorskip("torch") @@ -186,6 +188,7 @@ def test_gpu_cd_factorize_matches_dylkot_cpu_reference( gpu_check_every=1, gpu_compile_block=None, gpu_batch=N_REPLICATES, + gpu_row_tiling_ratio=row_tiling_ratio, ) actual = configure_nmf_engine(cNMF, args) expected = _reference_cnmf(generated_fixture.reference_root, "factorize") diff --git a/tests/test_tiling.py b/tests/test_tiling.py new file mode 100644 index 0000000..e9d3d38 --- /dev/null +++ b/tests/test_tiling.py @@ -0,0 +1,161 @@ +"""Solver-neutral row tiling: window sizing and host/device row windows.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +from utils import kernel, require_nmf_runtime + + +def _budget_never_used(): + raise AssertionError("only automatic sizing on CUDA may evaluate the budget") + + +@pytest.mark.parametrize("ratio,expected", [(None, 13), (1, 13), (0.5, 6), (1 / 3, 4), (1e-9, 1)]) +def test_explicit_ratios_never_query_cuda(kernel, ratio, expected): + # No CUDA API exists on this fake backend: explicit sizing must not use it. + assert kernel.tiling._resolve_window_rows( + SimpleNamespace(), ratio, "cuda:7", 13, _budget_never_used, "test solver" + ) == expected + + +def test_auto_ratio_off_cuda_keeps_all_rows(kernel): + assert kernel.tiling._resolve_window_rows( + SimpleNamespace(), 0, "cpu", 13, _budget_never_used, "test solver" + ) == 13 + + +def test_auto_ratio_fills_free_cuda_memory_after_the_reserve(kernel, monkeypatch): + tiling = kernel.tiling + fixed, per_row = 1000, 100 + reserve = tiling._VRAM_RESERVE_BYTES + free_bytes = reserve + fixed + 8 * per_row - 1 + devices = [] + + def mem_get_info(device): + devices.append(device) + return free_bytes, 16 * (1 << 30) + + backend = SimpleNamespace(cuda=SimpleNamespace(mem_get_info=mem_get_info)) + + def rows(): + return tiling._resolve_window_rows( + backend, 0, "cuda:7", 2065, lambda: (fixed, per_row), "test solver" + ) + + assert rows() == 7 + free_bytes = reserve + fixed + per_row + assert rows() == 1 + free_bytes -= 1 + with pytest.raises(MemoryError, match="GPU test solver .* enough VRAM for one row"): + rows() + free_bytes = 8 * (1 << 30) + assert rows() == 2065 # The full matrix stays resident when it fits. + assert devices == ["cuda:7"] * 4 + + # The fractional reserve applies once it exceeds the absolute minimum. + monkeypatch.setattr(tiling, "_VRAM_RESERVE_BYTES", 0) + monkeypatch.setattr(tiling, "_VRAM_RESERVE_FRACTION", 0.25) + free_bytes = (4 * (fixed + 7 * per_row) + 2) // 3 + assert rows() == 7 + + +@pytest.mark.parametrize("value", [-0.01, 1.01, np.nan, np.inf, -np.inf, "nan", "inf", "invalid"]) +def test_invalid_row_tiling_ratios_are_rejected(kernel, value): + with pytest.raises(ValueError, match="row tiling ratio"): + kernel.utils._resolve_gpu_opts({"row_tiling_ratio": value}) + + +@pytest.mark.parametrize("value,expected", [(None, None), (0, 0), ("0.5", 0.5), (1, 1)]) +def test_valid_row_tiling_ratios_are_parsed(kernel, value, expected): + assert kernel.utils._resolve_gpu_opts({"row_tiling_ratio": value})[ + "row_tiling_ratio" + ] == expected + + +def _host_arrays(n_rows): + X = np.arange(n_rows * 5, dtype=np.float64).reshape(n_rows, 5) + Wt = np.arange(2 * 3 * n_rows, dtype=np.float64).reshape(2, 3, n_rows) + return X, Wt + + +def _copying_backend(torch, uploads): + # CPU tensors would alias NumPy; copy so device and host stay distinct. + def as_tensor(value, **kwargs): + uploads.append(np.shape(value)) + return torch.as_tensor(value, **kwargs).clone() + + return SimpleNamespace(as_tensor=as_tensor) + + +@pytest.mark.parametrize("window_rows", [13, 20]) +def test_row_windows_keep_a_single_window_resident(kernel, window_rows): + torch = require_nmf_runtime() + X, Wt = _host_arrays(13) + original_Wt = Wt.copy() + uploads = [] + windows = kernel.tiling._RowWindows( + _copying_backend(torch, uploads), torch.empty(0, dtype=torch.float64), + 13, window_rows, {"X": (X, 0), "Wt": (Wt, 2)}, writable=("Wt",), + ) + + assert windows.resident + for _ in range(3): + bounds = [] + for start, stop, tensors in windows: + bounds.append((start, stop)) + tensors["Wt"].add_(1) + assert bounds == [(0, 13)] + assert uploads == [X.shape, Wt.shape] # Uploaded once, not per pass. + np.testing.assert_array_equal(windows.host("Wt"), original_Wt + 3) + np.testing.assert_array_equal(Wt, original_Wt) # Resident: host untouched. + + +@pytest.mark.parametrize("window_rows", [1, 4, 5, 12]) +def test_row_windows_stream_rows_and_write_back_writable_arrays(kernel, window_rows): + torch = require_nmf_runtime() + n_rows = 13 + X, Wt = _host_arrays(n_rows) + original_X, expected_Wt = X.copy(), Wt.copy() + uploads = [] + windows = kernel.tiling._RowWindows( + _copying_backend(torch, uploads), torch.empty(0, dtype=torch.float64), + n_rows, window_rows, {"X": (X, 0), "Wt": (Wt, 2)}, writable=("Wt",), + ) + + assert not windows.resident + for _ in range(2): + bounds = [] + for start, stop, tensors in windows: + bounds.append((start, stop)) + np.testing.assert_array_equal(tensors["X"].numpy(), X[start:stop]) + np.testing.assert_array_equal(tensors["Wt"].numpy(), expected_Wt[:, :, start:stop]) + tensors["X"].add_(100) # Not writable: must never reach the host. + tensors["Wt"].add_(1) + expected_Wt[:, :, start:stop] += 1 + held = tensors + assert held == {} # Released once the caller moved on. + assert bounds == [ + (start, min(start + window_rows, n_rows)) + for start in range(0, n_rows, window_rows) + ] + assert len(uploads) == 2 * 2 * len(bounds) # X and Wt, per window, per pass. + np.testing.assert_array_equal(windows.host("Wt"), expected_Wt) + np.testing.assert_array_equal(X, original_X) + + +def test_row_windows_write_back_even_when_the_caller_breaks(kernel): + torch = require_nmf_runtime() + X, Wt = _host_arrays(13) + expected = Wt.copy() + expected[:, :, :4] += 1 + windows = kernel.tiling._RowWindows( + _copying_backend(torch, []), torch.empty(0, dtype=torch.float64), + 13, 4, {"X": (X, 0), "Wt": (Wt, 2)}, writable=("Wt",), + ) + + for _, _, tensors in windows: + tensors["Wt"].add_(1) + break + np.testing.assert_array_equal(Wt, expected)