transport.samples.kernels.apply_fused_sqeuclid
module torchmatch.transport.samples.kernels.apply_fused_sqeuclidFused Schur complement matvec kernel for HVP CG acceleration.
This module implements a two-phase persistent kernel that fuses the axis1 and axis0
transport plan applications used in the HVP's CG linear operator, reducing kernel
launches from 2 to 1 per CG iteration.
The key innovation is using a spin-wait barrier with atomic counter for grid-wide
synchronization between phases, allowing both operations to run in a single kernel.
IMPORTANT: The spin barrier requires all blocks to be co-resident on the GPU.
If grid_size exceeds the maximum resident blocks (num_SMs * ~4), the kernel
will deadlock. This module automatically falls back to the two-kernel approach
when the grid size would be too large.
Members
function
- fused_schur_matvec_sqeuclidFused Schur complement matvec: out = denom * z - P^T @ (P @ z / diag_x).
Source: torchmatch/transport/samples/kernels/apply_fused_sqeuclid.py:1