module

transport.samples.kernels.apply_fused_sqeuclid

module torchmatch.transport.samples.kernels.apply_fused_sqeuclid

Fused 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

Source: torchmatch/transport/samples/kernels/apply_fused_sqeuclid.py:1