warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/test_linear_solvers.py

@@ -0,0 +1,223 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp._src.context import assert_conditional_graph_support
+from warp.optim.linear import bicgstab, cg, cr, gmres, preconditioner
+from warp.tests.unittest_utils import *
+
+
+def check_conditional_graph_support():
+    try:
+        assert_conditional_graph_support()
+    except RuntimeError:
+        return False
+    return True
+
+
+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 capturable graph
+    if A.device.is_cuda and check_conditional_graph_support():
+        x.zero_()
+        with wp.ScopedDevice(A.device):
+            with wp.ScopedCapture() as capture:
+                niter, err, atol = func(A, b, x, *args, use_cuda_graph=True, check_every=0, **kwargs)
+
+            wp.capture_launch(capture.graph)
+
+        niter = niter.numpy()[0]
+        err = np.sqrt(err.numpy()[0])
+        atol = np.sqrt(atol.numpy()[0])
+
+        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)
+
+    if isinstance(niter_warm, wp.array):
+        niter_warm = niter_warm.numpy()[0]
+        err = np.sqrt(err.numpy()[0])
+        atol = np.sqrt(atol.numpy()[0])
+
+    test.assertLessEqual(err, atol)
+
+    if func in [cr, gmres]:
+        # monotonic convergence
+        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):
+    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):
+    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_identity_system(n: int, seed: int, dtype, device):
+    rng = np.random.default_rng(seed)
+
+    A = np.eye(n)
+    b = rng.uniform(low=-1.0, high=1.0, size=(n,))
+
+    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)
+
+    A, b = _make_identity_system(n=5, seed=321, device=device, dtype=wp.float32)
+    _check_linear_solve(test, A, b, cg, maxiter=30)
+
+
+def test_cr(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, cr, maxiter=1000)
+    _check_linear_solve(test, A, b, cr, 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, cr, maxiter=1000)
+    _check_linear_solve(test, A, b, cr, M=M, maxiter=1000)
+
+    A, b = _make_identity_system(n=5, seed=321, device=device, dtype=wp.float32)
+    _check_linear_solve(test, A, b, cr, maxiter=30)
+
+
+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)
+
+    A, b = _make_identity_system(n=5, seed=321, device=device, dtype=wp.float32)
+    _check_linear_solve(test, A, b, bicgstab, maxiter=30)
+
+
+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)
+
+    A, b = _make_identity_system(n=5, seed=123, device=device, dtype=wp.float32)
+    _check_linear_solve(test, A, b, gmres, maxiter=120)
+
+
+class TestLinearSolvers(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestLinearSolvers, "test_cg", test_cg, devices=devices)
+add_function_test(TestLinearSolvers, "test_cr", test_cr, 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.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_lvalue.py

@@ -0,0 +1,427 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def rmw_array_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    foos[i] += wp.uint32(1)
+
+
+def test_rmw_array(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=rmw_array_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.struct
+class RmwFoo:
+    field: wp.uint32
+
+
+@wp.kernel
+def rmw_array_struct_kernel(foos: wp.array(dtype=RmwFoo)):
+    i = wp.tid()
+    foos[i].field += wp.uint32(1)
+
+
+def test_rmw_array_struct(test, device):
+    foos = wp.zeros((10,), dtype=RmwFoo, device=device)
+
+    wp.launch(kernel=rmw_array_struct_kernel, dim=(10,), inputs=[foos], device=device)
+
+    expected = RmwFoo()
+    expected.field = 1
+    for f in foos.list():
+        test.assertEqual(f.field, expected.field)
+
+
+@wp.func
+def lookup(foos: wp.array(dtype=wp.uint32), index: int):
+    return foos[index]
+
+
+@wp.kernel
+def lookup_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    x = lookup(foos, i)
+    foos[i] = x + wp.uint32(1)
+
+
+def test_lookup(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=lookup_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.func
+def lookup3(foos: wp.array(dtype=wp.float32), index: int):
+    return foos[index]
+
+
+@wp.kernel
+def grad_kernel(foos: wp.array(dtype=wp.float32), bars: wp.array(dtype=wp.float32)):
+    i = wp.tid()
+
+    x = lookup3(foos, i)
+    bars[i] = x * wp.float32(i) + 1.0
+
+
+def test_grad(test, device):
+    num = 10
+    data = np.linspace(20, 20 + num, num, endpoint=False, dtype=np.float32)
+    input = wp.array(data, device=device, requires_grad=True)
+    output = wp.zeros(num, dtype=wp.float32, device=device)
+
+    ones = wp.array(np.ones(len(output)), dtype=wp.float32, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel=grad_kernel, dim=(num,), inputs=[input], outputs=[output], device=device)
+
+    tape.backward(grads={output: ones})
+
+    # test forward results
+    for i, f in enumerate(output.list()):
+        test.assertEqual(f, data[i] * i + 1)
+
+    # test backward results
+    for i, f in enumerate(tape.gradients[input].list()):
+        test.assertEqual(f, i)
+
+
+@wp.func
+def lookup2(foos: wp.array(dtype=wp.uint32), index: int):
+    if index % 2 == 0:
+        x = foos[index]
+        x = wp.uint32(0)
+        return x
+    else:
+        return foos[index]
+
+
+@wp.kernel
+def lookup2_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    x = lookup2(foos, i)
+    foos[i] = x + wp.uint32(1)
+
+
+def test_lookup2(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=lookup2_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.kernel
+def unary_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    foos[i] = wp.uint32(-1)
+    x = -foos[i]
+    foos[i] = x
+
+
+def test_unary(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=unary_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.kernel
+def rvalue_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    if foos[i] < wp.uint32(1):
+        foos[i] = wp.uint32(1)
+
+
+def test_rvalue(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=rvalue_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+# Tests, among other things, that assigning a reference to a new variable does
+# not create a reference
+@wp.kernel
+def intermediate_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+
+    x = foos[i]
+    x = x + wp.uint32(1)
+    foos[i] = x
+
+
+def test_intermediate(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=intermediate_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.kernel
+def array_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+    foos[i] = wp.uint32(1)
+
+
+def test_array_assign(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=array_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.func
+def increment(arg: wp.uint32):
+    return arg + wp.uint32(1)
+
+
+@wp.kernel
+def array_call_kernel(foos: wp.array(dtype=wp.uint32)):
+    i = wp.tid()
+    foos[i] = increment(foos[i])
+
+
+def test_array_call_assign(test, device):
+    arr = wp.zeros((10,), dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=array_kernel, dim=(10,), inputs=[arr], device=device)
+
+    assert_np_equal(arr.numpy(), np.ones(10))
+
+
+@wp.struct
+class Foo:
+    field: wp.uint32
+
+
+@wp.kernel
+def array_struct_kernel(foos: wp.array(dtype=Foo)):
+    i = wp.tid()
+    foos[i].field = wp.uint32(1)
+
+
+def test_array_struct_assign(test, device):
+    foos = wp.zeros((10,), dtype=Foo, device=device)
+
+    wp.launch(kernel=array_struct_kernel, dim=(10,), inputs=[foos], device=device)
+
+    expected = Foo()
+    expected.field = 1
+
+    test.assertEqual(expected.field, 1)
+
+    for f in foos.list():
+        test.assertEqual(f.field, 1)
+
+
+@wp.struct
+class Bar:
+    field: wp.uint32
+
+
+@wp.struct
+class Baz:
+    bar: Bar
+
+
+@wp.kernel
+def array_struct_struct_kernel(foos: wp.array(dtype=Baz)):
+    i = wp.tid()
+    foos[i].bar.field = wp.uint32(1)
+
+
+def test_array_struct_struct_assign(test, device):
+    foos = wp.zeros((10,), dtype=Baz, device=device)
+
+    wp.launch(kernel=array_struct_struct_kernel, dim=(10,), inputs=[foos], device=device)
+
+    expected = Baz()
+    expected.bar.field = 1
+
+    test.assertEqual(expected.bar.field, 1)
+
+    for f in foos.list():
+        test.assertEqual(f.bar.field, 1)
+
+
+@wp.struct
+class S:
+    a: wp.uint32
+    b: wp.float32
+
+
+@wp.struct
+class F:
+    x: wp.float32
+    s: S
+    y: wp.int32
+
+
+@wp.kernel
+def complex_kernel(foos: wp.array(dtype=F)):
+    i = wp.tid()
+    foos[i].x += wp.float32(1.0)
+    foos[i].y = wp.int32(2)
+    foos[i].s.b += wp.float32(3.0)
+    foos[i].s.a = wp.uint32(foos[i].y)
+
+
+def test_complex(test, device):
+    foos = wp.zeros((10,), dtype=F, device=device)
+
+    wp.launch(kernel=complex_kernel, dim=(10,), inputs=[foos], device=device)
+
+    expected = F()
+    expected.x = 1.0
+    expected.y = 2
+    expected.s.b = 3.0
+    expected.s.a = expected.y
+    for f in foos.list():
+        test.assertEqual(f.x, expected.x)
+        test.assertEqual(f.y, expected.y)
+        test.assertEqual(f.s.a, expected.s.a)
+        test.assertEqual(f.s.b, expected.s.b)
+
+
+@wp.struct
+class Svec:
+    a: wp.uint32
+    b: wp.vec2f
+
+
+@wp.struct
+class Fvec:
+    x: wp.vec2f
+    s: Svec
+    y: wp.int32
+
+
+@wp.kernel
+def swizzle_kernel(foos: wp.array(dtype=Fvec)):
+    i = wp.tid()
+
+    foos[i].x += wp.vec2f(1.0, 2.0)
+    foos[i].y = wp.int32(3)
+    foos[i].s.b = wp.vec2f(4.0, 5.0)
+    foos[i].s.b.y = wp.float32(6.0)
+    foos[i].s.b.x = foos[i].x.y
+    foos[i].s.a = wp.uint32(foos[i].y)
+
+
+def test_swizzle(test, device):
+    foos = wp.zeros((10,), dtype=Fvec, device=device)
+
+    wp.launch(kernel=swizzle_kernel, dim=(10,), inputs=[foos], device=device)
+
+    expected = Fvec()
+    expected.x = wp.vec2f(1.0, 2.0)
+    expected.y = 3
+    expected.s.b = wp.vec2f(4.0, 5.0)
+    expected.s.b.y = 6.0
+    expected.s.b.x = expected.x.y
+    expected.s.a = expected.y
+
+    for f in foos.list():
+        test.assertEqual(f.x, expected.x)
+        test.assertEqual(f.y, expected.y)
+        test.assertEqual(f.s.a, expected.s.a)
+        test.assertEqual(f.s.b, expected.s.b)
+
+
+@wp.kernel
+def slice_kernel(a: wp.array2d(dtype=wp.vec3), b: wp.array2d(dtype=wp.vec3), c: wp.array2d(dtype=wp.vec3)):
+    tid = wp.tid()
+    c[tid][0] = a[tid][0] + b[tid][0]
+
+
+def test_slice(test, device):
+    a = wp.full((1, 1), value=1.0, dtype=wp.vec3, requires_grad=True, device=device)
+    b = wp.full((1, 1), value=1.0, dtype=wp.vec3, requires_grad=True, device=device)
+    c = wp.zeros((1, 1), dtype=wp.vec3, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel=slice_kernel, dim=1, inputs=[a, b], outputs=[c], device=device)
+
+    c.grad = wp.full((1, 1), value=1.0, dtype=wp.vec3, device=device)
+    tape.backward()
+
+    x = a.grad.list()[0]
+    y = b.grad.list()[0]
+
+    expected = wp.vec3(1.0)
+    test.assertEqual(x, expected)
+    test.assertEqual(y, expected)
+
+
+devices = get_test_devices()
+
+
+class TestLValue(unittest.TestCase):
+    def test_swizzle_error_invalid_attribute(self):
+        v = wp.vec3(1, 2, 3)
+        with self.assertRaisesRegex(AttributeError, r"'vec3f' object has no attribute 'foo'$"):
+            v.foo  # noqa: B018
+
+        try:
+            v.bar = 123
+        except AttributeError:
+            self.fail()
+
+
+add_function_test(TestLValue, "test_rmw_array", test_rmw_array, devices=devices)
+add_function_test(TestLValue, "test_rmw_array_struct", test_rmw_array_struct, devices=devices)
+add_function_test(TestLValue, "test_lookup", test_lookup, devices=devices)
+add_function_test(TestLValue, "test_lookup2", test_lookup2, devices=devices)
+add_function_test(TestLValue, "test_grad", test_grad, devices=devices)
+add_function_test(TestLValue, "test_unary", test_unary, devices=devices)
+add_function_test(TestLValue, "test_rvalue", test_rvalue, devices=devices)
+add_function_test(TestLValue, "test_intermediate", test_intermediate, devices=devices)
+add_function_test(TestLValue, "test_array_assign", test_array_assign, devices=devices)
+add_function_test(TestLValue, "test_array_struct_assign", test_array_struct_assign, devices=devices)
+add_function_test(TestLValue, "test_array_struct_struct_assign", test_array_struct_struct_assign, devices=devices)
+add_function_test(TestLValue, "test_complex", test_complex, devices=devices)
+add_function_test(TestLValue, "test_swizzle", test_swizzle, devices=devices)
+add_function_test(TestLValue, "test_slice", test_slice, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)