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

warp/tests/test_mlp.py

@@ -5,9 +5,12 @@
 # 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 *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -35,17 +38,17 @@ def loss_kernel(x: wp.array2d(dtype=float), loss: wp.array(dtype=float)):
 
 
 def test_mlp(test, device):
-    np.random.seed(0)
+    rng = np.random.default_rng(123)
 
     m = 10
     n = 200
 
     batches = 20000
 
-    weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device)
-    bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device)
+    weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+    bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
 
-    x = wp.array(np.random.rand(n, batches), dtype=float, device=device)
+    x = wp.array(rng.random(size=(n, batches)), dtype=float, device=device)
     y = wp.zeros(shape=(m, batches), device=device)
 
     with wp.ScopedTimer("warp", active=False):
@@ -86,13 +89,15 @@ def create_mlp(m, n):
 def create_golden():
     import torch
 
+    rng = np.random.default_rng(123)
+
     input_size = 32
     hidden_size = 16
     batch_size = 64
 
     network = create_mlp(input_size, hidden_size)
 
-    x = torch.Tensor(np.random.rand(batch_size, input_size))
+    x = torch.Tensor(rng.random(size=(batch_size, input_size)))
     x.requires_grad = True
 
     y = network.forward(x)
@@ -169,6 +174,8 @@ def test_mlp_grad(test, device):
 def profile_mlp_torch(device):
     import torch
 
+    rng = np.random.default_rng(123)
+
     m = 128
     n = 64
 
@@ -179,7 +186,7 @@ def profile_mlp_torch(device):
 
         network = create_mlp(m, n)
 
-        x = torch.Tensor(np.random.rand(b, m))
+        x = torch.Tensor(rng.random(size=(b, m)))
 
         with wp.ScopedTimer("torch_forward" + str(b)):
             y = network.forward(x)
@@ -190,7 +197,7 @@ def profile_mlp_torch(device):
 
         network = create_mlp(m, n)
 
-        x = torch.Tensor(np.random.rand(b, m))
+        x = torch.Tensor(rng.random(size=(b, m)))
         y = network.forward(x)
 
         loss = torch.norm(y)
@@ -204,6 +211,8 @@ def profile_mlp_torch(device):
 
 
 def profile_mlp_warp(device):
+    rng = np.random.default_rng(123)
+
     m = 128
     n = 64
 
@@ -212,10 +221,10 @@ def profile_mlp_warp(device):
     for i in range(steps):
         b = 2**i
 
-        weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device)
-        bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device)
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
 
-        x = wp.array(np.random.rand(n, b), dtype=float, device=device)
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device)
         y = wp.zeros(shape=(m, b), device=device)
 
         with wp.ScopedTimer("warp-forward" + str(b)):
@@ -225,10 +234,10 @@ def profile_mlp_warp(device):
     for i in range(steps):
         b = 2**i
 
-        weights = wp.array(np.random.rand(m, n) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
-        bias = wp.array(np.random.rand(m) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
 
-        x = wp.array(np.random.rand(n, b), dtype=float, device=device, requires_grad=True)
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device, requires_grad=True)
         y = wp.zeros(shape=(m, b), device=device, requires_grad=True)
 
         loss = wp.zeros(1, dtype=float, device=device)
@@ -251,18 +260,17 @@ def profile_mlp_warp(device):
 # profile_mlp_torch("cuda")
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestMLP(parent):
-        pass
+class TestMLP(unittest.TestCase):
+    pass
 
-    add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices)
-    add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices)
 
-    return TestMLP
+add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices)
+add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_model.py

@@ -5,103 +5,106 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import warp as wp
-from warp.tests.test_base import *
-from warp.sim import ModelBuilder
+import unittest
 
 import numpy as np
 
+import warp as wp
+from warp.sim import ModelBuilder
+from warp.tests.unittest_utils import *
+
 wp.init()
 
 
