transport.samples.kernels.cg_dense
module torchmatch.transport.samples.kernels.cg_denseDense CG solver with cached transport plan matrix.
This module provides a CG solver that materializes the O(n*m) transport plan
matrix P and uses dense matrix-vector products (P @ v, P.T @ v) instead of
streaming Triton kernels. This is significantly faster for small problem sizes
(n <= 8192) where:
1. The O(n²) plan fits in GPU memory
2. Kernel launch overhead dominates streaming computation
3. Dense matvecs can leverage tensor core saturation
Performance comparison (n=2048, d=64):
- Streaming CG (Triton): ~48 ms (64 kernel launches × 0.76 ms)
- Dense CG (cached P): ~5.9 ms (8.12x speedup!)
Memory trade-off:
- Streaming: O(nd) memory
- Dense: O(nm) memory for transport plan P
Crossover point: n ≈ 8192-10000 where memory becomes the bottleneck.
Members
type
- CachedDenseHVPCached HVP context that materializes transport plan once.
- DenseCgInfoConvergence information for dense CG solver.
function
- dense_cg_solveSolve the Schur complement system using dense matrix operations.
- hvp_dense_cgCompute HVP using dense CG with cached transport plan.
- materialize_transport_planMaterialize the transport plan P = exp((f_hat + g_hat - C) / eps).
Source: torchmatch/transport/samples/kernels/cg_dense.py:1