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

warp/tests/test_tape.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()
 
@@ -19,11 +22,17 @@ def mul_constant(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     y[tid] = x[tid] * 2.0
 
 
+@wp.struct
+class Multiplicands:
+    x: wp.array(dtype=float)
+    y: wp.array(dtype=float)
+
+
 @wp.kernel
-def mul_variable(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+def mul_variable(mutiplicands: Multiplicands, z: wp.array(dtype=float)):
     tid = wp.tid()
 
-    z[tid] = x[tid] * y[tid]
+    z[tid] = mutiplicands.x[tid] * mutiplicands.y[tid]
 
 
 @wp.kernel
@@ -65,12 +74,13 @@ def test_tape_mul_variable(test, device):
 
     # record onto tape
     with tape:
-        # input data
-        x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
-        y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
-        z = wp.zeros_like(x)
+        # input data (Note: We're intentionally testing structs in tapes here)
+        multiplicands = Multiplicands()
+        multiplicands.x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
+        multiplicands.y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
+        z = wp.zeros_like(multiplicands.x)
 
-        wp.launch(kernel=mul_variable, dim=dim, inputs=[x, y], outputs=[z], device=device)
+        wp.launch(kernel=mul_variable, dim=dim, inputs=[multiplicands], outputs=[z], device=device)
 
     # loss = wp.sum(x)
     z.grad = wp.array(np.ones(dim), device=device, dtype=wp.float32)
@@ -79,16 +89,21 @@ def test_tape_mul_variable(test, device):
     tape.backward()
 
     # grad_x=y, grad_y=x
-    assert_np_equal(tape.gradients[x].numpy(), y.numpy())
-    assert_np_equal(tape.gradients[y].numpy(), x.numpy())
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy())
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy())
 
     # run backward again with different incoming gradient
     # should accumulate the same gradients again onto output
     # so gradients = 2.0*prev
     tape.backward()
 
-    assert_np_equal(tape.gradients[x].numpy(), y.numpy() * 2.0)
-    assert_np_equal(tape.gradients[y].numpy(), x.numpy() * 2.0)
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy() * 2.0)
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy() * 2.0)
+
+    # Clear launches and zero out the gradients
+    tape.reset()
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), np.zeros_like(tape.gradients[multiplicands].x.numpy()))
+    test.assertFalse(tape.launches)
 
 
 def test_tape_dot_product(test, device):
@@ -112,19 +127,22 @@ def test_tape_dot_product(test, device):
     assert_np_equal(tape.gradients[y].numpy(), x.numpy())
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestTape(parent):
-        pass
+class TestTape(unittest.TestCase):
+    def test_tape_no_nested_tapes(self):
+        with self.assertRaises(RuntimeError):
+            with wp.Tape():
+                with wp.Tape():
+                    pass
 
-    add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
-    add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
-    add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
 
-    return TestTape
+add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
+add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
+add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_torch.py

@@ -7,11 +7,10 @@
 
 import unittest
 
-# include parent path
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -470,6 +469,8 @@ def test_torch_autograd(test, device):
 def test_torch_graph_torch_stream(test, device):
     """Capture Torch graph on Torch stream"""
 
+    wp.load_module(device=device)
+
     import torch
 
     torch_device = wp.device_to_torch(device)
@@ -551,12 +552,14 @@ def test_warp_graph_warp_stream(test, device):
 
     # capture graph
     with wp.ScopedDevice(device), torch.cuda.stream(torch_stream):
-        wp.capture_begin()
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        g = wp.capture_end()
+        wp.capture_begin(force_module_load=False)
+        try:
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+        finally:
+            g = wp.capture_end()
 
     # replay graph
     num_iters = 10
@@ -570,6 +573,8 @@ def test_warp_graph_warp_stream(test, device):
 def test_warp_graph_torch_stream(test, device):
     """Capture Warp graph on Torch stream"""
 
+    wp.load_module(device=device)
+
     import torch
 
     torch_device = wp.device_to_torch(device)
@@ -587,12 +592,14 @@ def test_warp_graph_torch_stream(test, device):
 
     # capture graph
     with wp.ScopedStream(warp_stream), torch.cuda.stream(torch_stream):
-        wp.capture_begin()
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        g = wp.capture_end()
+        wp.capture_begin(force_module_load=False)
+        try:
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+        finally:
+            g = wp.capture_end()
 
     # replay graph
     num_iters = 10
@@ -603,82 +610,79 @@ def test_warp_graph_torch_stream(test, device):
     assert passed.item()
 
 
