warp-lang 1.8.1__py3-none-win_amd64.whl → 1.9.1__py3-none-win_amd64.whl

warp/tests/tile/test_tile_mathdx.py

@@ -21,7 +21,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.is_mathdx_enabled()
+wp.init()  # For wp.context.runtime.core.wp_is_mathdx_enabled()
 
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
@@ -45,7 +45,6 @@ def tile_math_matmul_kernel(
     wp.tile_store(gc, c, offset=(i * TILE_M, j * TILE_N))
 
 
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
 def test_tile_math_matmul(test, device):
     rng = np.random.default_rng(42)
 
@@ -93,7 +92,7 @@ def tile_math_fft_kernel_vec2d(gx: wp.array2d(dtype=wp.vec2d), gy: wp.array2d(dt
     wp.tile_store(gy, xy)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
+@unittest.skipUnless(wp.context.runtime.core.wp_is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_math_fft(test, device, wp_dtype):
     np_real_dtype = {wp.vec2f: np.float32, wp.vec2d: np.float64}[wp_dtype]
     np_cplx_dtype = {wp.vec2f: np.complex64, wp.vec2d: np.complex128}[wp_dtype]
@@ -124,503 +123,6 @@ def test_tile_math_fft(test, device, wp_dtype):
     # TODO: implement and test backward pass
 
 
-@wp.kernel()
-def tile_math_cholesky(
-    gA: wp.array2d(dtype=wp.float64),
-    gD: wp.array1d(dtype=wp.float64),
-    gL: wp.array2d(dtype=wp.float64),
-    gy: wp.array1d(dtype=wp.float64),
-    gx: wp.array1d(dtype=wp.float64),
-):
-    i, j = wp.tid()
-    # Load A, D & y
-    a = wp.tile_load(gA, shape=(TILE_M, TILE_M), storage="shared")
-    d = wp.tile_load(gD, shape=TILE_M, storage="shared")
-    y = wp.tile_load(gy, shape=TILE_M, storage="shared")
-    # Ensure tile_diag_add() and tile_cholesky_solve() work with transposed matrices
-    a_t = wp.tile_transpose(a)
-    # Compute L st LL^T = A^T + diag(D)
-    b = wp.tile_diag_add(a_t, d)
-    l = wp.tile_cholesky(b)
-    # Solve for y in LL^T x = y
-    x = wp.tile_cholesky_solve(l, y)
-    # Store L & y
-    wp.tile_store(gL, l)
-    wp.tile_store(gx, x)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_cholesky(test, device):
-    A_h = np.ones((TILE_M, TILE_M), dtype=np.float64)
-    D_h = 8.0 * np.ones(TILE_M, dtype=np.float64)
-    L_h = np.zeros_like(A_h)
-    Y_h = np.arange(TILE_M, dtype=np.float64)
-    X_h = np.zeros_like(Y_h)
-
-    A_np = A_h.T + np.diag(D_h)
-    L_np = np.linalg.cholesky(A_np)
-    X_np = np.linalg.solve(A_np, Y_h)
-
-    A_wp = wp.array2d(A_h, requires_grad=True, dtype=wp.float64, device=device)
-    D_wp = wp.array2d(D_h, requires_grad=True, dtype=wp.float64, device=device)
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    Y_wp = wp.array2d(Y_h, requires_grad=True, dtype=wp.float64, device=device)
-    X_wp = wp.array2d(X_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    wp.launch_tiled(
-        tile_math_cholesky, dim=[1, 1], inputs=[A_wp, D_wp, L_wp, Y_wp, X_wp], block_dim=TILE_DIM, device=device
-    )
-    wp.synchronize_device(device)
-
-    np.testing.assert_allclose(X_wp.numpy(), X_np)
-    np.testing.assert_allclose(L_wp.numpy(), L_np)
-
-    # TODO: implement and test backward pass
-
-
-@wp.kernel()
-def tile_math_cholesky_multiple_rhs(
-    gA: wp.array2d(dtype=wp.float64),
-    gD: wp.array1d(dtype=wp.float64),
-    gL: wp.array2d(dtype=wp.float64),
-    gy: wp.array2d(dtype=wp.float64),
-    gx: wp.array2d(dtype=wp.float64),
-    gz: wp.array2d(dtype=wp.float64),
-):
-    i, j = wp.tid()
-    # Load A, D & y
-    a = wp.tile_load(gA, shape=(TILE_M, TILE_M), storage="shared")
-    d = wp.tile_load(gD, shape=TILE_M, storage="shared")
-    y = wp.tile_load(gy, shape=(TILE_M, TILE_M), storage="shared")
-    # Ensure tile_diag_add() and tile_cholesky_solve() work with transposed matrices
-    a_t = wp.tile_transpose(a)
-    # Compute L st LL^T = A.T + diag(D)
-    b = wp.tile_diag_add(a_t, d)
-    l = wp.tile_cholesky(b)
-    # Solve for y in LL^T x = y.T
-    y_t = wp.tile_transpose(y)
-    x = wp.tile_cholesky_solve(l, y_t)
-    # Ensure matmul receives correct layout information
-    z = wp.tile_matmul(x, x)
-    # Store L & y
-    wp.tile_store(gL, l)
-    wp.tile_store(gx, x)
-    wp.tile_store(gz, z)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_cholesky_multiple_rhs(test, device):
-    A_h = np.ones((TILE_M, TILE_M), dtype=np.float64)
-    D_h = 8.0 * np.ones(TILE_M, dtype=np.float64)
-    L_h = np.zeros_like(A_h)
-    Y_h = np.arange((TILE_M, TILE_M), dtype=np.float64)
-    X_h = np.zeros_like(Y_h)
-    Z_h = np.zeros_like(Y_h)
-
-    A_np = A_h.T + np.diag(D_h)
-    L_np = np.linalg.cholesky(A_np)
-    X_np = np.linalg.solve(A_np, Y_h.T)
-    Z_np = X_np @ X_np
-
-    A_wp = wp.array2d(A_h, requires_grad=True, dtype=wp.float64, device=device)
-    D_wp = wp.array2d(D_h, requires_grad=True, dtype=wp.float64, device=device)
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    Y_wp = wp.array2d(Y_h, requires_grad=True, dtype=wp.float64, device=device)
-    X_wp = wp.array2d(X_h, requires_grad=True, dtype=wp.float64, device=device)
-    Z_wp = wp.array2d(Z_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    wp.launch_tiled(
-        tile_math_cholesky_multiple_rhs,
-        dim=[1, 1],
-        inputs=[A_wp, D_wp, L_wp, Y_wp, X_wp, Z_wp],
-        block_dim=TILE_DIM,
-        device=device,
-    )
-    wp.synchronize_device(device)
-
-    np.testing.assert_allclose(L_wp.numpy(), L_np)
-    np.testing.assert_allclose(X_wp.numpy(), X_np)
-    np.testing.assert_allclose(Z_wp.numpy(), Z_np)
-
-    # TODO: implement and test backward pass
-
-
-@wp.kernel
-def tile_math_forward_substitution(
-    gL: wp.array2d(dtype=wp.float64), gx: wp.array1d(dtype=wp.float64), gz: wp.array1d(dtype=wp.float64)
-):
-    i, j = wp.tid()
-    # Load L & x
-    L = wp.tile_load(gL, shape=(TILE_M, TILE_M), storage="shared")
-    x = wp.tile_load(gx, shape=TILE_M, storage="shared")
-    # Solve for z in Lz = x
-    # Transpose because we loaded an upper triangular matrix
-    z = wp.tile_lower_solve(wp.tile_transpose(L), x)
-    # Store z
-    wp.tile_store(gz, z)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_forward_substitution(test, device):
-    # Create test data
-    rng = np.random.default_rng(42)
-    L_h = np.triu(rng.random((TILE_M, TILE_M)))  # Upper triangular matrix
-    x_h = rng.random(TILE_M)
-    z_h = np.zeros_like(x_h)
-
-    # Compute reference solution using numpy
-    z_np = np.linalg.solve(L_h.T, x_h)
-
-    # Create Warp arrays
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    x_wp = wp.array1d(x_h, requires_grad=True, dtype=wp.float64, device=device)
-    z_wp = wp.array1d(z_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    # Run kernel
-    wp.launch_tiled(
-        tile_math_forward_substitution, dim=[1, 1], inputs=[L_wp, x_wp, z_wp], block_dim=TILE_DIM, device=device
-    )
-    wp.synchronize_device(device)
-
-    # Verify results
-    np.testing.assert_allclose(z_wp.numpy(), z_np)
-
-    # TODO: implement and test backward pass
-
-
-@wp.kernel
-def tile_math_back_substitution(
-    gL: wp.array2d(dtype=wp.float64), gx: wp.array1d(dtype=wp.float64), gz: wp.array1d(dtype=wp.float64)
-):
-    i, j = wp.tid()
-    # Load L & x
-    L = wp.tile_load(gL, shape=(TILE_M, TILE_M), storage="shared")
-    x = wp.tile_load(gx, shape=TILE_M, storage="shared")
-    # Solve for z in L^T z = x
-    # Transpose because we loaded a lower triangular matrix
-    z = wp.tile_upper_solve(wp.tile_transpose(L), x)
-    # Store z
-    wp.tile_store(gz, z)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_back_substitution(test, device):
-    # Create test data
-    rng = np.random.default_rng(42)
-    L_h = np.tril(rng.random((TILE_M, TILE_M)))  # Lower triangular matrix
-    x_h = rng.random(TILE_M)
-    z_h = np.zeros_like(x_h)
-
-    # Compute reference solution using numpy
-    z_np = np.linalg.solve(L_h.T, x_h)
-
-    # Create Warp arrays
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    x_wp = wp.array1d(x_h, requires_grad=True, dtype=wp.float64, device=device)
-    z_wp = wp.array1d(z_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    # Run kernel
-    wp.launch_tiled(
-        tile_math_back_substitution, dim=[1, 1], inputs=[L_wp, x_wp, z_wp], block_dim=TILE_DIM, device=device
-    )
-    wp.synchronize_device(device)
-
-    # Verify results
-    np.testing.assert_allclose(z_wp.numpy(), z_np)
-
-    # TODO: implement and test backward pass
-
-
-@wp.kernel
-def tile_math_forward_substitution_multiple_rhs(
-    gL: wp.array2d(dtype=wp.float64),
-    gx: wp.array2d(dtype=wp.float64),
-    gz: wp.array2d(dtype=wp.float64),
-    gc: wp.array2d(dtype=wp.float64),
-):
-    i, j = wp.tid()
-    # Load L & x
-    L = wp.tile_load(gL, shape=(TILE_M, TILE_M), storage="shared")
-    x = wp.tile_load(gx, shape=(TILE_M, TILE_M), storage="shared")
-    # Solve for z in Lz = x.T
-    x_t = wp.tile_transpose(x)
-    z = wp.tile_lower_solve(L, x_t)
-    # Ensure matmul receives correct layout information
-    c = wp.tile_matmul(z, z)
-    # Store z and c
-    wp.tile_store(gz, z)
-    wp.tile_store(gc, c)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_forward_substitution_multiple_rhs(test, device):
-    # Create test data
-    rng = np.random.default_rng(42)
-    L_h = np.tril(rng.random((TILE_M, TILE_M)))  # Lower triangular matrix
-    x_h = rng.random((TILE_M, TILE_M))  # Multiple right-hand sides
-    z_h = np.zeros_like(x_h)
-    c_h = np.zeros_like(x_h)
-
-    # Compute reference solution using numpy
-    z_np = np.linalg.solve(L_h, x_h.T)
-    c_np = z_np @ z_np
-
-    # Create Warp arrays
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    x_wp = wp.array2d(x_h, requires_grad=True, dtype=wp.float64, device=device)
-    z_wp = wp.array2d(z_h, requires_grad=True, dtype=wp.float64, device=device)
-    c_wp = wp.array2d(c_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    # Run kernel
-    wp.launch_tiled(
-        tile_math_forward_substitution_multiple_rhs,
-        dim=[1, 1],
-        inputs=[L_wp, x_wp, z_wp, c_wp],
-        block_dim=TILE_DIM,
-        device=device,
-    )
-    wp.synchronize_device()
-
-    # Verify results
-    assert np.allclose(z_wp.numpy(), z_np)
-    assert np.allclose(c_wp.numpy(), c_np)
-
-    # TODO: implement and test backward pass
-
-
-@wp.kernel
-def tile_math_back_substitution_multiple_rhs(
-    gL: wp.array2d(dtype=wp.float64),
-    gx: wp.array2d(dtype=wp.float64),
-    gz: wp.array2d(dtype=wp.float64),
-    gc: wp.array2d(dtype=wp.float64),
-):
-    i, j = wp.tid()
-    # Load L & x
-    L = wp.tile_load(gL, shape=(TILE_M, TILE_M), storage="shared")
-    x = wp.tile_load(gx, shape=(TILE_M, TILE_M), storage="shared")
-    # Solve for z in L^T z = x.T
-    x_t = wp.tile_transpose(x)
-    z = wp.tile_upper_solve(wp.tile_transpose(L), x_t)
-    # Ensure matmul receives correct layout information
-    c = wp.tile_matmul(z, z)
-    # Store z and c
-    wp.tile_store(gz, z)
-    wp.tile_store(gc, c)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_back_substitution_multiple_rhs(test, device):
-    # Create test data
-    rng = np.random.default_rng(42)
-    L_h = np.tril(rng.random((TILE_M, TILE_M)))  # Lower triangular matrix
-    x_h = rng.random((TILE_M, TILE_M))  # Multiple right-hand sides
-    z_h = np.zeros_like(x_h)
-    c_h = np.zeros_like(x_h)
-
-    # Compute reference solution using numpy
-    z_np = np.linalg.solve(L_h.T, x_h.T)
-    c_np = z_np @ z_np
-
-    # Create Warp arrays
-    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
-    x_wp = wp.array2d(x_h, requires_grad=True, dtype=wp.float64, device=device)
-    z_wp = wp.array2d(z_h, requires_grad=True, dtype=wp.float64, device=device)
-    c_wp = wp.array2d(c_h, requires_grad=True, dtype=wp.float64, device=device)
-
-    # Run kernel
-    wp.launch_tiled(
-        tile_math_back_substitution_multiple_rhs,
-        dim=[1, 1],
-        inputs=[L_wp, x_wp, z_wp, c_wp],
-        block_dim=TILE_DIM,
-        device=device,
-    )
-    wp.synchronize_device()
-
-    # Verify results
-    assert np.allclose(z_wp.numpy(), z_np)
-    assert np.allclose(c_wp.numpy(), c_np)
-
-    # TODO: implement and test backward pass
-
-
-# tests a complex composition of most libmathdx calls
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_block_cholesky(test, device):
-    BLOCK_SIZE = wp.constant(TILE_M // 2)
-
-    @wp.kernel(module="unique")
-    def block_cholesky_kernel(
-        A: wp.array2d(dtype=float),
-        L: wp.array2d(dtype=float),
-    ):
-        """
-        Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
-        It returns a lower-triangular matrix L such that A = L L^T.
-        """
-
-        # Process the matrix in blocks along its leading dimension.
-        for k in range(0, TILE_M, BLOCK_SIZE):
-            end = k + BLOCK_SIZE
-
-            # Load current diagonal block A[k:end, k:end]
-            # and update with contributions from previously computed blocks.
-            A_kk_tile = wp.tile_load(A, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(k, k), storage="shared")
-
-            for j in range(0, k, BLOCK_SIZE):
-                L_block = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(k, j))
-                L_block_T = wp.tile_transpose(L_block)
-                L_L_T_block = wp.tile_matmul(L_block, L_block_T)
-                A_kk_tile -= L_L_T_block
-
-            # Compute the Cholesky factorization for the block
-            # print(A_kk_tile)
-            L_kk_tile = wp.tile_cholesky(A_kk_tile)
-            wp.tile_store(L, L_kk_tile, offset=(k, k))
-
-            # Process the blocks below the current block
-            for i in range(end, TILE_M, BLOCK_SIZE):
-                A_ik_tile = wp.tile_load(A, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(i, k), storage="shared")
-
-                for j in range(0, k, BLOCK_SIZE):
-                    L_tile = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(i, j))
-                    L_2_tile = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(k, j))
-                    L_T_tile = wp.tile_transpose(L_2_tile)
-                    L_L_T_tile = wp.tile_matmul(L_tile, L_T_tile)
-                    A_ik_tile -= L_L_T_tile
-
-                A_ik_T_tile = wp.tile_transpose(A_ik_tile)
-                sol_T_tile = wp.tile_lower_solve(L_kk_tile, A_ik_T_tile)
-                sol_tile = wp.tile_transpose(sol_T_tile)
-
-                wp.tile_store(L, sol_tile, offset=(i, k))
-
-    @wp.kernel(module="unique")
-    def block_cholesky_solve_kernel(
-        L: wp.array2d(dtype=float),
-        b: wp.array2d(dtype=float),
-        scratch: wp.array2d(dtype=float),
-        x: wp.array2d(dtype=float),
-    ):
-        """
-        Solves A x = b given the Cholesky factor L (A = L L^T) using
-        blocked forward and backward substitution.
-        """
-
-        # Forward substitution: solve L y = b
-        for i in range(0, TILE_M, BLOCK_SIZE):
-            i_end = i + BLOCK_SIZE
-            rhs_tile = wp.tile_load(b, shape=(BLOCK_SIZE, 1), offset=(i, 0))
-            for j in range(0, i, BLOCK_SIZE):
-                L_block = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(i, j))
-                y_block = wp.tile_load(scratch, shape=(BLOCK_SIZE, 1), offset=(j, 0))
-                Ly_block = wp.tile_matmul(L_block, y_block)
-                rhs_tile -= Ly_block
-            L_tile = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(i, i))
-            y_tile = wp.tile_lower_solve(L_tile, rhs_tile)
-            wp.tile_store(scratch, y_tile, offset=(i, 0))
-
-        # Backward substitution: solve L^T x = y
-        for i in range(TILE_M - BLOCK_SIZE, -1, -BLOCK_SIZE):
-            i_start = i
-            i_end = i_start + BLOCK_SIZE
-            rhs_tile = wp.tile_load(scratch, shape=(BLOCK_SIZE, 1), offset=(i_start, 0))
-            for j in range(i_end, TILE_M, BLOCK_SIZE):
-                L_tile = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(j, i_start))
-                L_T_tile = wp.tile_transpose(L_tile)
-                x_tile = wp.tile_load(x, shape=(BLOCK_SIZE, 1), offset=(j, 0))
-                L_T_x_tile = wp.tile_matmul(L_T_tile, x_tile)
-                rhs_tile -= L_T_x_tile
-            L_tile = wp.tile_load(L, shape=(BLOCK_SIZE, BLOCK_SIZE), offset=(i_start, i_start))
-            x_tile = wp.tile_upper_solve(wp.tile_transpose(L_tile), rhs_tile)
-            wp.tile_store(x, x_tile, offset=(i_start, 0))
-
-    # check block cholesky decomposition
-
-    rng = np.random.default_rng(42)
-
-    M = np.array(rng.random((TILE_M, TILE_M)), dtype=float)
-
-    A_np = M.T @ M + np.eye(TILE_M, TILE_M)
-    L_np = np.linalg.cholesky(A_np)
-
-    A_wp = wp.array2d(A_np, dtype=float, device=device)
-    L_wp = wp.zeros_like(A_wp)
-
-    wp.launch_tiled(block_cholesky_kernel, dim=1, inputs=[A_wp], outputs=[L_wp], block_dim=TILE_DIM, device=device)
-
-    # check block cholesky solve
-
-    assert_np_equal(L_wp.numpy(), L_np, tol=1e-6)
-
-    b_np = np.array(rng.random((TILE_M, 1)), dtype=float)
-    b_wp = wp.array(b_np, dtype=float, device=device)
-
-    scratch = wp.zeros_like(b_wp)
-
-    x_np = np.linalg.solve(L_np.T, np.linalg.solve(L_np, b_np))
-    x_wp = wp.zeros_like(b_wp)
-
-    wp.launch_tiled(
-        block_cholesky_solve_kernel,
-        dim=1,
-        inputs=[L_wp, b_wp, scratch],
-        outputs=[x_wp],
-        block_dim=TILE_DIM,
-        device=device,
-    )
-
-    assert_np_equal(x_wp.numpy(), x_np, tol=1e-6)
-
-
-@wp.kernel
-def test_tile_lower_solve(L: wp.array2d(dtype=float), y: wp.array(dtype=float), x: wp.array(dtype=float)):
-    L_tile = wp.tile_load(L, shape=(TILE_M, TILE_M))
-    y_tile = wp.tile_load(x, shape=(TILE_M,))
-    sol = wp.tile_lower_solve(L_tile, y_tile)
-    wp.tile_store(x, sol)
-
-
-@wp.kernel
-def test_tile_upper_solve(L: wp.array2d(dtype=float), y: wp.array(dtype=float), x: wp.array(dtype=float)):
-    L_tile = wp.tile_load(L, shape=(TILE_M, TILE_M))
-    y_tile = wp.tile_load(x, shape=(TILE_M,))
-    sol = wp.tile_upper_solve(L_tile, y_tile)
-    wp.tile_store(x, sol)
-
-
-@unittest.skipUnless(wp.context.runtime.core.cuda_toolkit_version() >= 12060, "CUDA toolkit version is less than 12.6")
-def test_tile_math_singular_matrices(test, device):
-    rng = np.random.default_rng(42)
-    L_np = np.tril(rng.random((TILE_M, TILE_M)))  # Lower triangular matrix
-    L_np[-1, -1] = 0.0  # Make it singular
-    y_np = rng.random(TILE_M)
-
-    L_wp = wp.array2d(L_np, dtype=float, device=device)
-    y_wp = wp.array(y_np, dtype=float, device=device)
-    x_wp = wp.zeros_like(y_wp)
-
-    wp.launch_tiled(
-        test_tile_lower_solve, dim=1, inputs=[L_wp, y_wp], outputs=[x_wp], block_dim=TILE_DIM, device=device
-    )
-
-    assert np.isnan(x_wp.numpy()).any()
-
-    L_np = np.triu(rng.random((TILE_M, TILE_M)))  # Upper triangular matrix
-    L_np[-1, -1] = 0.0  # Make it singular
-
-    L_wp = wp.array2d(L_np, dtype=float, device=device)
-    y_wp = wp.array(y_np, dtype=float, device=device)
-    x_wp = wp.zeros_like(y_wp)
-
-    wp.launch_tiled(
-        test_tile_upper_solve, dim=1, inputs=[L_wp, y_wp], outputs=[x_wp], block_dim=TILE_DIM, device=device
-    )
-
-    assert np.isnan(x_wp.numpy()).any()
-
-
 all_devices = get_test_devices()
 cuda_devices = get_cuda_test_devices()
 
@@ -633,16 +135,6 @@ class TestTileMathDx(unittest.TestCase):
 add_function_test(
     TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=all_devices, check_output=False
 )
-add_function_test(
-    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=all_devices, check_output=False
-)
-add_function_test(
-    TestTileMathDx,
-    "tile_math_cholesky_multiple_rhs",
-    tile_math_cholesky_multiple_rhs,
-    devices=all_devices,
-    check_output=False,
-)
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2f",
@@ -658,54 +150,6 @@ add_function_test(
     check_output=False,
 )
 
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_forward_substitution",
-    test_tile_math_forward_substitution,
-    devices=cuda_devices,
-    check_output=False,
-)
-
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_back_substitution",
-    test_tile_math_back_substitution,
-    devices=cuda_devices,
-    check_output=False,
-)
-
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_forward_substitution_multiple_rhs",
-    test_tile_math_forward_substitution_multiple_rhs,
-    devices=cuda_devices,
-    check_output=False,
-)
-
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_back_substitution_multiple_rhs",
-    test_tile_math_back_substitution_multiple_rhs,
-    devices=cuda_devices,
-    check_output=False,
-)
-
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_block_cholesky",
-    test_tile_math_block_cholesky,
-    devices=cuda_devices,
-    check_output=False,
-)
-
-add_function_test(
-    TestTileMathDx,
-    "test_tile_math_singular_matrices",
-    test_tile_math_singular_matrices,
-    devices=cuda_devices,
-    check_output=False,
-)
-
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/tile/test_tile_matmul.py

