diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py index 785e069..ec352fd 100644 --- a/test/test_optimal_transport.py +++ b/test/test_optimal_transport.py @@ -1,7 +1,12 @@ +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + import numpy as np import pytest import torch +import torchmorph.optimal_transport as ot_impl from torchmorph import SinkhornSolver, build_cost_matrix try: @@ -189,6 +194,96 @@ def test_fused_short_run_without_cuda_graph_matches_cpu(log_space): assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5) +@pytest.mark.parametrize("log_space", LOG_SPACE) +@pytest.mark.parametrize("method", ["forward", "plan", "potentials"]) +def test_graph_reuses_buffers_with_new_inputs(log_space, method, monkeypatch): + """New tensors, costs and budgets must not replay stale data or recapture.""" + _require_cuda() + monkeypatch.setattr(ot_impl, "_GRAPH_CACHE", SimpleNamespace()) + graph = Mock(wraps=torch.cuda.CUDAGraph) + monkeypatch.setattr(torch.cuda, "CUDAGraph", graph) + solver = SinkhornSolver(epsilon=1.5, log_space=log_space) + solve = getattr(solver, method) + + for seed, max_iter in enumerate((100, 101, 124, 125, 200), start=1500): + solver.max_iter = max_iter + source = _make_positive((3, 6), device="cpu", seed=seed) + target = _make_positive((3, 6), device="cpu", seed=seed + 100) + cost = _make_positive((6, 6), device="cpu", seed=seed + 200) + expected = solve(source, target, cost) + actual = solve(source.cuda(), target.cuda(), cost.cuda()) + actual_parts = actual if isinstance(actual, tuple) else (actual,) + expected_parts = expected if isinstance(expected, tuple) else (expected,) + for cuda_value, cpu_value in zip(actual_parts, expected_parts): + assert torch.allclose(cuda_value.cpu(), cpu_value, rtol=1e-4, atol=1e-4) + assert graph.call_count == 1 + + +def test_graph_iteration_count_and_output_lifetime(monkeypatch): + """Warmup and previous replays must not affect this call or earlier outputs.""" + _require_cuda() + monkeypatch.setattr(ot_impl, "_GRAPH_CACHE", SimpleNamespace()) + saved = [] + + def launch(inputs, outputs, k): + outputs[0].add_(inputs[0], alpha=k) + + for value, max_iter in enumerate((99, 100, 101, 124, 125, 200), start=1): + inputs = (torch.full((2,), float(value), device="cuda"),) + outputs = (torch.zeros_like(inputs[0]),) + ot_impl._run_fused(launch, max_iter, inputs[0].device, inputs, outputs, "count") + saved.append((outputs[0], value * max_iter)) + for result, expected in saved: + assert torch.equal(result, torch.full_like(result, expected)) + + +def test_graph_rebuilds_for_configuration_and_stream(monkeypatch): + _require_cuda() + monkeypatch.setattr(ot_impl, "_GRAPH_CACHE", SimpleNamespace()) + graph = Mock(wraps=torch.cuda.CUDAGraph) + monkeypatch.setattr(torch.cuda, "CUDAGraph", graph) + + def launch(inputs, outputs, k): + outputs[0].add_(inputs[0], alpha=k) + + def run(size, signature): + inputs = (torch.ones(size, device="cuda"),) + outputs = (torch.zeros_like(inputs[0]),) + ot_impl._run_fused(launch, 100, inputs[0].device, inputs, outputs, signature) + assert torch.equal(outputs[0], torch.full_like(outputs[0], 100)) + + run(2, ("log", 1.5)) + run(3, ("log", 1.5)) + run(3, ("log", 2.0)) + run(3, "scaling") + assert graph.call_count == 4 + with torch.cuda.stream(torch.cuda.Stream()): + run(3, "scaling") + assert graph.call_count == 5 + + +@pytest.mark.parametrize("stage", ["warmup", "capture", "replay"]) +def test_graph_errors_propagate_without_retry(stage, monkeypatch): + """An error in any graph stage must never trigger a full-budget retry.""" + monkeypatch.setattr(ot_impl, "_GRAPH_CACHE", SimpleNamespace()) + stream = SimpleNamespace(cuda_stream=1, wait_stream=lambda _: None) + monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr(torch.cuda, "Stream", lambda: stream) + monkeypatch.setattr(torch.cuda, "stream", lambda _: nullcontext()) + graph = SimpleNamespace(replay=Mock(side_effect=RuntimeError("replay failed"))) + monkeypatch.setattr(torch.cuda, "CUDAGraph", lambda: graph) + monkeypatch.setattr(torch.cuda, "graph", lambda *args, **kwargs: nullcontext()) + error = RuntimeError(f"{stage} failed") + effects = {"warmup": [error], "capture": [None, error], "replay": [None, None]} + launch = Mock(side_effect=effects[stage]) + inputs, outputs = (torch.ones(2),), (torch.zeros(2),) + with pytest.raises(RuntimeError, match=f"{stage} failed"): + ot_impl._run_fused(launch, 101, "cuda:0", inputs, outputs, "error") + assert [call.args[2] for call in launch.call_args_list] == [25] * len(effects[stage]) + + @pytest.mark.parametrize("log_space", LOG_SPACE) def test_dtype_device_combinations_agree(log_space): """float32/float64 on CPU/CUDA must all reach the same solution. diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py index 3652845..ccb67e7 100644 --- a/torchmorph/optimal_transport.py +++ b/torchmorph/optimal_transport.py @@ -1,3 +1,4 @@ +from threading import local from typing import Optional import torch @@ -31,41 +32,63 @@ def _transpose_or_self(matrix: Tensor) -> Tensor: return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous() -_GRAPH_MIN_ITER = 100 # below this, plain launches beat the graph-capture setup cost +_GRAPH_MIN_ITER = 100 _GRAPH_CHUNK = 25 +_GRAPH_CACHE = local() # one workspace per thread; changing configuration replaces it -def _run_fused(launch, max_iter, device): - """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total. +def _run_fused(launch, max_iter, device, inputs, outputs, signature): + """Replay a cached chunk on fixed buffers, then copy out this call's result. - Long runs are launch-latency-bound (two tiny kernels per iteration), so a - chunk of iterations is captured into a CUDA graph once and replayed. The - iterations are fixed-point steps on ping-pong buffers, so a fallback that - runs extra iterations is always safe. + `launch(inputs, outputs, k)` enqueues k in-place iterations. Only the most + recent configuration is cached per thread, with streams isolated by the + cache key. Warmup uses scratch state and never counts as solver iterations. + Errors propagate instead of silently rerunning an already modified state. """ - if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): - launch(max_iter) - return - try: - with torch.cuda.device(device): - # Warm up on a side stream (required before capture), doing the - # first chunk of real iterations in the process. + with torch.cuda.device(device): + if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): + launch(inputs, outputs, max_iter) + return + + stream = torch.cuda.current_stream() + key = ( + device, + stream.cuda_stream, + signature, + tuple((tuple(t.shape), t.dtype, t.stride()) for t in (*inputs, *outputs)), + ) + entry = getattr(_GRAPH_CACHE, "entry", None) + if entry is None or entry[0] != key: + # Release the old workspace before allocating its replacement. + _GRAPH_CACHE.entry = None + entry = None + static_inputs = tuple(t.detach().clone() for t in inputs) + static_outputs = tuple(t.clone() for t in outputs) side_stream = torch.cuda.Stream() - side_stream.wait_stream(torch.cuda.current_stream()) + side_stream.wait_stream(stream) with torch.cuda.stream(side_stream): - launch(_GRAPH_CHUNK) - torch.cuda.current_stream().wait_stream(side_stream) + launch(static_inputs, static_outputs, _GRAPH_CHUNK) + stream.wait_stream(side_stream) graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - launch(_GRAPH_CHUNK) - for _ in range(max_iter // _GRAPH_CHUNK - 1): - graph.replay() - remainder = max_iter % _GRAPH_CHUNK - if remainder: - launch(remainder) - except RuntimeError: - launch(max_iter) # graphs unavailable; extra iterations are harmless + with torch.cuda.graph(graph, stream=side_stream): + launch(static_inputs, static_outputs, _GRAPH_CHUNK) + _GRAPH_CACHE.entry = (key, graph, static_inputs, static_outputs) + else: + _, graph, static_inputs, static_outputs = entry + for dst, src in zip(static_inputs, inputs): + dst.copy_(src) + + # Discard warmup or the previous call's solution before any replay. + for dst, src in zip(static_outputs, outputs): + dst.copy_(src) + for _ in range(max_iter // _GRAPH_CHUNK): + graph.replay() + remainder = max_iter % _GRAPH_CHUNK + if remainder: + launch(static_inputs, static_outputs, remainder) + for dst, src in zip(outputs, static_outputs): + dst.copy_(src) def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): @@ -78,7 +101,14 @@ def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): K_T = _transpose_or_self(K) u = torch.ones_like(a) v = torch.ones_like(b) - _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device) + _run_fused( + lambda ins, outs, k: _C.sinkhorn_fastiter(*ins, *outs, k), + max_iter, + a.device, + (a, b, K, K_T), + (u, v), + "scaling", + ) return u.log(), v.log() u = torch.ones_like(a) @@ -104,11 +134,12 @@ def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): log_u = torch.zeros_like(a) log_v = torch.zeros_like(b) _run_fused( - lambda k: _C.sinkhorn_logiter( - log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon - ), + lambda ins, outs, k: _C.sinkhorn_logiter(*ins, *outs, k, epsilon), max_iter, a.device, + (log_a, log_b, cost_matrix, cost_T), + (log_u, log_v), + ("log", epsilon), ) return log_u, log_v @@ -164,6 +195,9 @@ class SinkhornSolver(nn.Module): the input automatically: CUDA float32 tensors use the fused kernels, everything else uses pure torch ops (where `threshold` > 0 enables early stopping; the fused kernels always run `max_iter` iterations). + + Long fused runs reuse a CUDA graph with fixed scratch buffers. Each thread + caches only its most recent device, stream, shape, and kernel configuration. """ def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False):