warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/tests/test_multigpu.py

@@ -5,16 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
-import warp as wp
+import unittest
 
-import math
+import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
-
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -31,157 +27,134 @@ def arange(start: int, step: int, a: wp.array(dtype=int)):
     a[tid] = start + step * tid
 
 
-def test_multigpu_set_device(test, device):
-    assert len(wp.get_cuda_devices()) > 1, "At least two CUDA devices are required"
-
-    # save default device
-    saved_device = wp.get_device()
-
-    n = 32
-
-    wp.set_device("cuda:0")
-    a0 = wp.empty(n, dtype=int)
-    wp.launch(arange, dim=a0.size, inputs=[0, 1, a0])
-
-    wp.set_device("cuda:1")
-    a1 = wp.empty(n, dtype=int)
-    wp.launch(arange, dim=a1.size, inputs=[0, 1, a1])
-
-    # restore default device
-    wp.set_device(saved_device)
+class TestMultiGPU(unittest.TestCase):
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_multigpu_set_device(self):
+        # save default device
+        saved_device = wp.get_device()
 
-    assert a0.device == "cuda:0"
-    assert a1.device == "cuda:1"
+        n = 32
 
-    expected = np.arange(n, dtype=int)
-
-    assert_np_equal(a0.numpy(), expected)
-    assert_np_equal(a1.numpy(), expected)
-
-
-def test_multigpu_scoped_device(test, device):
-    assert len(wp.get_cuda_devices()) > 1, "At least two CUDA devices are required"
-
-    n = 32
-
-    with wp.ScopedDevice("cuda:0"):
+        wp.set_device("cuda:0")
         a0 = wp.empty(n, dtype=int)
         wp.launch(arange, dim=a0.size, inputs=[0, 1, a0])
 
-    with wp.ScopedDevice("cuda:1"):
+        wp.set_device("cuda:1")
         a1 = wp.empty(n, dtype=int)
         wp.launch(arange, dim=a1.size, inputs=[0, 1, a1])
 
-    assert a0.device == "cuda:0"
-    assert a1.device == "cuda:1"
-
-    expected = np.arange(n, dtype=int)
+        # restore default device
+        wp.set_device(saved_device)
 
-    assert_np_equal(a0.numpy(), expected)
-    assert_np_equal(a1.numpy(), expected)
+        assert a0.device == "cuda:0"
+        assert a1.device == "cuda:1"
 
+        expected = np.arange(n, dtype=int)
 
-def test_multigpu_nesting(test, device):
-    assert len(wp.get_cuda_devices()) > 1, "At least two CUDA devices are required"
+        assert_np_equal(a0.numpy(), expected)
+        assert_np_equal(a1.numpy(), expected)
 
-    initial_device = wp.get_device()
-    initial_cuda_device = wp.get_cuda_device()
-
-    with wp.ScopedDevice("cuda:1"):
-        assert wp.get_device() == "cuda:1"
-        assert wp.get_cuda_device() == "cuda:1"
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_multigpu_scoped_device(self):
+        n = 32
 
         with wp.ScopedDevice("cuda:0"):
-            assert wp.get_device() == "cuda:0"
-            assert wp.get_cuda_device() == "cuda:0"
-
-            with wp.ScopedDevice("cpu"):
-                assert wp.get_device() == "cpu"
-                assert wp.get_cuda_device() == "cuda:0"
+            a0 = wp.empty(n, dtype=int)
+            wp.launch(arange, dim=a0.size, inputs=[0, 1, a0])
 
-                wp.set_device("cuda:1")
+        with wp.ScopedDevice("cuda:1"):
+            a1 = wp.empty(n, dtype=int)
+            wp.launch(arange, dim=a1.size, inputs=[0, 1, a1])
 
-                assert wp.get_device() == "cuda:1"
-                assert wp.get_cuda_device() == "cuda:1"
+        assert a0.device == "cuda:0"
+        assert a1.device == "cuda:1"
 
-            assert wp.get_device() == "cuda:0"
-            assert wp.get_cuda_device() == "cuda:0"
+        expected = np.arange(n, dtype=int)
 
-        assert wp.get_device() == "cuda:1"
-        assert wp.get_cuda_device() == "cuda:1"
+        assert_np_equal(a0.numpy(), expected)
+        assert_np_equal(a1.numpy(), expected)
 