-def register(parent):
-    class TestModel(parent):
-        def test_add_triangles(self):
-            pts = np.array(
-                [
-                    [-0.00585869, 0.34189449, -1.17415233],
-                    [-1.894547, 0.1788074, 0.9251329],
-                    [-1.26141048, 0.16140787, 0.08823282],
-                    [-0.08609255, -0.82722546, 0.65995427],
-                    [0.78827592, -1.77375711, -0.55582718],
-                ]
-            )
-            tris = np.array([[0, 3, 4], [0, 2, 3], [2, 1, 3], [1, 4, 3]])
-
-            builder1 = ModelBuilder()
-            builder2 = ModelBuilder()
-            for pt in pts:
-                builder1.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-                builder2.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-
-            # test add_triangle(s) with default arguments:
-            areas = builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2])
-            for i, t in enumerate(tris):
-                area = builder1.add_triangle(t[0], t[1], t[2])
-                self.assertAlmostEqual(area, areas[i], places=6)
-
-            # test add_triangle(s) with non default arguments:
-            tri_ke = np.random.randn(pts.shape[0])
-            tri_ka = np.random.randn(pts.shape[0])
-            tri_kd = np.random.randn(pts.shape[0])
-            tri_drag = np.random.randn(pts.shape[0])
-            tri_lift = np.random.randn(pts.shape[0])
-            for i, t in enumerate(tris):
-                builder1.add_triangle(
-                    t[0],
-                    t[1],
-                    t[2],
-                    tri_ke[i],
-                    tri_ka[i],
-                    tri_kd[i],
-                    tri_drag[i],
-                    tri_lift[i],
-                )
-            builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
-
-            assert_np_equal(np.array(builder1.tri_indices), np.array(builder2.tri_indices))
-            assert_np_equal(np.array(builder1.tri_poses), np.array(builder2.tri_poses), tol=1.0e-6)
-            assert_np_equal(np.array(builder1.tri_activations), np.array(builder2.tri_activations))
-            assert_np_equal(np.array(builder1.tri_materials), np.array(builder2.tri_materials))
-
-        def test_add_edges(self):
-            pts = np.array(
-                [
-                    [-0.00585869, 0.34189449, -1.17415233],
-                    [-1.894547, 0.1788074, 0.9251329],
-                    [-1.26141048, 0.16140787, 0.08823282],
-                    [-0.08609255, -0.82722546, 0.65995427],
-                    [0.78827592, -1.77375711, -0.55582718],
-                ]
+class TestModel(unittest.TestCase):
+    def test_add_triangles(self):
+        rng = np.random.default_rng(123)
+
+        pts = np.array(
+            [
+                [-0.00585869, 0.34189449, -1.17415233],
+                [-1.894547, 0.1788074, 0.9251329],
+                [-1.26141048, 0.16140787, 0.08823282],
+                [-0.08609255, -0.82722546, 0.65995427],
+                [0.78827592, -1.77375711, -0.55582718],
+            ]
+        )
+        tris = np.array([[0, 3, 4], [0, 2, 3], [2, 1, 3], [1, 4, 3]])
+
+        builder1 = ModelBuilder()
+        builder2 = ModelBuilder()
+        for pt in pts:
+            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+
+        # test add_triangle(s) with default arguments:
+        areas = builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2])
+        for i, t in enumerate(tris):
+            area = builder1.add_triangle(t[0], t[1], t[2])
+            self.assertAlmostEqual(area, areas[i], places=6)
+
+        # test add_triangle(s) with non default arguments:
+        tri_ke = rng.standard_normal(size=pts.shape[0])
+        tri_ka = rng.standard_normal(size=pts.shape[0])
+        tri_kd = rng.standard_normal(size=pts.shape[0])
+        tri_drag = rng.standard_normal(size=pts.shape[0])
+        tri_lift = rng.standard_normal(size=pts.shape[0])
+        for i, t in enumerate(tris):
+            builder1.add_triangle(
+                t[0],
+                t[1],
+                t[2],
+                tri_ke[i],
+                tri_ka[i],
+                tri_kd[i],
+                tri_drag[i],
+                tri_lift[i],
             )
-            edges = np.array([[0, 4, 3, 1], [3, 2, 4, 1]])
-
-            builder1 = ModelBuilder()
-            builder2 = ModelBuilder()
-            for pt in pts:
-                builder1.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-                builder2.add_particle(pt, [0.0, 0.0, 0.0], 1.0)
-
-            # test defaults:
-            for i in range(2):
-                builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3])
-            builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3])
-
-            # test non defaults:
-            rest = np.random.randn(2)
-            edge_ke = np.random.randn(2)
-            edge_kd = np.random.randn(2)
-            for i in range(2):
-                builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3], rest[i], edge_ke[i], edge_kd[i])
-            builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3], rest, edge_ke, edge_kd)
-
-            assert_np_equal(np.array(builder1.edge_indices), np.array(builder2.edge_indices))
-            assert_np_equal(np.array(builder1.edge_rest_angle), np.array(builder2.edge_rest_angle), tol=1.0e-4)
-            assert_np_equal(np.array(builder1.edge_bending_properties), np.array(builder2.edge_bending_properties))
-
-    return TestModel
+        builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
+
+        assert_np_equal(np.array(builder1.tri_indices), np.array(builder2.tri_indices))
+        assert_np_equal(np.array(builder1.tri_poses), np.array(builder2.tri_poses), tol=1.0e-6)
+        assert_np_equal(np.array(builder1.tri_activations), np.array(builder2.tri_activations))
+        assert_np_equal(np.array(builder1.tri_materials), np.array(builder2.tri_materials))
+
+    def test_add_edges(self):
+        rng = np.random.default_rng(123)
+
+        pts = np.array(
+            [
+                [-0.00585869, 0.34189449, -1.17415233],
+                [-1.894547, 0.1788074, 0.9251329],
+                [-1.26141048, 0.16140787, 0.08823282],
+                [-0.08609255, -0.82722546, 0.65995427],
+                [0.78827592, -1.77375711, -0.55582718],
+            ]
+        )
+        edges = np.array([[0, 4, 3, 1], [3, 2, 4, 1]])
+
+        builder1 = ModelBuilder()
+        builder2 = ModelBuilder()
+        for pt in pts:
+            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
+
+        # test defaults:
+        for i in range(2):
+            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3])
+        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3])
+
+        # test non defaults:
+        rest = rng.standard_normal(size=2)
+        edge_ke = rng.standard_normal(size=2)
+        edge_kd = rng.standard_normal(size=2)
+        for i in range(2):
+            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3], rest[i], edge_ke[i], edge_kd[i])
+        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3], rest, edge_ke, edge_kd)
+
+        assert_np_equal(np.array(builder1.edge_indices), np.array(builder2.edge_indices))
+        assert_np_equal(np.array(builder1.edge_rest_angle), np.array(builder2.edge_rest_angle), tol=1.0e-4)
+        assert_np_equal(np.array(builder1.edge_bending_properties), np.array(builder2.edge_bending_properties))
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_modules_lite.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+devices = get_test_devices()
+
+
+class TestModuleLite(unittest.TestCase):
+    def test_module_lite_load(self):
+        # Load current module
+        wp.load_module()
+
+        # Load named module
+        wp.load_module(wp.config)
+
+        # Load named module (string)
+        wp.load_module(wp.config, recursive=True)
+
+    def test_module_lite_options(self):
+        wp.set_module_options({"max_unroll": 8})
+        module_options = wp.get_module_options()
+        self.assertIsInstance(module_options, dict)
+        self.assertEqual(module_options["max_unroll"], 8)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

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,156 +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__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)