transport.samples.kernels.apply_fused_sqeuclid.fused_schur_matvec_sqeuclid
def fused_schur_matvec_sqeuclid(x: torch.Tensor, y: torch.Tensor, f: torch.Tensor, g: torch.Tensor, z: torch.Tensor, diag_x: torch.Tensor, denom: torch.Tensor, eps: float, x2: torch.Tensor | None = None, y2: torch.Tensor | None = None, piz_buffer: torch.Tensor | None = None, block_m: int = 64, block_n: int = 64, block_k: int = 64, num_warps: int = 4, num_stages: int = 2, use_exp2: bool = True, allow_tf32: bool = False) -> tuple[torch.Tensor, torch.Tensor]Fused Schur complement matvec: out = denom * z - P^T @ (P @ z / diag_x).
This is a single-kernel implementation that fuses two transport plan applications
(axis1 and axis0) with a grid-wide barrier for synchronization.
Args:
x: Source points (n, d)
y: Target points (m, d)
f: Source potential (n,) - raw-form: P = exp((f+g-C)/eps)
g: Target potential (m,) - raw-form: P = exp((f+g-C)/eps)
z: Input CG vector (m,)
diag_x: Source diagonal D_x = diag_factor_x * a_hat (n,)
denom: Target denominator D_y = diag_factor_y * b_hat + eps*tau2 (m,)
eps: Entropy regularization
x2: Precomputed ||x||^2 (n,) [optional]
y2: Precomputed ||y||^2 (m,) [optional]
piz_buffer: Reusable buffer for P @ z intermediate (n,) [optional]
Returns:
out: Result denom * z - P^T @ (P @ z / diag_x), shape (m,)
piz_buffer: The intermediate buffer (for reuse)
Note:
This kernel uses a spin-wait barrier for grid-wide synchronization.
The grid size must not exceed the number of co-resident blocks on the GPU.
If the grid would be too large, this function automatically falls back to
a two-kernel approach using apply_plan_vec_shifted (axis=1, then axis=0).
The fallback path converts raw potentials to shifted form and uses shifted-form apply kernels.
Source: torchmatch/transport/samples/kernels/apply_fused_sqeuclid.py:338