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

warp/tests/test_runlength_encode.py

@@ -1,11 +1,18 @@
-from functools import partial
+# 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
+from functools import partial
 
 import numpy as np
-import warp as wp
 
+import warp as wp
+from warp.tests.unittest_utils import *
 from warp.utils import runlength_encode
-from warp.tests.test_base import *
 
 wp.init()
 
@@ -29,39 +36,10 @@ def test_runlength_encode_int(test, device, n):
     assert_np_equal(unique_counts.numpy()[:run_count], unique_counts_np[:run_count])
 
 
-def test_runlength_encode_error_devices_mismatch(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty_like(values, device="cuda:0")
-    run_lengths = wp.empty_like(values, device="cuda:0")
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Array storage devices do not match$",
-    ):
-        runlength_encode(values, run_values, run_lengths)
-
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty_like(values, device="cpu")
-    run_lengths = wp.empty_like(values, device="cuda:0")
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Array storage devices do not match$",
-    ):
-        runlength_encode(values, run_values, run_lengths)
-
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty_like(values, device="cuda:0")
-    run_lengths = wp.empty_like(values, device="cpu")
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Array storage devices do not match$",
-    ):
-        runlength_encode(values, run_values, run_lengths)
-
-
 def test_runlength_encode_error_insufficient_storage(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty(1, dtype=int, device="cpu")
-    run_lengths = wp.empty(123, dtype=int, device="cpu")
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(1, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=int, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
         r"Output array storage sizes must be at least equal to value_count$",
@@ -79,9 +57,9 @@ def test_runlength_encode_error_insufficient_storage(test, device):
 
 
 def test_runlength_encode_error_dtypes_mismatch(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty(123, dtype=float, device="cpu")
-    run_lengths = wp.empty_like(values)
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=float, device=device)
+    run_lengths = wp.empty_like(values, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
         r"values and run_values data types do not match$",
@@ -90,9 +68,9 @@ def test_runlength_encode_error_dtypes_mismatch(test, device):
 
 
 def test_runlength_encode_error_run_length_unsupported_dtype(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty(123, dtype=int, device="cpu")
-    run_lengths = wp.empty(123, dtype=float, device="cpu")
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=float, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
         r"run_lengths array must be of type int32$",
@@ -100,23 +78,11 @@ def test_runlength_encode_error_run_length_unsupported_dtype(test, device):
         runlength_encode(values, run_values, run_lengths)
 
 
-def test_runlength_encode_error_run_count_device_mismatch(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty_like(values, device="cpu")
-    run_lengths = wp.empty_like(values, device="cpu")
-    run_count = wp.empty(shape=(1,), dtype=int, device="cuda:0")
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"run_count storage device does not match other arrays$",
-    ):
-        runlength_encode(values, run_values, run_lengths, run_count=run_count)
-
-
 def test_runlength_encode_error_run_count_unsupported_dtype(test, device):
-    values = wp.zeros(123, dtype=int, device="cpu")
-    run_values = wp.empty_like(values, device="cpu")
-    run_lengths = wp.empty_like(values, device="cpu")
-    run_count = wp.empty(shape=(1,), dtype=float, device="cpu")
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty_like(values, device=device)
+    run_lengths = wp.empty_like(values, device=device)
+    run_count = wp.empty(shape=(1,), dtype=float, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
         r"run_count array must be of type int32$",
@@ -135,57 +101,90 @@ def test_runlength_encode_error_unsupported_dtype(test, device):
         runlength_encode(values, run_values, run_lengths)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestRunlengthEncode(parent):
-        pass
-
-    add_function_test(
-        TestRunlengthEncode, "test_runlength_encode_int", partial(test_runlength_encode_int, n=100), devices=devices
-    )
-    add_function_test(
-        TestRunlengthEncode, "test_runlength_encode_empty", partial(test_runlength_encode_int, n=0), devices=devices
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_devices_mismatch",
-        test_runlength_encode_error_devices_mismatch,
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_insufficient_storage",
-        test_runlength_encode_error_insufficient_storage,
-    )
-    add_function_test(
-        TestRunlengthEncode, "test_runlength_encode_error_dtypes_mismatch", test_runlength_encode_error_dtypes_mismatch
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_run_length_unsupported_dtype",
-        test_runlength_encode_error_run_length_unsupported_dtype,
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_run_count_device_mismatch",
-        test_runlength_encode_error_run_count_device_mismatch,
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_run_count_unsupported_dtype",
-        test_runlength_encode_error_run_count_unsupported_dtype,
-    )
-    add_function_test(
-        TestRunlengthEncode,
-        "test_runlength_encode_error_unsupported_dtype",
-        test_runlength_encode_error_unsupported_dtype,
-        devices=devices,
-    )
-
-    return TestRunlengthEncode
+devices = get_test_devices()
+
+
+class TestRunlengthEncode(unittest.TestCase):
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_devices_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cpu")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_run_count_device_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cpu")
+        run_count = wp.empty(shape=(1,), dtype=int, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"run_count storage device does not match other arrays$",
+        ):
+            runlength_encode(values, run_values, run_lengths, run_count=run_count)
+
+
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_int", partial(test_runlength_encode_int, n=100), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_empty", partial(test_runlength_encode_int, n=0), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_insufficient_storage",
+    test_runlength_encode_error_insufficient_storage,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_dtypes_mismatch",
+    test_runlength_encode_error_dtypes_mismatch,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_length_unsupported_dtype",
+    test_runlength_encode_error_run_length_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_count_unsupported_dtype",
+    test_runlength_encode_error_run_count_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_unsupported_dtype",
+    test_runlength_encode_error_unsupported_dtype,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_smoothstep.py

@@ -5,14 +5,14 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
 from dataclasses import dataclass
 from typing import Any
-import unittest
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 
 @dataclass
@@ -153,17 +153,16 @@ def test_smoothstep(test, device):
                 )
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
+
+class TestSmoothstep(unittest.TestCase):
+    pass
 
-    class TestSmoothstep(parent):
-        pass
 
-    add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=devices)
-    return TestSmoothstep
+add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_snippet.py

@@ -1,8 +1,10 @@
-import warp as wp
-from warp.tests.test_base import *
-import numpy as np
 import unittest
 
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
 wp.init()
 
 
@@ -101,11 +103,11 @@ def test_shared_memory(test, device):
 
 
 def test_cpu_snippet(test, device):
-    
     snippet = """
     int inc = 1;
     out[tid] = x[tid] + inc;
     """
+
     @wp.func_native(snippet)
     def increment_snippet(
         x: wp.array(dtype=wp.int32),
@@ -115,10 +117,7 @@ def test_cpu_snippet(test, device):
         ...
 
     @wp.kernel
-    def increment(
-        x: wp.array(dtype=wp.int32),
-        out: wp.array(dtype=wp.int32)
-    ):
+    def increment(x: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.int32)):
         tid = wp.tid()
         increment_snippet(x, out, tid)
 
@@ -128,25 +127,17 @@ def test_cpu_snippet(test, device):
 
     wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
 
-    assert_np_equal(out.numpy(), np.arange(1, N+1, 1, dtype=np.int32))
+    assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
 
 
-def register(parent):
-    
-    class TestSnippets(parent):
-        pass
+class TestSnippets(unittest.TestCase):
+    pass
 
-    if wp.is_cuda_available():
-        cuda_device = [wp.get_cuda_device()]
-        add_function_test(TestSnippets, "test_basic", test_basic, devices=cuda_device)
-        add_function_test(TestSnippets, "test_shared_memory", test_shared_memory, devices=cuda_device)
-    
-    if wp.is_cpu_available():
-        add_function_test(TestSnippets, "test_cpu_snippet", test_cpu_snippet, devices=["cpu"])
 
-    return TestSnippets
+add_function_test(TestSnippets, "test_basic", test_basic, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_shared_memory", test_shared_memory, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_cpu_snippet", test_cpu_snippet, devices=["cpu"])
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_sparse.py

@@ -1,13 +1,20 @@
+# 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 numpy as np
-import warp as wp
-import unittest
 
+import warp as wp
 from warp.sparse import bsr_zeros, bsr_set_from_triplets, bsr_get_diag, bsr_diag, bsr_identity, bsr_copy, bsr_scale
 from warp.sparse import bsr_set_transpose, bsr_transposed
 from warp.sparse import bsr_axpy, bsr_mm, bsr_axpy_work_arrays, bsr_mm_work_arrays, bsr_mv
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
+
 
 wp.init()
 
@@ -419,41 +426,35 @@ def make_test_bsr_mv(block_shape, scalar_type):
     return test_bsr_mv
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestSparse(parent):
-        pass
 
-    add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
-    add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
-    add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
-    add_function_test(TestSparse, "test_bsr_copy_scale", test_bsr_copy_scale, devices=devices)
+class TestSparse(unittest.TestCase):
+    pass
 
-    add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
-    add_function_test(
-        TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices
-    )
-    add_function_test(
-        TestSparse, "test_bsr_transpose_3_3", make_test_bsr_transpose((3, 3), wp.float64), devices=devices
-    )
 
-    add_function_test(TestSparse, "test_csr_axpy", make_test_bsr_axpy((1, 1), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_axpy_1_3", make_test_bsr_axpy((1, 3), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_axpy_3_3", make_test_bsr_axpy((3, 3), wp.float64), devices=devices)
+add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
+add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
+add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
+add_function_test(TestSparse, "test_bsr_copy_scale", test_bsr_copy_scale, devices=devices)
+
+add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_transpose_3_3", make_test_bsr_transpose((3, 3), wp.float64), devices=devices)
 
-    add_function_test(TestSparse, "test_csr_mm", make_test_bsr_mm((1, 1), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mm_1_3", make_test_bsr_mm((1, 3), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mm_3_3", make_test_bsr_mm((3, 3), wp.float64), devices=devices)
+add_function_test(TestSparse, "test_csr_axpy", make_test_bsr_axpy((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_axpy_1_3", make_test_bsr_axpy((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_axpy_3_3", make_test_bsr_axpy((3, 3), wp.float64), devices=devices)
 
-    add_function_test(TestSparse, "test_csr_mv", make_test_bsr_mv((1, 1), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mv_1_3", make_test_bsr_mv((1, 3), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mv_3_3", make_test_bsr_mv((3, 3), wp.float64), devices=devices)
+add_function_test(TestSparse, "test_csr_mm", make_test_bsr_mm((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mm_1_3", make_test_bsr_mm((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mm_3_3", make_test_bsr_mm((3, 3), wp.float64), devices=devices)
 
-    return TestSparse
+add_function_test(TestSparse, "test_csr_mv", make_test_bsr_mv((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mv_1_3", make_test_bsr_mv((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mv_3_3", make_test_bsr_mv((3, 3), wp.float64), devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)