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

warp/tests/test_torch.py

@@ -5,13 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
-import numpy as np
 import unittest
-import sys
+
+import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -103,7 +102,7 @@ def test_from_torch(test, device):
     wrap_scalar_tensor_implicit(torch.int16, wp.int16)
     wrap_scalar_tensor_implicit(torch.int8, wp.int8)
     wrap_scalar_tensor_implicit(torch.uint8, wp.uint8)
-    wrap_scalar_tensor_implicit(torch.bool, wp.uint8)
+    wrap_scalar_tensor_implicit(torch.bool, wp.bool)
 
     # explicitly specify warp dtype
     def wrap_scalar_tensor_explicit(torch_dtype, expected_warp_dtype):
@@ -127,6 +126,7 @@ def test_from_torch(test, device):
     wrap_scalar_tensor_explicit(torch.uint8, wp.int8)
     wrap_scalar_tensor_explicit(torch.bool, wp.uint8)
     wrap_scalar_tensor_explicit(torch.bool, wp.int8)
+    wrap_scalar_tensor_explicit(torch.bool, wp.bool)
 
     def wrap_vec_tensor(n, desired_warp_dtype):
         t = torch.zeros((10, n), dtype=torch.float32, device=torch_device)
@@ -151,6 +151,29 @@ def test_from_torch(test, device):
     wrap_mat_tensor(4, 4, wp.mat44)
     wrap_mat_tensor(6, 6, wp.spatial_matrix)
 
+    def wrap_vec_tensor_with_grad(n, desired_warp_dtype):
+        t = torch.zeros((10, n), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, desired_warp_dtype, requires_grad=True)
+        assert a.dtype == desired_warp_dtype
+        assert a.shape == (10,)
+
+    wrap_vec_tensor_with_grad(2, wp.vec2)
+    wrap_vec_tensor_with_grad(3, wp.vec3)
+    wrap_vec_tensor_with_grad(4, wp.vec4)
+    wrap_vec_tensor_with_grad(6, wp.spatial_vector)
+    wrap_vec_tensor_with_grad(7, wp.transform)
+
+    def wrap_mat_tensor_with_grad(n, m, desired_warp_dtype):
+        t = torch.zeros((10, n, m), dtype=torch.float32, device=torch_device)
+        a = wp.from_torch(t, desired_warp_dtype, requires_grad=True)
+        assert a.dtype == desired_warp_dtype
+        assert a.shape == (10,)
+
+    wrap_mat_tensor_with_grad(2, 2, wp.mat22)
+    wrap_mat_tensor_with_grad(3, 3, wp.mat33)
+    wrap_mat_tensor_with_grad(4, 4, wp.mat44)
+    wrap_mat_tensor_with_grad(6, 6, wp.spatial_matrix)
+
 
 def test_to_torch(test, device):
     import torch
@@ -169,6 +192,7 @@ def test_to_torch(test, device):
     wrap_scalar_array(wp.int16, torch.int16)
     wrap_scalar_array(wp.int8, torch.int8)
     wrap_scalar_array(wp.uint8, torch.uint8)
+    wrap_scalar_array(wp.bool, torch.bool)
 
     # not supported by torch
     # wrap_scalar_array(wp.uint64, torch.int64)
@@ -445,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)
@@ -526,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
@@ -545,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)
@@ -562,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
@@ -578,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)
+test_devices = get_test_devices()
 
-    except Exception as e:
-        print(f"Skipping Torch tests due to exception: {e}")
+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,13 +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 warp.tests.unittest_utils import *
 
 CODE = """# -*- coding: utf-8 -*-
 
@@ -45,8 +45,8 @@ def load_code_as_module(code, name):
         with os.fdopen(file, "w") as f:
             f.write(code)
 
-        spec = importlib.util.spec_from_file_location(name, file_path)
-        module = importlib.util.module_from_spec(spec)
+        spec = util.spec_from_file_location(name, file_path)
+        module = util.module_from_spec(spec)
         spec.loader.exec_module(module)
     finally:
         os.remove(file_path)
@@ -63,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
 
-    add_function_test(TestTransientModule, "test_transient_module", test_transient_module, devices=devices)
-    return TestTransientModule
+class TestTransientModule(unittest.TestCase):
+    pass
 
 
+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)