-    assert wp.get_device() == initial_device
-    assert wp.get_cuda_device() == initial_cuda_device
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_multigpu_nesting(self):
+        initial_device = wp.get_device()
+        initial_cuda_device = wp.get_cuda_device()
 
+        with wp.ScopedDevice("cuda:1"):
+            assert wp.get_device() == "cuda:1"
+            assert wp.get_cuda_device() == "cuda:1"
 
-def test_multigpu_pingpong(test, device):
-    assert len(wp.get_cuda_devices()) > 1, "At least two CUDA devices are required"
+            with wp.ScopedDevice("cuda:0"):
+                assert wp.get_device() == "cuda:0"
+                assert wp.get_cuda_device() == "cuda:0"
 
-    n = 1024 * 1024
+                with wp.ScopedDevice("cpu"):
+                    assert wp.get_device() == "cpu"
+                    assert wp.get_cuda_device() == "cuda:0"
 
-    a0 = wp.zeros(n, dtype=float, device="cuda:0")
-    a1 = wp.zeros(n, dtype=float, device="cuda:1")
+                    wp.set_device("cuda:1")
 
-    iters = 10
+                    assert wp.get_device() == "cuda:1"
+                    assert wp.get_cuda_device() == "cuda:1"
 
-    for _ in range(iters):
-        wp.launch(inc, dim=a0.size, inputs=[a0], device=a0.device)
-        wp.synchronize_device(a0.device)
-        wp.copy(a1, a0)
+                assert wp.get_device() == "cuda:0"
+                assert wp.get_cuda_device() == "cuda:0"
 
-        wp.launch(inc, dim=a1.size, inputs=[a1], device=a1.device)
-        wp.synchronize_device(a1.device)
-        wp.copy(a0, a1)
+            assert wp.get_device() == "cuda:1"
+            assert wp.get_cuda_device() == "cuda:1"
 
-    expected = np.full(n, iters * 2, dtype=np.float32)
+        assert wp.get_device() == initial_device
+        assert wp.get_cuda_device() == initial_cuda_device
 
-    assert_np_equal(a0.numpy(), expected)
-    assert_np_equal(a1.numpy(), expected)
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_multigpu_pingpong(self):
+        n = 1024 * 1024
 
+        a0 = wp.zeros(n, dtype=float, device="cuda:0")
+        a1 = wp.zeros(n, dtype=float, device="cuda:1")
 
-def test_multigpu_pingpong_streams(test, device):
-    assert len(wp.get_cuda_devices()) > 1, "At least two CUDA devices are required"
+        iters = 10
 
-    n = 1024 * 1024
+        for _ in range(iters):
+            wp.launch(inc, dim=a0.size, inputs=[a0], device=a0.device)
+            wp.synchronize_device(a0.device)
+            wp.copy(a1, a0)
 
-    a0 = wp.zeros(n, dtype=float, device="cuda:0")
-    a1 = wp.zeros(n, dtype=float, device="cuda:1")
+            wp.launch(inc, dim=a1.size, inputs=[a1], device=a1.device)
+            wp.synchronize_device(a1.device)
+            wp.copy(a0, a1)
 
-    stream0 = wp.get_stream("cuda:0")
-    stream1 = wp.get_stream("cuda:1")
+        expected = np.full(n, iters * 2, dtype=np.float32)
 
-    iters = 10
+        assert_np_equal(a0.numpy(), expected)
+        assert_np_equal(a1.numpy(), expected)
 
-    for _ in range(iters):
-        wp.launch(inc, dim=a0.size, inputs=[a0], stream=stream0)
-        stream1.wait_stream(stream0)
-        wp.copy(a1, a0, stream=stream1)
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_multigpu_pingpong_streams(self):
+        n = 1024 * 1024
 
-        wp.launch(inc, dim=a1.size, inputs=[a1], stream=stream1)
-        stream0.wait_stream(stream1)
-        wp.copy(a0, a1, stream=stream0)
+        a0 = wp.zeros(n, dtype=float, device="cuda:0")
+        a1 = wp.zeros(n, dtype=float, device="cuda:1")
 
-    expected = np.full(n, iters * 2, dtype=np.float32)
+        stream0 = wp.get_stream("cuda:0")
+        stream1 = wp.get_stream("cuda:1")
 
