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

warp/tests/test_struct.py

@@ -5,15 +5,16 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
 from typing import Any
+
 import numpy as np
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 from warp.fem import Sample as StructFromAnotherModule
 
-import unittest
-
 wp.init()
 
 
@@ -291,7 +292,7 @@ def test_struct_math_conversions(test, device):
     s.m5 = [10, 20, 30, 40]
     s.m6 = np.array([100, 200, 300, 400])
 
-    wp.launch(check_math_conversions, dim=1, inputs=[s])
+    wp.launch(check_math_conversions, dim=1, inputs=[s], device=device)
 
 
 @wp.struct
@@ -416,9 +417,9 @@ def test_nested_array_struct(test, device):
     var2.i = 2
 
     struct = ArrayStruct()
-    struct.array = wp.array([var1, var2], dtype=InnerStruct)
+    struct.array = wp.array([var1, var2], dtype=InnerStruct, device=device)
 
-    wp.launch(struct2_reader, dim=2, inputs=[struct])
+    wp.launch(struct2_reader, dim=2, inputs=[struct], device=device)
 
 
 @wp.struct
@@ -564,81 +565,111 @@ def test_dependent_module_import(c: DependentModuleImport_C):
     wp.tid()  # nop, we're just testing codegen
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestStruct(parent):
-        pass
 
-    add_function_test(TestStruct, "test_step", test_step, devices=devices)
-    add_function_test(TestStruct, "test_step_grad", test_step_grad, devices=devices)
-    add_kernel_test(TestStruct, kernel=test_empty, name="test_empty", dim=1, inputs=[Empty()], devices=devices)
-    add_kernel_test(
-        TestStruct,
-        kernel=test_uninitialized,
-        name="test_uninitialized",
-        dim=1,
-        inputs=[Uninitialized()],
-        devices=devices,
-    )
-    add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
-    add_function_test(TestStruct, "test_struct_attribute_error", test_struct_attribute_error, devices=devices)
-    add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
-    add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
-    add_function_test(TestStruct, "test_nested_empty_struct", test_nested_empty_struct, devices=devices)
-    add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
-    add_function_test(
-        TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
-    )
-    add_kernel_test(
-        TestStruct,
-        name="test_struct_default_attributes",
-        kernel=test_struct_default_attributes_kernel,
-        dim=1,
-        inputs=[],
-        devices=devices,
-    )
+class TestStruct(unittest.TestCase):
+    pass
+
 
+add_function_test(TestStruct, "test_step", test_step, devices=devices)
+add_function_test(TestStruct, "test_step_grad", test_step_grad, devices=devices)
+add_kernel_test(TestStruct, kernel=test_empty, name="test_empty", dim=1, inputs=[Empty()], devices=devices)
+add_kernel_test(
+    TestStruct,
+    kernel=test_uninitialized,
+    name="test_uninitialized",
+    dim=1,
+    inputs=[Uninitialized()],
+    devices=devices,
+)
+add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
+add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_empty_struct", test_nested_empty_struct, devices=devices)
+add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
+add_function_test(
+    TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
+)
+add_kernel_test(
+    TestStruct,
+    name="test_struct_default_attributes",
+    kernel=test_struct_default_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+add_kernel_test(
+    TestStruct,
+    name="test_struct_mutate_attributes",
+    kernel=test_struct_mutate_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+add_kernel_test(
+    TestStruct,
+    kernel=test_uninitialized,
+    name="test_uninitialized",
+    dim=1,
+    inputs=[Uninitialized()],
+    devices=devices,
+)
+add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
+add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_empty_struct", test_nested_empty_struct, devices=devices)
+add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
+add_function_test(
+    TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
+)
+add_kernel_test(
+    TestStruct,
+    name="test_struct_default_attributes",
+    kernel=test_struct_default_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+add_kernel_test(
+    TestStruct,
+    name="test_struct_mutate_attributes",
+    kernel=test_struct_mutate_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+for device in devices:
     add_kernel_test(
         TestStruct,
-        name="test_struct_mutate_attributes",
-        kernel=test_struct_mutate_attributes_kernel,
+        kernel=test_struct_instantiate,
+        name="test_struct_instantiate",
         dim=1,
-        inputs=[],
-        devices=devices,
+        inputs=[wp.array([1], dtype=int, device=device)],
+        devices=[device],
     )
-
-    for device in devices:
-        add_kernel_test(
-            TestStruct,
-            kernel=test_struct_instantiate,
-            name="test_struct_instantiate",
-            dim=1,
-            inputs=[wp.array([1], dtype=int, device=device)],
-            devices=[device],
-        )
-        add_kernel_test(
-            TestStruct,
-            kernel=test_return_struct,
-            name="test_return_struct",
-            dim=1,
-            inputs=[wp.zeros(10, dtype=int, device=device)],
-            devices=[device],
-        )
-
     add_kernel_test(
         TestStruct,
-        kernel=test_dependent_module_import,
-        name="test_dependent_module_import",
+        kernel=test_return_struct,
+        name="test_return_struct",
         dim=1,
-        inputs=[DependentModuleImport_C()],
-        devices=devices,
+        inputs=[wp.zeros(10, dtype=int, device=device)],
+        devices=[device],
     )
 
-    return TestStruct
+add_kernel_test(
+    TestStruct,
+    kernel=test_dependent_module_import,
+    name="test_dependent_module_import",
+    dim=1,
+    inputs=[DependentModuleImport_C()],
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_tape.py

@@ -10,7 +10,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -127,28 +127,22 @@ def test_tape_dot_product(test, device):
     assert_np_equal(tape.gradients[y].numpy(), x.numpy())
 
 
-def test_tape_no_nested_tapes(test, device):
-    with test.assertRaises(RuntimeError):
-        with wp.Tape():
-            with wp.Tape():
-                pass
-
+devices = get_test_devices()
 
-def register(parent):
-    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)
-    add_function_test(TestTape, "test_tape_no_nested_tapes", test_tape_no_nested_tapes, 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__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     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,83 +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__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     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,27 +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__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)