warp.tile_lower_solve#

warp.tile_lower_solve( L: Tile[Float, tuple[int, int]], y: Tile[Float, tuple[int, ...]], ) → Tile[Float, tuple[int, ...]]#

Kernel: true
Python: false
Differentiable: true

Solve L z = y for a lower-triangular L and return z in a new tile.

Only the lower triangle of L, including its diagonal, is read; the strictly upper triangle is ignored. L must be non-singular with nonzero diagonal entries.

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.

Backward propagation computes gradients with respect to y and the lower-triangular parameterization of L; the strictly upper triangle of L receives zero gradient.

Parameters:

L – A square, non-singular, lower-triangular tile with shape=(M, M).
y – The right-hand side, a tile with shape=(M,) or shape=(M, NRHS).

Returns:

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

Example

@wp.kernel
def forward_substitution(ell: wp.array2d[float], y: wp.array[float], z: wp.array[float]):
    l_tile = wp.tile_load(ell, shape=(3, 3))
    y_tile = wp.tile_load(y, shape=3)
    wp.tile_store(z, wp.tile_lower_solve(l_tile, y_tile))

ell = wp.array(np.array([[2.0, 0.0, 0.0], [1.0, 3.0, 0.0], [4.0, 5.0, 6.0]], dtype=np.float32))
y = wp.array(np.array([2.0, 7.0, 32.0], dtype=np.float32))
z = wp.zeros(3, dtype=float)

# one full warp cooperates on the triangular solve
wp.launch_tiled(forward_substitution, dim=1, inputs=[ell, y], outputs=[z], block_dim=32)

print(z.numpy())

[1. 2. 3.]