-    assert_np_equal(a0.numpy(), expected)
-    assert_np_equal(a1.numpy(), expected)
+        iters = 10
 
+        for _ in range(iters):
+            wp.launch(inc, dim=a0.size, inputs=[a0], stream=stream0)
+            stream1.wait_stream(stream0)
+            wp.copy(a1, a0, stream=stream1)
 
-def register(parent):
-    class TestMultigpu(parent):
-        pass
+            wp.launch(inc, dim=a1.size, inputs=[a1], stream=stream1)
+            stream0.wait_stream(stream1)
+            wp.copy(a0, a1, stream=stream0)
 
-    if wp.get_cuda_device_count() > 1:
-        add_function_test(TestMultigpu, "test_multigpu_set_device", test_multigpu_set_device)
-        add_function_test(TestMultigpu, "test_multigpu_scoped_device", test_multigpu_scoped_device)
-        add_function_test(TestMultigpu, "test_multigpu_nesting", test_multigpu_nesting)
-        add_function_test(TestMultigpu, "test_multigpu_pingpong", test_multigpu_pingpong)
-        add_function_test(TestMultigpu, "test_multigpu_pingpong_streams", test_multigpu_pingpong_streams)
+        expected = np.full(n, iters * 2, dtype=np.float32)
 
-    return TestMultigpu
+        assert_np_equal(a0.numpy(), expected)
+        assert_np_equal(a1.numpy(), expected)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_noise.py

@@ -7,11 +7,11 @@
 
 import unittest
 
-import warp as wp
-from warp.tests.test_base import *
-
 import numpy as np
 
+import warp as wp
+from warp.tests.unittest_utils import *
+
 # import matplotlib.pyplot as plt
 
 wp.init()
@@ -231,20 +231,18 @@ def test_adj_noise(test, device):
     test.assertTrue(err < 1.0e-8)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestNoise(parent):
-        pass
+class TestNoise(unittest.TestCase):
+    pass
 
-    add_function_test(TestNoise, "test_pnoise", test_pnoise, devices=devices)
-    add_function_test(TestNoise, "test_curlnoise", test_curlnoise, devices=devices)
-    add_function_test(TestNoise, "test_adj_noise", test_adj_noise, devices=devices)
 
