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

warp/tests/test_copy.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 unittest
+
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -200,19 +199,17 @@ def test_copy_indexed(test, device):
         assert_np_equal(a4.numpy(), expected4 * s)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestCopy(parent):
-        pass
+class TestCopy(unittest.TestCase):
+    pass
 
-    add_function_test(TestCopy, "test_copy_strided", test_copy_strided, devices=devices)
-    add_function_test(TestCopy, "test_copy_indexed", test_copy_indexed, devices=devices)
 
-    return TestCopy
+add_function_test(TestCopy, "test_copy_strided", test_copy_strided, devices=devices)
+add_function_test(TestCopy, "test_copy_indexed", test_copy_indexed, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_ctypes.py

@@ -5,14 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -24,11 +22,11 @@ def add_vec2(dest: wp.array(dtype=wp.vec2), c: wp.vec2):
 
 
 @wp.kernel
-def transform_vec2(dest: wp.array(dtype=wp.vec2), m: wp.mat22, v: wp.vec2):
+def transform_vec2(dest_right: wp.array(dtype=wp.vec2), dest_left: wp.array(dtype=wp.vec2), m: wp.mat22, v: wp.vec2):
     tid = wp.tid()
 
-    p = wp.mul(m, v)
-    dest[tid] = p
+    dest_right[tid] = wp.mul(m, v)
+    dest_left[tid] = wp.mul(v, m)
 
 
 @wp.kernel
@@ -38,11 +36,11 @@ def add_vec3(dest: wp.array(dtype=wp.vec3), c: wp.vec3):
 
 
 @wp.kernel
-def transform_vec3(dest: wp.array(dtype=wp.vec3), m: wp.mat33, v: wp.vec3):
+def transform_vec3(dest_right: wp.array(dtype=wp.vec3), dest_left: wp.array(dtype=wp.vec3), m: wp.mat33, v: wp.vec3):
     tid = wp.tid()
 
-    p = wp.mul(m, v)
-    dest[tid] = p
+    dest_right[tid] = wp.mul(m, v)
+    dest_left[tid] = wp.mul(v, m)
 
 
 @wp.kernel
@@ -63,12 +61,14 @@ def test_vec2_arg(test, device, n):
 
 
 def test_vec2_transform(test, device, n):
-    dest = wp.zeros(n=n, dtype=wp.vec2, device=device)
+    dest_right = wp.zeros(n=n, dtype=wp.vec2, device=device)
+    dest_left = wp.zeros(n=n, dtype=wp.vec2, device=device)
     c = np.array((1.0, 2.0))
     m = np.array(((3.0, -1.0), (2.5, 4.0)))
 
-    wp.launch(transform_vec2, dim=n, inputs=[dest, m, c], device=device)
-    test.assertTrue(np.array_equal(dest.numpy(), np.tile(m @ c, (n, 1))))
+    wp.launch(transform_vec2, dim=n, inputs=[dest_right, dest_left, m, c], device=device)
+    test.assertTrue(np.array_equal(dest_right.numpy(), np.tile(m @ c, (n, 1))))
+    test.assertTrue(np.array_equal(dest_left.numpy(), np.tile(c @ m, (n, 1))))
 
 
 def test_vec3_arg(test, device, n):
@@ -80,12 +80,14 @@ def test_vec3_arg(test, device, n):
 
 
 def test_vec3_transform(test, device, n):
-    dest = wp.zeros(n=n, dtype=wp.vec3, device=device)
+    dest_right = wp.zeros(n=n, dtype=wp.vec3, device=device)
+    dest_left = wp.zeros(n=n, dtype=wp.vec3, device=device)
     c = np.array((1.0, 2.0, 3.0))
     m = np.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)))
 
-    wp.launch(transform_vec3, dim=n, inputs=[dest, m, c], device=device)
-    test.assertTrue(np.array_equal(dest.numpy(), np.tile(m @ c, (n, 1))))
+    wp.launch(transform_vec3, dim=n, inputs=[dest_right, dest_left, m, c], device=device)
+    test.assertTrue(np.array_equal(dest_right.numpy(), np.tile(m @ c, (n, 1))))
+    test.assertTrue(np.array_equal(dest_left.numpy(), np.tile(c @ m, (n, 1))))
 
 
 def test_transform_multiply(test, device, n):