@@ -159,7 +159,7 @@ def test_tile_transpose_matmul(test, device):
         test_tile_transpose_matmul_kernel, dim=[1], inputs=[input, output], block_dim=TILE_DIM, device=device
     )
 
-    assert_np_equal(output.numpy(), input.numpy().T @ input.numpy())
+    assert_np_equal(output.numpy(), input.numpy().T @ input.numpy(), 1e-6)
 
 
 class TestTileMatmul(unittest.TestCase):

warp/tests/tile/test_tile_mlp.py

@@ -43,7 +43,7 @@ def create_array(rng, dim_in, dim_hid, dtype=float):
 def test_multi_layer_nn(test, device):
     import torch as tc
 
-    if device.is_cuda and not wp.context.runtime.core.is_mathdx_enabled():
+    if device.is_cuda and not wp.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
 
     NUM_FREQ = wp.constant(8)

warp/tests/tile/test_tile_shared_memory.py

@@ -110,13 +110,13 @@ def test_tile_shared_mem_graph(test, device):
 
     out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
 
-    compute.module.load(device)
+    # preload the unique module
+    wp.load_module(compute.module, device=device, block_dim=BLOCK_DIM)
 
-    wp.capture_begin(device, force_module_load=False)
-    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
-    graph = wp.capture_end(device)
+    with wp.ScopedCapture(device, force_module_load=False) as capture:
+        wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
 