-    return TestNoise
+add_function_test(TestNoise, "test_pnoise", test_pnoise, devices=devices)
+add_function_test(TestNoise, "test_curlnoise", test_curlnoise, devices=devices)
+add_function_test(TestNoise, "test_adj_noise", test_adj_noise, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_operators.py

@@ -7,11 +7,8 @@
 
 import unittest
 
-import numpy as np
-import math
-
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -229,29 +226,25 @@ def test_operators_mat44():
     expect_eq(r0[3], wp.vec4(39.0, 42.0, 45.0, 48.0))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestOperators(parent):
-        pass
+class TestOperators(unittest.TestCase):
+    pass
 
-    add_kernel_test(TestOperators, test_operators_scalar_float, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_scalar_int, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_matrix_index, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_vector_index, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_vec3, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_vec4, dim=1, devices=devices)
 
-    add_kernel_test(TestOperators, test_operators_mat22, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_mat33, dim=1, devices=devices)
-    add_kernel_test(TestOperators, test_operators_mat44, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_scalar_float, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_scalar_int, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_matrix_index, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_vector_index, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_vec3, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_vec4, dim=1, devices=devices)
 
-    return TestOperators
+add_kernel_test(TestOperators, test_operators_mat22, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_mat33, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_mat44, dim=1, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    wp.force_load()
-
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_options.py

@@ -8,7 +8,7 @@
 import unittest
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -93,20 +93,19 @@ def test_options_4(test, device):
     assert_np_equal(tape.gradients[x].numpy(), np.array(0.0))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestOptions(parent):
-        pass
 
-    add_function_test(TestOptions, "test_options_1", test_options_1, devices=devices)
-    add_function_test(TestOptions, "test_options_2", test_options_2, devices=devices)
-    add_function_test(TestOptions, "test_options_3", test_options_3, devices=devices)
-    add_function_test(TestOptions, "test_options_4", test_options_4, devices=devices)
-    return TestOptions
+class TestOptions(unittest.TestCase):
+    pass
+
+
+add_function_test(TestOptions, "test_options_1", test_options_1, devices=devices)
+add_function_test(TestOptions, "test_options_2", test_options_2, devices=devices)
+add_function_test(TestOptions, "test_options_3", test_options_3, devices=devices)
+add_function_test(TestOptions, "test_options_4", test_options_4, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_pinned.py

@@ -5,16 +5,17 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
+
 import numpy as np
-import warp as wp
-from warp.tests.test_base import *
 
-import unittest
+import warp as wp
+from warp.tests.unittest_utils import *
 
 wp.init()
 
 
-def test_pinned(test, device):
+def test_pinned(test: unittest.TestCase, device):
     assert wp.get_device(device).is_cuda, "Test device must be a CUDA device"
 
     n = 1024 * 1024
@@ -25,20 +26,20 @@ def test_pinned(test, device):
     a_pageable1 = wp.array(ones, dtype=float, device="cpu")
     a_pageable2 = wp.zeros_like(a_pageable1)
 
-    assert a_pageable1.pinned == False
-    assert a_pageable2.pinned == False
+    test.assertFalse(a_pageable1.pinned)
+    test.assertFalse(a_pageable2.pinned)
 
     # pinned host arrays for asynchronous transfers
     a_pinned1 = wp.array(ones, dtype=float, device="cpu", pinned=True)
     a_pinned2 = wp.zeros_like(a_pinned1)
 
-    assert a_pinned1.pinned == True
-    assert a_pinned2.pinned == True
+    test.assertTrue(a_pinned1.pinned)
+    test.assertTrue(a_pinned2.pinned)
 
     # device array
     a_device = wp.zeros(n, dtype=float, device=device)
 
-    assert a_device.pinned == False
+    test.assertFalse(a_device.pinned)
 
     wp.synchronize_device(device)
 
@@ -59,22 +60,19 @@ def test_pinned(test, device):
     assert_np_equal(a_pinned2.numpy(), ones)
 
     # ensure that launching asynchronous transfers took less CPU time
-    assert pinned_timer.elapsed < pageable_timer.elapsed, "Pinned transfers did not take less CPU time"
+    test.assertTrue(pinned_timer.elapsed < pageable_timer.elapsed, "Pinned transfers did not take less CPU time")
+
 
+devices = get_unique_cuda_test_devices()
 
-def register(parent):
-    cuda_devices = wp.get_cuda_devices()
 
-    class TestPinned(parent):
-        pass
+class TestPinned(unittest.TestCase):
+    pass
 
-    if cuda_devices:
-        add_function_test(TestPinned, "test_pinned", test_pinned, devices=cuda_devices)
 
-    return TestPinned
+add_function_test(TestPinned, "test_pinned", test_pinned, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_print.py

@@ -5,13 +5,10 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import os
 import unittest
 
 import warp as wp
-
-from warp.tests.test_base import *
-
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -25,33 +22,32 @@ def test_print_kernel():
 
 
 def test_print(test, device):
+    wp.load_module(device=device)
     capture = StdOutCapture()
     capture.begin()
     wp.launch(kernel=test_print_kernel, dim=1, inputs=[], device=device)
-    wp.synchronize()
+    wp.synchronize_device(device)
     s = capture.end()
 
-    test.assertRegex(
-        s,
-        rf"1{os.linesep}"
-        rf"this is a string{os.linesep}"
-        rf"this is a float 457\.500000{os.linesep}"
-        rf"this is an int 123",
-    )
+    # The CPU kernel printouts don't get captured by StdOutCapture()
+    if device.is_cuda:
+        test.assertRegex(
+            s,
+            rf"1{os.linesep}"
+            rf"this is a string{os.linesep}"
+            rf"this is a float 457\.500000{os.linesep}"
+            rf"this is an int 123",
+        )
 
 
-def register(parent):
-    devices = get_test_devices()
-    devices = tuple(x for x in devices if not x.is_cpu)
+class TestPrint(unittest.TestCase):
+    pass
 
-    class TestPrint(parent):
-        pass
 
-    add_function_test(TestPrint, "test_print", test_print, devices=devices, check_output=False)
-    return TestPrint
+devices = get_test_devices()
+add_function_test(TestPrint, "test_print", test_print, devices=devices, check_output=False)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)