warp-lang 0.9.0__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)