@@ -552,89 +554,79 @@ def test_transform_matrix():
     wp.expect_near(r_2, r_0 - t, 1.0e-4)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestCTypes(parent):
-        pass
-
-    inputs = [
-        wp.vec2(1.0, 2.0),
-        wp.vec3(1.0, 2.0, 3.0),
-        wp.vec4(1.0, 2.0, 3.0, 4.0),
-        wp.mat22(1.0, 2.0, 3.0, 4.0),
-        wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0),
-        wp.mat44(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0),
-    ]
-
-    add_function_test(TestCTypes, "test_mat22", test_mat22, devices=devices)
-    add_function_test(TestCTypes, "test_mat33", test_mat33, devices=devices)
-    add_function_test(TestCTypes, "test_mat44", test_mat44, devices=devices)
-    add_kernel_test(
-        TestCTypes,
-        name="test_transformation_constructor",
-        kernel=test_transformation_constructor,
-        dim=1,
-        devices=devices,
-    )
-    add_kernel_test(
-        TestCTypes,
-        name="test_spatial_vector_constructor",
-        kernel=test_spatial_vector_constructor,
-        dim=1,
-        devices=devices,
-    )
-    add_kernel_test(
-        TestCTypes,
-        name="test_scalar_arg_types",
-        kernel=test_scalar_arg_types,
-        dim=1,
-        inputs=[-64, 255, -64, 255, -64, 255, -64, 255, 3.14159, 3.14159],
-        devices=devices,
-    )
-    add_kernel_test(
-        TestCTypes,
-        name="test_scalar_arg_types_explicit",
-        kernel=test_scalar_arg_types,
-        dim=1,
-        inputs=[
-            wp.int8(-64),
-            wp.uint8(255),
-            wp.int16(-64),
-            wp.uint16(255),
-            wp.int32(-64),
-            wp.uint32(255),
-            wp.int64(-64),
-            wp.uint64(255),
-            wp.float32(3.14159),
-            wp.float64(3.14159),
-        ],
-        devices=devices,
-    )
-    add_kernel_test(
-        TestCTypes, name="test_vector_arg_types", kernel=test_vector_arg_types, dim=1, inputs=inputs, devices=devices
-    )
-    add_kernel_test(TestCTypes, name="test_type_convesrions", kernel=test_type_conversions, dim=1, devices=devices)
-
-    add_function_test(
-        TestCTypes, "test_scalar_array_load", test_scalar_array_types, devices=devices, load=True, store=False
-    )
-    add_function_test(
-        TestCTypes, "test_scalar_array_store", test_scalar_array_types, devices=devices, load=False, store=True
-    )
-    add_function_test(TestCTypes, "test_vec2_arg", test_vec2_arg, devices=devices, n=8)
-    add_function_test(TestCTypes, "test_vec2_transform", test_vec2_transform, devices=devices, n=8)
-    add_function_test(TestCTypes, "test_vec3_arg", test_vec3_arg, devices=devices, n=8)
-    add_function_test(TestCTypes, "test_vec3_transform", test_vec3_transform, devices=devices, n=8)
-    add_function_test(TestCTypes, "test_transform_multiply", test_transform_multiply, devices=devices, n=8)
-    add_kernel_test(TestCTypes, name="test_transform_matrix", kernel=test_transform_matrix, dim=1, devices=devices)
-    add_function_test(TestCTypes, "test_scalar_array", test_scalar_array, devices=devices)
-    add_function_test(TestCTypes, "test_vector_array", test_vector_array, devices=devices)
-
-    return TestCTypes
+devices = get_test_devices()
+
+
+class TestCTypes(unittest.TestCase):
+    pass
+
+
+inputs = [
+    wp.vec2(1.0, 2.0),
+    wp.vec3(1.0, 2.0, 3.0),
+    wp.vec4(1.0, 2.0, 3.0, 4.0),
+    wp.mat22(1.0, 2.0, 3.0, 4.0),
+    wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0),
+    wp.mat44(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0),
+]
+
+add_function_test(TestCTypes, "test_mat22", test_mat22, devices=devices)
+add_function_test(TestCTypes, "test_mat33", test_mat33, devices=devices)
+add_function_test(TestCTypes, "test_mat44", test_mat44, devices=devices)
+add_kernel_test(
+    TestCTypes, name="test_transformation_constructor", kernel=test_transformation_constructor, dim=1, devices=devices
+)
+add_kernel_test(
+    TestCTypes, name="test_spatial_vector_constructor", kernel=test_spatial_vector_constructor, dim=1, devices=devices
+)
+add_kernel_test(
+    TestCTypes,
+    name="test_scalar_arg_types",
+    kernel=test_scalar_arg_types,
+    dim=1,
+    inputs=[-64, 255, -64, 255, -64, 255, -64, 255, 3.14159, 3.14159],
+    devices=devices,
+)
+add_kernel_test(
+    TestCTypes,
+    name="test_scalar_arg_types_explicit",
+    kernel=test_scalar_arg_types,
+    dim=1,
+    inputs=[
+        wp.int8(-64),
+        wp.uint8(255),
+        wp.int16(-64),
+        wp.uint16(255),
+        wp.int32(-64),
+        wp.uint32(255),
+        wp.int64(-64),
+        wp.uint64(255),
+        wp.float32(3.14159),
+        wp.float64(3.14159),
+    ],
+    devices=devices,
+)
+add_kernel_test(
+    TestCTypes, name="test_vector_arg_types", kernel=test_vector_arg_types, dim=1, inputs=inputs, devices=devices
+)
+add_kernel_test(TestCTypes, name="test_type_convesrions", kernel=test_type_conversions, dim=1, devices=devices)
+
+add_function_test(
+    TestCTypes, "test_scalar_array_load", test_scalar_array_types, devices=devices, load=True, store=False
+)
+add_function_test(
+    TestCTypes, "test_scalar_array_store", test_scalar_array_types, devices=devices, load=False, store=True
+)
+add_function_test(TestCTypes, "test_vec2_arg", test_vec2_arg, devices=devices, n=8)
+add_function_test(TestCTypes, "test_vec2_transform", test_vec2_transform, devices=devices, n=8)
+add_function_test(TestCTypes, "test_vec3_arg", test_vec3_arg, devices=devices, n=8)
+add_function_test(TestCTypes, "test_vec3_transform", test_vec3_transform, devices=devices, n=8)
+add_function_test(TestCTypes, "test_transform_multiply", test_transform_multiply, devices=devices, n=8)
+add_kernel_test(TestCTypes, name="test_transform_matrix", kernel=test_transform_matrix, dim=1, devices=devices)
+add_function_test(TestCTypes, "test_scalar_array", test_scalar_array, devices=devices)
+add_function_test(TestCTypes, "test_vector_array", test_vector_array, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_dense.py

@@ -1,11 +1,15 @@
-import numpy as np
-import math
-
-import warp as wp
-from warp.tests.test_base import *
+# Copyright (c) 2022 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()
 
 
@@ -42,20 +46,22 @@ def eval_dense_solve(
     wp.dense_solve(n, A, L, b, x)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestDense(parent):
-        pass
-
+def test_dense_compilation(test, device):
     # just testing compilation of the dense matrix routines
     # most are deprecated / WIP
-    wp.force_load()
+    wp.load_module(device=device)
+
+
+devices = get_test_devices()
+
+
+class TestDense(unittest.TestCase):
+    pass
+
 
-    return TestDense
+add_function_test(TestDense, "test_dense_compilation", test_dense_compilation, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_devices.py

@@ -8,33 +8,31 @@
 import unittest
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
 
-def test_devices_get_device_functions(test, device):
-    # save default device
-    saved_device = wp.get_device()
-
-    test.assertTrue(saved_device.is_cuda)
+def test_devices_get_cuda_device_functions(test, device):
+    test.assertTrue(device.is_cuda)
     test.assertTrue(wp.is_device_available(device))
 
     device_ordinal = device.ordinal
     current_device = wp.get_cuda_device(device_ordinal)
     test.assertEqual(current_device, device)
-
     current_device = wp.get_cuda_device()  # No-ordinal version
     test.assertTrue(wp.is_device_available(current_device))
 
+    if device == current_device:
+        test.assertEqual(device, "cuda")
+    else:
+        test.assertNotEqual(device, "cuda")
+
     preferred_device = wp.get_preferred_device()
     test.assertTrue(wp.is_device_available(preferred_device))
 
-    # restore default device
-    wp.set_device(saved_device)
 
-
-def test_devices_map_device(test, device):
+def test_devices_map_cuda_device(test, device):
     with wp.ScopedDevice(device):
         saved_alias = device.alias
         # Map alias twice to check code path
@@ -58,42 +56,43 @@ def test_devices_verify_cuda_device(test, device):
     wp.config.verify_cuda = verify_cuda_saved
 
 
+@unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
 def test_devices_can_access_self(test, device):
     test.assertTrue(device.can_access(device))
 
-    # Also test CPU access
-    cpu_device = wp.get_device("cpu")
-    if device != cpu_device:
-        test.assertFalse(device.can_access(cpu_device))
+    for warp_device in wp.get_devices():
+        device_str = str(warp_device)
+
+        if (device.is_cpu and warp_device.is_cuda) or (device.is_cuda and warp_device.is_cpu):
+            test.assertFalse(device.can_access(warp_device))
+            test.assertNotEqual(device, warp_device)
+            test.assertNotEqual(device, device_str)
 
-    test.assertNotEqual(cpu_device, "cuda")
 
+devices = get_test_devices()
 
-def register(parent):
-    devices = get_test_devices()
 
-    class TestDevices(parent):
-        pass
+class TestDevices(unittest.TestCase):
+    pass
 
-    add_function_test(
-        TestDevices,
-        "test_devices_get_device_functions",
-        test_devices_get_device_functions,
-        devices=wp.get_cuda_devices(),
-    )
-    add_function_test(TestDevices, "test_devices_map_device", test_devices_map_device, devices=wp.get_cuda_devices())
-    add_function_test(
-        TestDevices, "test_devices_unmap_imaginary_device", test_devices_unmap_imaginary_device, devices=devices
-    )
-    add_function_test(TestDevices, "test_devices_verify_cuda_device", test_devices_verify_cuda_device, devices=devices)
 
-    if wp.is_cuda_available():
-        add_function_test(TestDevices, "test_devices_can_access_self", test_devices_can_access_self, devices=devices)
+add_function_test(
+    TestDevices,
+    "test_devices_get_cuda_device_functions",
+    test_devices_get_cuda_device_functions,
+    devices=get_unique_cuda_test_devices(),
+)
+add_function_test(
+    TestDevices, "test_devices_map_cuda_device", test_devices_map_cuda_device, devices=get_unique_cuda_test_devices()
+)
+add_function_test(
+    TestDevices, "test_devices_unmap_imaginary_device", test_devices_unmap_imaginary_device, devices=devices
+)
+add_function_test(TestDevices, "test_devices_verify_cuda_device", test_devices_verify_cuda_device, devices=devices)
 
-    return TestDevices
+add_function_test(TestDevices, "test_devices_can_access_self", test_devices_can_access_self, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_dlpack.py

@@ -5,13 +5,14 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
-import unittest
-import os
 import ctypes
+import os
+import unittest
+
+import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -299,79 +300,77 @@ def test_dlpack_jax_to_warp(test, device):
         assert_np_equal(a2.numpy(), np.asarray(j))
 
 
-def register(parent):
-    class TestDLPack(parent):
-        pass
-
-    devices = get_test_devices()
-
-    add_function_test(TestDLPack, "test_dlpack_warp_to_warp", test_dlpack_warp_to_warp, devices=devices)
-    add_function_test(TestDLPack, "test_dlpack_dtypes_and_shapes", test_dlpack_dtypes_and_shapes, devices=devices)
-
-    # torch interop via dlpack
-    try:
-        import torch
-        import torch.utils.dlpack
-
-        # 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 = []
-        for d in test_devices:
-            try:
-                t = torch.arange(10, device=wp.device_to_torch(d))
-                t += 1
-                torch_compatible_devices.append(d)
-            except Exception as e:
-                print(f"Skipping Torch DLPack tests on device '{d}' due to exception: {e}")
-
-        if torch_compatible_devices:
-            add_function_test(
-                TestDLPack, "test_dlpack_warp_to_torch", test_dlpack_warp_to_torch, devices=torch_compatible_devices
-            )
-            add_function_test(
-                TestDLPack, "test_dlpack_torch_to_warp", test_dlpack_torch_to_warp, devices=torch_compatible_devices
-            )
-
-    except Exception as e:
-        print(f"Skipping Torch DLPack tests due to exception: {e}")
-
-    # jax interop via dlpack
-    try:
-        # prevent Jax from gobbling up GPU memory
-        os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
-
-        import jax
-        import jax.dlpack
-
-        # check which Warp devices work with Jax
-        # CUDA devices may fail if Jax cannot find a CUDA Toolkit
-        test_devices = get_test_devices()
-        jax_compatible_devices = []
-        for d in test_devices:
-            try:
-                with jax.default_device(wp.device_to_jax(d)):
-                    j = jax.numpy.arange(10, dtype=jax.numpy.float32)
-                    j += 1
-                jax_compatible_devices.append(d)
-            except Exception as e:
-                print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
-
-        if jax_compatible_devices:
-            add_function_test(
-                TestDLPack, "test_dlpack_warp_to_jax", test_dlpack_warp_to_jax, devices=jax_compatible_devices
-            )
-            add_function_test(
-                TestDLPack, "test_dlpack_jax_to_warp", test_dlpack_jax_to_warp, devices=jax_compatible_devices
-            )
-
-    except Exception as e:
-        print(f"Skipping Jax DLPack tests due to exception: {e}")
-
-    return TestDLPack
+class TestDLPack(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestDLPack, "test_dlpack_warp_to_warp", test_dlpack_warp_to_warp, devices=devices)
+add_function_test(TestDLPack, "test_dlpack_dtypes_and_shapes", test_dlpack_dtypes_and_shapes, devices=devices)
+
+# torch interop via dlpack
+try:
+    import torch
+    import torch.utils.dlpack
+
+    # 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 = []
+    for d in test_devices:
+        try:
+            t = torch.arange(10, device=wp.device_to_torch(d))
+            t += 1
+            torch_compatible_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Torch DLPack tests on device '{d}' due to exception: {e}")
+
+    if torch_compatible_devices:
+        add_function_test(
+            TestDLPack, "test_dlpack_warp_to_torch", test_dlpack_warp_to_torch, devices=torch_compatible_devices
+        )
+        add_function_test(
+            TestDLPack, "test_dlpack_torch_to_warp", test_dlpack_torch_to_warp, devices=torch_compatible_devices
+        )
+
+except Exception as e:
+    print(f"Skipping Torch DLPack tests due to exception: {e}")
+
+# jax interop via dlpack
+try:
+    # prevent Jax from gobbling up GPU memory
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+    os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+
+    import jax
+    import jax.dlpack
+
+    # check which Warp devices work with Jax
+    # CUDA devices may fail if Jax cannot find a CUDA Toolkit
+    test_devices = get_test_devices()
+    jax_compatible_devices = []
+    for d in test_devices:
+        try:
+            with jax.default_device(wp.device_to_jax(d)):
+                j = jax.numpy.arange(10, dtype=jax.numpy.float32)
+                j += 1
+            jax_compatible_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
+
+    if jax_compatible_devices:
+        add_function_test(
+            TestDLPack, "test_dlpack_warp_to_jax", test_dlpack_warp_to_jax, devices=jax_compatible_devices
+        )
+        add_function_test(
+            TestDLPack, "test_dlpack_jax_to_warp", test_dlpack_jax_to_warp, devices=jax_compatible_devices
+        )
+
+except Exception as e:
+    print(f"Skipping Jax DLPack tests due to exception: {e}")
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)