-    wp.capture_launch(graph)
+    wp.capture_launch(capture.graph)
 
     # check output
     assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)

warp/tests/unittest_suites.py

@@ -164,6 +164,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_linear_solvers import TestLinearSolvers
     from warp.tests.test_lvalue import TestLValue
     from warp.tests.test_mat import TestMat
+    from warp.tests.test_mat_constructors import TestMatConstructors
     from warp.tests.test_mat_lite import TestMatLite
     from warp.tests.test_mat_scalar_ops import TestMatScalarOps
     from warp.tests.test_math import TestMath
@@ -198,6 +199,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_vec_scalar_ops import TestVecScalarOps
     from warp.tests.test_verify_fp import TestVerifyFP
     from warp.tests.tile.test_tile import TestTile
+    from warp.tests.tile.test_tile_cholesky import TestTileCholesky
     from warp.tests.tile.test_tile_load import TestTileLoad
     from warp.tests.tile.test_tile_mathdx import TestTileMathDx
     from warp.tests.tile.test_tile_matmul import TestTileMatmul
@@ -261,6 +263,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestLValue,
         TestMarchingCubes,
         TestMat,
+        TestMatConstructors,
         TestMatLite,
         TestMatScalarOps,
         TestMath,
@@ -298,6 +301,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestStruct,
         TestTape,
         TestTile,
+        TestTileCholesky,
         TestTileLoad,
         TestTileMathDx,
         TestTileMatmul,
@@ -360,6 +364,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     from warp.tests.test_lvalue import TestLValue
     from warp.tests.test_mat_lite import TestMatLite
     from warp.tests.test_math import TestMath
+    from warp.tests.test_module_aot import TestModuleAOT
     from warp.tests.test_module_hashing import TestModuleHashing
     from warp.tests.test_modules_lite import TestModuleLite
     from warp.tests.test_noise import TestNoise
@@ -406,6 +411,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
+        TestModuleAOT,
         TestModuleHashing,
         TestModuleLite,
         TestNoise,

warp/tests/walkthrough_debug.py

@@ -68,7 +68,7 @@ wp.init()
 wp.config.mode = "debug"
 
 # Make sure Warp was built with `build_lib.py --mode=debug`
-assert wp.context.runtime.core.is_debug_enabled(), "Warp must be built in debug mode to enable debugging kernels"
+assert wp.context.runtime.core.wp_is_debug_enabled(), "Warp must be built in debug mode to enable debugging kernels"
 
 
 @wp.kernel