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

warp/tests/test_hash_grid.py

@@ -5,12 +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 *
-
-np.random.seed(532)
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -41,9 +41,7 @@ def count_neighbors(grid: wp.uint64, radius: float, points: wp.array(dtype=wp.ve
     count = int(0)
 
     # construct query around point p
-    neighbors = wp.hash_grid_query(grid, p, radius)
-
-    for index in neighbors:
+    for index in wp.hash_grid_query(grid, p, radius):
         # compute distance to point
         d = wp.length(p - points[index])
 
@@ -74,6 +72,9 @@ def count_neighbors_reference(
 
 
 def test_hashgrid_query(test, device):
+    wp.load_module(device=device)
+    rng = np.random.default_rng(123)
+
     grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
 
     for i in range(num_runs):
@@ -81,15 +82,14 @@ def test_hashgrid_query(test, device):
             print(f"Run: {i+1}")
             print("---------")
 
-        np.random.seed(532)
-        points = np.random.rand(num_points, 3) * scale - np.array((scale, scale, scale)) * 0.5
+        points = rng.random(size=(num_points, 3)) * scale - np.array((scale, scale, scale)) * 0.5
 
         def particle_grid(dim_x, dim_y, dim_z, lower, radius, jitter):
             points = np.meshgrid(
                 np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z)
             )
             points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
-            points_t = points_t + np.random.rand(*points_t.shape) * radius * jitter
+            points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
 
             return points_t.reshape((-1, 3))
 
@@ -99,27 +99,28 @@ def test_hashgrid_query(test, device):
         counts_arr = wp.zeros(len(points), dtype=int, device=device)
         counts_arr_ref = wp.zeros(len(points), dtype=int, device=device)
 
-        with wp.ScopedTimer("brute", active=print_enabled):
-            wp.launch(
-                kernel=count_neighbors_reference,
-                dim=len(points) * len(points),
-                inputs=[query_radius, points_arr, counts_arr_ref, len(points)],
-                device=device,
-            )
-            wp.synchronize()
-
-        with wp.ScopedTimer("grid build", active=print_enabled):
-            grid.build(points_arr, cell_radius)
-            wp.synchronize()
-
-        with wp.ScopedTimer("grid query", active=print_enabled):
-            wp.launch(
-                kernel=count_neighbors,
-                dim=len(points),
-                inputs=[grid.id, query_radius, points_arr, counts_arr],
-                device=device,
-            )
-            wp.synchronize()
+        profiler = {}
+
+        with wp.ScopedTimer("grid operations", print=print_enabled, dict=profiler, synchronize=True):
+            with wp.ScopedTimer("brute", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors_reference,
+                    dim=len(points) * len(points),
+                    inputs=[query_radius, points_arr, counts_arr_ref, len(points)],
+                    device=device,
+                )
+                wp.synchronize()
+
+            with wp.ScopedTimer("grid build", print=print_enabled, dict=profiler, synchronize=True):
+                grid.build(points_arr, cell_radius)
+
+            with wp.ScopedTimer("grid query", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors,
+                    dim=len(points),
+                    inputs=[grid.id, query_radius, points_arr, counts_arr],
+                    device=device,
+                )
 
         counts = counts_arr.numpy()
         counts_ref = counts_arr_ref.numpy()
@@ -133,17 +134,31 @@ def test_hashgrid_query(test, device):
         test.assertTrue(np.array_equal(counts, counts_ref))
 
 
-def register(parent):
-    devices = get_test_devices()
+def test_hashgrid_codegen_adjoints_with_select(test, device):
+    def kernel_fn(
+        grid: wp.uint64,
+    ):
+        v = wp.vec3(0.0, 0.0, 0.0)
+
+        if True:
+            query = wp.hash_grid_query(grid, v, 0.0)
+        else:
+            query = wp.hash_grid_query(grid, v, 0.0)
+
+    wp.Kernel(func=kernel_fn)
+
+
+devices = get_test_devices()
+
 
-    class TestHashGrid(parent):
-        pass
+class TestHashGrid(unittest.TestCase):
+    pass
 
-    add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
 
-    return TestHashGrid
+add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
+add_function_test(TestHashGrid, "test_hashgrid_codegen_adjoints_with_select", test_hashgrid_codegen_adjoints_with_select, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_import.py

@@ -5,20 +5,14 @@
 # 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 math
+import unittest
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+import warp.tests.test_func as test_func
+from warp.tests.unittest_utils import *
 
 wp.init()
 
-# from test_func import sqr
-import warp.tests.test_func as test_func
-
 
 @wp.kernel
 def test_import_func():
@@ -30,23 +24,16 @@ def test_import_func():
     wp.expect_eq(y, 8.0)
 
 
-def register(parent):
-    devices = get_test_devices()
-
-    class TestImport(parent):
-        pass
+devices = get_test_devices()
 
-    add_kernel_test(TestImport, kernel=test_import_func, name="test_import_func", dim=1, devices=devices)
 
-    return TestImport
+class TestImport(unittest.TestCase):
+    pass
 
 
-if __name__ == "__main__":
-    c = register(unittest.TestCase)
-    # unittest.main(verbosity=2)
+add_kernel_test(TestImport, kernel=test_import_func, name="test_import_func", dim=1, devices=devices)
 
-    wp.force_load()
 
-    loader = unittest.defaultTestLoader
-    testSuite = loader.loadTestsFromTestCase(c)
-    testSuite.debug()
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_indexedarray.py

@@ -5,15 +5,14 @@
 # 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
 from typing import Any
 
-import warp as wp
-from warp.tests.test_base import *
-from .test_array import FillStruct
+import numpy as np
 
-import unittest
+import warp as wp
+from warp.tests.test_array import FillStruct
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -37,7 +36,7 @@ def test_indexedarray_1d(test, device):
 
     indices = wp.array([1, 3, 5, 7, 9], dtype=int, device=device)
 
-    iarr = wp.indexedarray(arr, [indices])
+    iarr = wp.indexedarray1d(arr, [indices])
 
     test.assertEqual(iarr.dtype, arr.dtype)
     test.assertEqual(iarr.ndim, 1)
@@ -73,7 +72,7 @@ def test_indexedarray_2d(test, device):
     indices0 = wp.array([1, 3], dtype=int, device=device)
     indices1 = wp.array([2, 4, 8], dtype=int, device=device)
 
-    iarr = wp.indexedarray(arr, [indices0, indices1])
+    iarr = wp.indexedarray2d(arr, [indices0, indices1])
 
     test.assertEqual(iarr.dtype, arr.dtype)
     test.assertEqual(iarr.ndim, 2)
@@ -113,7 +112,7 @@ def test_indexedarray_3d(test, device):
     indices1 = wp.array([2, 4, 8], dtype=int, device=device)
     indices2 = wp.array([0, 5], dtype=int, device=device)
 
-    iarr = wp.indexedarray(arr, [indices0, indices1, indices2])
+    iarr = wp.indexedarray3d(arr, [indices0, indices1, indices2])
 
     test.assertEqual(iarr.dtype, arr.dtype)
     test.assertEqual(iarr.ndim, 3)
@@ -158,7 +157,7 @@ def test_indexedarray_4d(test, device):
     indices2 = wp.array([0, 5], dtype=int, device=device)
     indices3 = wp.array([6, 7, 9], dtype=int, device=device)
 
-    iarr = wp.indexedarray(arr, [indices0, indices1, indices2, indices3])
+    iarr = wp.indexedarray4d(arr, [indices0, indices1, indices2, indices3])
 
     test.assertEqual(iarr.dtype, arr.dtype)
     test.assertEqual(iarr.ndim, 4)
@@ -919,7 +918,10 @@ def test_indexedarray_fill_matrix(test, device):
             assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
 
             # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
-            fill_arr1 = np.arange(mat_len, dtype=nptype)
+            if wptype != wp.bool:
+                fill_arr1 = np.arange(mat_len, dtype=nptype)
+            else:
+                fill_arr1 = np.ones(mat_len, dtype=nptype)
             fill_arr2 = fill_arr1.reshape(mat_shape)
             fill_list1 = list(fill_arr1)
             fill_list2 = [list(row) for row in fill_arr2]
@@ -1104,30 +1106,29 @@ def test_indexedarray_fill_struct(test, device):
     assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestIndexedArray(parent):
-        pass
+class TestIndexedArray(unittest.TestCase):
+    pass
 
-    add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
-    add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
 
-    return TestIndexedArray
+add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_intersect.py

@@ -1,6 +1,16 @@
-import warp as wp
+# 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 numpy as np
-from warp.tests.test_base import *
+
+import warp as wp
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -42,17 +52,16 @@ def test_intersect_tri(test, device):
     assert_np_equal(result.numpy(), np.array([0]))
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestIntersect(parent):
-        pass
+class TestIntersect(unittest.TestCase):
+    pass
 
-    add_function_test(TestIntersect, "test_intersect_tri", test_intersect_tri, devices=devices)
 
-    return TestIntersect
+add_function_test(TestIntersect, "test_intersect_tri", test_intersect_tri, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)

warp/tests/test_large.py

@@ -0,0 +1,141 @@
+# 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 math
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+@wp.kernel
+def conditional_sum(result: wp.array(dtype=wp.uint64)):
+    i, j, k = wp.tid()
+
+    if i == 0:
+        wp.atomic_add(result, 0, wp.uint64(1))
+
+
+def test_large_launch_large_kernel(test, device):
+    """Test tid() on kernel launch of 2**33 threads.
+
+    The function conditional sum will add 1 to result for every thread that has an i index of 0.
+    Due to the size of the grid, this test is not run on CPUs
+    """
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+
+    large_dim_length = 2**16
+    half_result = large_dim_length * large_dim_length
+
+    wp.launch(kernel=conditional_sum, dim=[2, large_dim_length, large_dim_length], inputs=[test_result], device=device)
+    test.assertEqual(test_result.numpy()[0], half_result)
+
+
+@wp.kernel
+def count_elements(result: wp.array(dtype=wp.uint64)):
+    wp.atomic_add(result, 0, wp.uint64(1))
+
+
+def test_large_launch_max_blocks(test, device):
+    # Loop over 1000x1x1 elements using a grid of 256 threads
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+    wp.launch(count_elements, (1000,), inputs=[test_result], max_blocks=1, device=device)
+    test.assertEqual(test_result.numpy()[0], 1000)
+
+    # Loop over 2x10x10 elements using a grid of 256 threads, using the tid() index to count half the elements
+    test_result.zero_()
+    wp.launch(
+        conditional_sum,
+        (
+            2,
+            50,
+            10,
+        ),
+        inputs=[test_result],
+        max_blocks=1,
+        device=device,
+    )
+    test.assertEqual(test_result.numpy()[0], 500)
+
+
+def test_large_launch_very_large_kernel(test, device):
+    """Due to the size of the grid, this test is not run on CPUs"""
+
+    # Dim is chosen to be larger than the maximum CUDA one-dimensional grid size (total threads)
+    dim = (2**31 - 1) * 256 + 1
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+    wp.launch(count_elements, (dim,), inputs=[test_result], device=device)
+    test.assertEqual(test_result.numpy()[0], dim)
+
+
+def test_large_arrays_slow(test, device):
+    # The goal of this test is to use arrays just large enough to know
+    # if there's a flaw in handling arrays with more than 2**31-1 elements
+    # Unfortunately, it takes a long time to run so it won't be run automatically
+    # without changes to support how frequently a test may be run
+    total_elements = 2**31 + 8
+
+    # 1-D to 4-D arrays: test zero_, fill_, then zero_ for scalar data types:
+    for total_dims in range(1, 5):
+        dim_x = math.ceil(total_elements ** (1 / total_dims))
+        shape_tuple = tuple([dim_x] * total_dims)
+
+        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+            a1 = wp.zeros(shape_tuple, dtype=wptype, device=device)
+            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+            a1.fill_(127)
+            assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
+
+            a1.zero_()
+            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+
+def test_large_arrays_fast(test, device):
+    # A truncated version of test_large_arrays_slow meant to catch basic errors
+    total_elements = 2**31 + 8
+
+    nptype = np.dtype(np.int8)
+    wptype = wp.types.np_dtype_to_warp_type[nptype]
+
+    a1 = wp.zeros((total_elements,), dtype=wptype, device=device)
+    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+    a1.fill_(127)
+    assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
+
+    a1.zero_()
+    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+
+devices = get_test_devices()
+
+
+class TestLarge(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestLarge, "test_large_launch_large_kernel", test_large_launch_large_kernel, devices=get_unique_cuda_test_devices()
+)
+
+add_function_test(TestLarge, "test_large_launch_max_blocks", test_large_launch_max_blocks, devices=devices)
+add_function_test(
+    TestLarge,
+    "test_large_launch_very_large_kernel",
+    test_large_launch_very_large_kernel,
+    devices=get_unique_cuda_test_devices(),
+)
+
+add_function_test(TestLarge, "test_large_arrays_fast", test_large_arrays_fast, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_launch.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()
 
@@ -306,51 +305,25 @@ def test_launch_tuple_args(test, device):
     assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
 
 
-@wp.kernel
-def conditional_sum(result: wp.array(dtype=wp.uint64)):
-    i, j, k = wp.tid()
-
-    if i == 0:
-        wp.atomic_add(result, 0, wp.uint64(1))
-
-
-def test_launch_large_kernel(test, device):
-    """Test tid() on kernel launch of 2**33 threads.
-
-    The function conditional sum will add 1 to result for every thread that has an i index of 0.
-    Due to the size of the grid, this test is not run on CPUs
-    """
-    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
-
-    large_dim_length = 2**16
-    half_result = large_dim_length * large_dim_length
-
-    wp.launch(kernel=conditional_sum, dim=[2, large_dim_length, large_dim_length], inputs=[test_result], device=device)
-    test.assertEqual(test_result.numpy()[0], half_result)
-
-
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestLaunch(parent):
-        pass
 
-    add_function_test(TestLaunch, "test_launch_1d", test1d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_2d", test2d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_3d", test3d, devices=devices)
-    add_function_test(TestLaunch, "test_launch_4d", test4d, devices=devices)
+class TestLaunch(unittest.TestCase):
+    pass
 
-    add_function_test(TestLaunch, "test_launch_cmd", test_launch_cmd, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_param, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
-    add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
 
-    add_function_test(TestLaunch, "test_launch_large_kernel", test_launch_large_kernel, devices=wp.get_cuda_devices())
+add_function_test(TestLaunch, "test_launch_1d", test1d, devices=devices)
+add_function_test(TestLaunch, "test_launch_2d", test2d, devices=devices)
+add_function_test(TestLaunch, "test_launch_3d", test3d, devices=devices)
+add_function_test(TestLaunch, "test_launch_4d", test4d, devices=devices)
 
-    return TestLaunch
+add_function_test(TestLaunch, "test_launch_cmd", test_launch_cmd, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_param, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)