-def register(parent):
-    class TestTorch(parent):
-        pass
-
-    try:
-        import torch
-
-        # check which Warp devices work with Torch
-        # CUDA devices may fail if Torch was not compiled with CUDA support
-        test_devices = get_test_devices()
-        torch_compatible_devices = []
-        torch_compatible_cuda_devices = []
-
-        for d in test_devices:
-            try:
-                t = torch.arange(10, device=wp.device_to_torch(d))
-                t += 1
-                torch_compatible_devices.append(d)
-                if d.is_cuda:
-                    torch_compatible_cuda_devices.append(d)
-            except Exception as e:
-                print(f"Skipping Torch tests on device '{d}' due to exception: {e}")
-
-        if torch_compatible_devices:
-            add_function_test(TestTorch, "test_from_torch", test_from_torch, devices=torch_compatible_devices)
-            add_function_test(
-                TestTorch, "test_from_torch_slices", test_from_torch_slices, devices=torch_compatible_devices
-            )
-            add_function_test(
-                TestTorch,
-                "test_from_torch_zero_strides",
-                test_from_torch_zero_strides,
-                devices=torch_compatible_devices,
-            )
-            add_function_test(TestTorch, "test_to_torch", test_to_torch, devices=torch_compatible_devices)
-            add_function_test(TestTorch, "test_torch_zerocopy", test_torch_zerocopy, devices=torch_compatible_devices)
-            add_function_test(TestTorch, "test_torch_autograd", test_torch_autograd, devices=torch_compatible_devices)
-
-        if torch_compatible_cuda_devices:
-            add_function_test(
-                TestTorch,
-                "test_torch_graph_torch_stream",
-                test_torch_graph_torch_stream,
-                devices=torch_compatible_cuda_devices,
-            )
-            add_function_test(
-                TestTorch,
-                "test_torch_graph_warp_stream",
-                test_torch_graph_warp_stream,
-                devices=torch_compatible_cuda_devices,
-            )
-            add_function_test(
-                TestTorch,
-                "test_warp_graph_warp_stream",
-                test_warp_graph_warp_stream,
-                devices=torch_compatible_cuda_devices,
-            )
-            add_function_test(
-                TestTorch,
-                "test_warp_graph_torch_stream",
-                test_warp_graph_torch_stream,
-                devices=torch_compatible_cuda_devices,
-            )
+class TestTorch(unittest.TestCase):
+    pass
 
-        # multi-GPU tests
-        if len(torch_compatible_cuda_devices) > 1:
-            add_function_test(TestTorch, "test_torch_mgpu_from_torch", test_torch_mgpu_from_torch)
-            add_function_test(TestTorch, "test_torch_mgpu_to_torch", test_torch_mgpu_to_torch)
-            add_function_test(TestTorch, "test_torch_mgpu_interop", test_torch_mgpu_interop)
 
-    except Exception as e:
-        print(f"Skipping Torch tests due to exception: {e}")
+test_devices = get_test_devices()
+
+try:
+    import torch
 
-    return TestTorch
+    # check which Warp devices work with Torch
+    # CUDA devices may fail if Torch was not compiled with CUDA support
+    torch_compatible_devices = []
+    torch_compatible_cuda_devices = []
+
+    for d in test_devices:
+        try:
+            t = torch.arange(10, device=wp.device_to_torch(d))
+            t += 1
+            torch_compatible_devices.append(d)
+            if d.is_cuda:
+                torch_compatible_cuda_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Torch tests on device '{d}' due to exception: {e}")
+
+    if torch_compatible_devices:
+        add_function_test(TestTorch, "test_from_torch", test_from_torch, devices=torch_compatible_devices)
+        add_function_test(TestTorch, "test_from_torch_slices", test_from_torch_slices, devices=torch_compatible_devices)
+        add_function_test(
+            TestTorch,
+            "test_from_torch_zero_strides",
+            test_from_torch_zero_strides,
+            devices=torch_compatible_devices,
+        )
+        add_function_test(TestTorch, "test_to_torch", test_to_torch, devices=torch_compatible_devices)
+        add_function_test(TestTorch, "test_torch_zerocopy", test_torch_zerocopy, devices=torch_compatible_devices)
+        add_function_test(TestTorch, "test_torch_autograd", test_torch_autograd, devices=torch_compatible_devices)
+
+    if torch_compatible_cuda_devices:
+        add_function_test(
+            TestTorch,
+            "test_torch_graph_torch_stream",
+            test_torch_graph_torch_stream,
+            devices=torch_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestTorch,
+            "test_torch_graph_warp_stream",
+            test_torch_graph_warp_stream,
+            devices=torch_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestTorch,
+            "test_warp_graph_warp_stream",
+            test_warp_graph_warp_stream,
+            devices=torch_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestTorch,
+            "test_warp_graph_torch_stream",
+            test_warp_graph_torch_stream,
+            devices=torch_compatible_cuda_devices,
+        )
+
+    # multi-GPU tests
+    if len(torch_compatible_cuda_devices) > 1:
+        add_function_test(TestTorch, "test_torch_mgpu_from_torch", test_torch_mgpu_from_torch)
+        add_function_test(TestTorch, "test_torch_mgpu_to_torch", test_torch_mgpu_to_torch)
+        add_function_test(TestTorch, "test_torch_mgpu_interop", test_torch_mgpu_interop)
+
+except Exception as e:
+    print(f"Skipping Torch tests due to exception: {e}")
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_transient_module.py

@@ -5,14 +5,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import importlib
 import os
 import tempfile
 import unittest
+from importlib import util
 
 import warp as wp
-from warp.tests.test_base import *
-from importlib import util
+from warp.tests.unittest_utils import *
 
 CODE = """# -*- coding: utf-8 -*-
 
@@ -64,26 +63,25 @@ def test_transient_module(test, device):
     assert len(module.compute.module.functions) == 1
 
     data = module.Data()
-    data.x = wp.array([123], dtype=int)
+    data.x = wp.array([123], dtype=int, device=device)
 
     wp.set_module_options({"foo": "bar"}, module=module)
     assert wp.get_module_options(module=module).get("foo") == "bar"
     assert module.compute.module.options.get("foo") == "bar"
 
-    wp.launch(module.compute, dim=1, inputs=[data])
+    wp.launch(module.compute, dim=1, inputs=[data], device=device)
     assert_np_equal(data.x.numpy(), np.array([124]))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestTransientModule(parent):
-        pass
+class TestTransientModule(unittest.TestCase):
+    pass
 
-    add_function_test(TestTransientModule, "test_transient_module", test_transient_module, devices=devices)
-    return TestTransientModule
 
+add_function_test(TestTransientModule, "test_transient_module", test_transient_module, devices=devices)
 
 if __name__ == "__main__":
-    _ = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)