Jax process runs ahead of solve

[meliao]

Hi, thanks for providing this package! I am encountering an issue when using the sparse solver repeatedly in an iterative loop on the GPU with jax. I believe what I am seeing is the effects of jax's asynchronous dispatch, which is causing the jax process to "run ahead" of the cholespy.CholeskySolverD.solve() call. At a high level, the loop looks like this:

for _ in range(N):
    x = jnp.zeros_like(b)
    cholespy_solver_obj.solve(b, x)
    
    b = do_other_stuff(x)    # x is still 0!

Do you have any advice about how to control this behavior? Calling .block_until_ready() on x does not work.

Here is a "minimal" example. I apologize about the amount of code devoted to generating the matrix A. It's generating a discrete difference operator used in TV denoising. When I tried to reproduce this with a simpler A such as a diagonal matrix, the behavior did not happen.

import cholespy
import jax.numpy as jnp
import jax
import scipy.sparse as sp
from cholespy import CholeskySolverD, MatrixType

jax.config.update("jax_enable_x64", True)


def _J_TV_indices(n: int) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Returns the data, row indices, and column indices for specifying
    CSC matrix for J_TV. This is a discrete difference operator applied
    to an n x n image flattened in column-major order. I'm using it for
    TV denoising.

    J_TV has shape (n^2 -> 2n^2 - 4n +2)
    """

    n_diff = n**2 - 2 * n + 1
    # Each row has two entries. The vertical diffs are in the first n**2 - n rows.
    row_idxes_firsthalf = jnp.repeat(jnp.arange(n_diff), 2)

    col_indices_firsthalf = []
    for i in range(n - 1):
        # Each iteration of this loop corresponds to a column of the image,
        # and (n-1) rows of the operator J_TV
        # Need col idxes like 0, n, 1, n +1, ...
        col_idxes_i = jnp.arange(i * n, i * n + n - 1)
        col_idxes_i = jnp.stack([col_idxes_i, n + col_idxes_i], axis=1)
        # Flatten the column indices
        col_idxes_i = col_idxes_i.flatten()

        col_indices_firsthalf.append(col_idxes_i)
    col_indices_firsthalf = jnp.concatenate(col_indices_firsthalf)

    # Each row has two entries. The horizontal diffs are in the last n**2 - n rows.
    row_idxes_secondhalf = jnp.repeat(jnp.arange(n_diff, 2 * n_diff), 2)
    col_indices_secondhalf = []
    for i in range(n - 1):
        # Each iteration of this loop corresponds to a row of the image, and (n-1) rows of the operator J_TV
        col_idxes_i = jnp.arange(i * n, (i + 1) * n)
        col_idxes_i = jnp.repeat(col_idxes_i, 2)

        # Remove first and last entry
        col_idxes_i = col_idxes_i[1:-1]
        col_indices_secondhalf.append(col_idxes_i)
    col_indices_secondhalf = jnp.concatenate(col_indices_secondhalf)

    all_row_idxes = jnp.concatenate(
        [row_idxes_firsthalf, row_idxes_secondhalf]
    )
    all_col_indices = jnp.concatenate(
        [col_indices_firsthalf, col_indices_secondhalf]
    )
    all_data = jnp.array([-1, 1] * (all_col_indices.shape[0] // 2))

    return all_data, all_row_idxes, all_col_indices


def generate_A(n: int) -> sp.csc_matrix:
    """A = I + J^\\top J"""
    J_TV_data, J_TV_row_idxes, J_TV_col_idxes = _J_TV_indices(n)
    J_TV = sp.csc_matrix(
        (J_TV_data, (J_TV_row_idxes, J_TV_col_idxes)),
        shape=(2 * (n**2 - 2 * n + 1), n**2),
    )
    I = sp.eye(n**2, format="csc", dtype=jnp.float32)

    # Compute I + J^\\top J
    J_TV_T = J_TV.transpose()
    J_TV_T_J_TV = J_TV_T @ J_TV
    A = I + J_TV_T_J_TV

    return A


@jax.jit
def jitted_normalization_step(z: jax.Array) -> jax.Array:
    norm = jnp.linalg.norm(z)
    return z / norm


def main() -> None:
    n = 100
    A = generate_A(n)
    gpu_bool = False
    if gpu_bool:
        device = jax.devices()[0]
    else:
        device = jax.devices("cpu")[0]
    A_indptr = jax.device_put(jnp.array(A.indptr), device=device)
    A_indices = jax.device_put(jnp.array(A.indices), device=device)
    A_data = jax.device_put(
        jnp.array(A.data, dtype=jnp.float64), device=device
    )
    x_t = jax.device_put(jnp.sin(jnp.arange(n**2)), device=device)
    print("A has %d rows and %d nonzero elements" % (A.shape[0], A.nnz))


    factor_obj = CholeskySolverD(
        n**2, A_indptr, A_indices, A_data, MatrixType.CSC
    )
    print("Factored I+J^T J")

    n_iter = 10
    xt_norms = jnp.zeros((n_iter,), dtype=jnp.float64)
    Ainv_xt_norms = jnp.zeros((n_iter,), dtype=jnp.float64)

    # Loop over interleaved linear solves with other operations.
    for t in range(n_iter):
        Ainv_xt = jnp.zeros_like(x_t)
        factor_obj.solve(x_t, Ainv_xt)
        Ainv_xt_norms = Ainv_xt_norms.at[t].set(jnp.linalg.norm(Ainv_xt))

        # At the time this is called, Ainv_xt is not ready.
        x_t = jitted_normalization_step(Ainv_xt)
        xt_norms = xt_norms.at[t].set(jnp.linalg.norm(x_t))
    print("xt norms:", xt_norms)
    print("Ainv_xt norms:", Ainv_xt_norms)


if __name__ == "__main__":
    main()

When I run this with gpu_bool=True, I see:

A has 10000 rows and 49204 nonzero elements
Factored I+J^T J
xt norms: [ 1. nan nan nan nan nan nan nan nan nan]
Ainv_xt norms: [33.07609835  0.          0.          0.          0.          0.
  0.          0.          0.          0.        ]

And when I run it on the cpu, I see:

A has 10000 rows and 49204 nonzero elements
Factored I+J^T J
xt norms: [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
Ainv_xt norms: [33.07609835  0.4844604   0.51777055  0.57615485  0.65323766  0.72506883
  0.77696109  0.8129019   0.84109767  0.86632467]