import numpy as np
import warp as wp


def tile_cholesky(
  M: wp.array(dtype=wp.float32, ndim=2),
  L: wp.array(dtype=wp.float32, ndim=2),
  tilesize: int,
):
  @wp.kernel
  def cholesky(
    M: wp.array(dtype=wp.float32, ndim=2), L: wp.array(dtype=wp.float32, ndim=2)
  ):
    M_tile = wp.tile_load(M, shape=(tilesize, tilesize))
    L_tile = wp.tile_cholesky(M_tile)
    wp.tile_store(L, L_tile)

  wp.launch_tiled(cholesky, dim=(1,), inputs=[M, L], block_dim=32)


dim = 5
Mnp = np.eye(dim) * 2.0
Lnp = np.linalg.cholesky(Mnp)

Mwp = wp.array(Mnp, dtype=wp.float32)
Lwp = wp.zeros((dim, dim), dtype=wp.float32)

tilesize = 4
tile_cholesky(Mwp, Lwp, tilesize)

err = Lnp[:tilesize, :tilesize] - Lwp.numpy()[:tilesize, :tilesize]

print(f"Lnp - Lwp:\n{err}")
print(f"error(Lnp - Lwp):{np.linalg.norm(err)}")