Fix Sinkhorn CUDA graph caching

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.

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):