warp.tile_cholesky_solve#

warp.tile_cholesky_solve(
L: Tile[Float, tuple[int, int]],
y: Tile[Float, tuple[int, ...]],
fill_mode: str = 'lower',
) → Tile[Float, tuple[int, ...]]#
  • Kernel

Solve A x = y given the Cholesky factor of A and return x in a new tile.

With fill_mode="lower" L is the lower-triangular factor such that L @ L^T = A; with fill_mode="upper" L instead holds the upper-triangular factor U such that U^T @ U = A, despite the parameter name. Only the triangle selected by fill_mode is read; the opposite triangle is ignored. fill_mode must be a compile-time constant.

This is a cooperative operation, so every thread in the block must call it.

The right-hand side may be a 1D tile with shape=(M,) or a 2D tile with shape=(M, NRHS), which solves for all NRHS columns at once. L and y must have the same dtype, either float32 or float64.

On CUDA devices, Warp uses cuSolverDx when it is built with libmathdx and warp.config.enable_mathdx_solver (or the enable_mathdx_solver module option) is True. Otherwise, and on the CPU, it uses a cooperative fallback implementation, which is slower than cuSolverDx on CUDA devices.

Parameters:

  • L – A square triangular Cholesky factor of A, with shape=(M, M).

  • y – The right-hand side, a tile with shape=(M,) or shape=(M, NRHS).

  • fill_mode – "lower" or "upper". Must be a compile-time constant.

Returns:

A new tile in shared memory with the same shape and dtype as y, holding the solution x of A x = y.

See tile_cholesky() for a usage example.