warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_lerp.py

@@ -5,14 +5,16 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
 from dataclasses import dataclass
 from typing import Any
-import unittest
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
+
+wp.init()
 
 
 @dataclass
@@ -162,8 +164,6 @@ TEST_DATA = {
     ),
 }
 
-wp.init()
-
 
 def test_lerp(test, device):
     def make_kernel_fn(data_type):
@@ -179,84 +179,83 @@ def test_lerp(test, device):
 
     for data_type in TEST_DATA:
         kernel_fn = make_kernel_fn(data_type)
-        module = wp.get_module(kernel_fn.__module__)
         kernel = wp.Kernel(
             func=kernel_fn,
             key=f"test_lerp_{data_type.__name__}_kernel",
-            module=module,
         )
 
-        for test_data in TEST_DATA[data_type]:
-            a = wp.array(
-                [test_data.a],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            b = wp.array(
-                [test_data.b],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            t = wp.array(
-                [test_data.t],
-                dtype=float,
-                device=device,
-                requires_grad=True,
-            )
-            out = wp.array(
-                [0] * wp.types.type_length(data_type),
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-
-            tape = wp.Tape()
-            with tape:
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[a, b, t, out],
+        with test.subTest(data_type=data_type):
+            for test_data in TEST_DATA[data_type]:
+                a = wp.array(
+                    [test_data.a],
+                    dtype=data_type,
                     device=device,
+                    requires_grad=True,
                 )
-
-            assert_np_equal(
-                out.numpy(),
-                np.array([test_data.expected]),
-                tol=1e-6,
-            )
-
-            if test_data.check_backwards():
-                tape.backward(out)
-
-                assert_np_equal(
-                    tape.gradients[a].numpy(),
-                    np.array([test_data.expected_adj_a]),
-                    tol=1e-6,
+                b = wp.array(
+                    [test_data.b],
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
-                assert_np_equal(
-                    tape.gradients[b].numpy(),
-                    np.array([test_data.expected_adj_b]),
-                    tol=1e-6,
+                t = wp.array(
+                    [test_data.t],
+                    dtype=float,
+                    device=device,
+                    requires_grad=True,
+                )
+                out = wp.array(
+                    [0] * wp.types.type_length(data_type),
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
+
+                tape = wp.Tape()
+                with tape:
+                    wp.launch(
+                        kernel,
+                        dim=1,
+                        inputs=[a, b, t, out],
+                        device=device,
+                    )
+
                 assert_np_equal(
-                    tape.gradients[t].numpy(),
-                    np.array([test_data.expected_adj_t]),
+                    out.numpy(),
+                    np.array([test_data.expected]),
                     tol=1e-6,
                 )
 
+                if test_data.check_backwards():
+                    tape.backward(out)
+
+                    assert_np_equal(
+                        tape.gradients[a].numpy(),
+                        np.array([test_data.expected_adj_a]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[b].numpy(),
+                        np.array([test_data.expected_adj_b]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[t].numpy(),
+                        np.array([test_data.expected_adj_t]),
+                        tol=1e-6,
+                    )
+
+
+devices = get_test_devices()
+
 
-def register(parent):
-    devices = get_test_devices()
+class TestLerp(unittest.TestCase):
+    pass
 
-    class TestLerp(parent):
-        pass
 
-    add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
-    return TestLerp
+add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
 
 
 if __name__ == "__main__":
-    _ = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_linear_solvers.py

@@ -0,0 +1,154 @@
+import unittest
+
+import numpy as np
+import warp as wp
+import unittest
+
+from warp.optim.linear import preconditioner, cg, bicgstab, gmres
+from warp.tests.unittest_utils import *
+
+
+wp.init()
+
+
+def _check_linear_solve(test, A, b, func, *args, **kwargs):
+    # test from zero
+    x = wp.zeros_like(b)
+    with wp.ScopedDevice(A.device):
+        niter, err, atol = func(A, b, x, *args, use_cuda_graph=True, **kwargs)
+
+    test.assertLessEqual(err, atol)
+
+    # test with warm start
+    with wp.ScopedDevice(A.device):
+        niter_warm, err, atol = func(A, b, x, *args, use_cuda_graph=False, **kwargs)
+
+    test.assertLessEqual(err, atol)
+    test.assertLess(niter_warm, niter)
+
+    # In CG and BiCGSTAB residual norm is evaluating from running residual
+    # rather then being computed from scratch as Ax - b
+    # This can lead to accumulated inaccuracies over iterations, esp in float32
+    residual = A.numpy() @ x.numpy() - b.numpy()
+    err_np = np.linalg.norm(residual)
+    if A.dtype == wp.float64:
+        test.assertLessEqual(err_np, 2.0 * atol)
+    else:
+        test.assertLessEqual(err_np, 32.0 * atol)
+
+
+def _least_square_system(rng, n: int):
+    C = rng.uniform(low=-100, high=100, size=(n, n))
+    f = rng.uniform(low=-100, high=100, size=(n,))
+
+    A = C @ C.T
+    b = C @ f
+
+    return A, b
+
+
+def _make_spd_system(n: int, seed: int, dtype, device):
+    rng = np.random.default_rng(seed)
+
+    A, b = _least_square_system(rng, n)
+
+    return wp.array(A, dtype=dtype, device=device), wp.array(b, dtype=dtype, device=device)
+
+
+def _make_nonsymmetric_system(n: int, seed: int, dtype, device, spd=False):
+    rng = np.random.default_rng(seed)
+    s = rng.uniform(low=0.1, high=10, size=(n,))
+
+    A, b = _least_square_system(rng, n)
+    A = A @ np.diag(s)
+
+    return wp.array(A, dtype=dtype, device=device), wp.array(b, dtype=dtype, device=device)
+
+
+def _make_indefinite_system(n: int, seed: int, dtype, device, spd=False):
+    rng = np.random.default_rng(seed)
+    s = rng.uniform(low=0.1, high=10, size=(n,))
+
+    A, b = _least_square_system(rng, n)
+    A = A @ np.diag(s)
+
+    return wp.array(A, dtype=dtype, device=device), wp.array(b, dtype=dtype, device=device)
+
+
+def test_cg(test, device):
+    A, b = _make_spd_system(n=64, seed=123, device=device, dtype=wp.float64)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, cg, maxiter=1000)
+    _check_linear_solve(test, A, b, cg, M=M, maxiter=1000)
+
+    A, b = _make_spd_system(n=16, seed=321, device=device, dtype=wp.float32)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, cg, maxiter=1000)
+    _check_linear_solve(test, A, b, cg, M=M, maxiter=1000)
+
+
+def test_bicgstab(test, device):
+    A, b = _make_nonsymmetric_system(n=64, seed=123, device=device, dtype=wp.float64)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, bicgstab, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000, is_left_preconditioner=True)
+
+    A, b = _make_nonsymmetric_system(n=16, seed=321, device=device, dtype=wp.float32)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, bicgstab, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000, is_left_preconditioner=True)
+
+    A, b = _make_indefinite_system(n=64, seed=121, device=device, dtype=wp.float64)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, bicgstab, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000)
+    _check_linear_solve(test, A, b, bicgstab, M=M, maxiter=1000, is_left_preconditioner=True)
+
+
+def test_gmres(test, device):
+    A, b = _make_nonsymmetric_system(n=64, seed=456, device=device, dtype=wp.float64)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, gmres, maxiter=1000, tol=1.0e-3)
+    _check_linear_solve(test, A, b, gmres, M=M, maxiter=1000, tol=1.0e-5)
+    _check_linear_solve(test, A, b, gmres, M=M, maxiter=1000, tol=1.0e-5, is_left_preconditioner=True)
+
+    A, b = _make_nonsymmetric_system(n=64, seed=654, device=device, dtype=wp.float64)
+    M = preconditioner(A, "diag")
+
+    _check_linear_solve(test, A, b, gmres, maxiter=1000, tol=1.0e-3)
+    _check_linear_solve(test, A, b, gmres, M=M, maxiter=1000, tol=1.0e-5)
+    _check_linear_solve(test, A, b, gmres, M=M, maxiter=1000, tol=1.0e-5, is_left_preconditioner=True)
+
+
+class TestLinearSolvers(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+from warp.context import runtime
+
+if not runtime.core.is_cutlass_enabled():
+    devices = [d for d in devices if not d.is_cuda]
+    print("Skipping CUDA linear solver tests because CUTLASS is not supported in this build")
+
+if runtime.core.is_debug_enabled():
+    # cutlass-based matmul is *very* slow in debug mode -- skip
+    devices = [d for d in devices if not d.is_cuda]
+    print("Skipping CUDA linear solver tests in debug mode")
+
+add_function_test(TestLinearSolvers, "test_cg", test_cg, devices=devices)
+add_function_test(TestLinearSolvers, "test_bicgstab", test_bicgstab, devices=devices)
+add_function_test(TestLinearSolvers, "test_gmres", test_gmres, devices=devices)
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)