warp-lang 1.1.0__py3-none-manylinux2014_x86_64.whl → 1.2.1__py3-none-manylinux2014_x86_64.whl

warp/tests/test_async.py

@@ -13,8 +13,6 @@ import warp as wp
 from warp.tests.unittest_utils import *
 from warp.utils import check_iommu
 
-wp.init()
-
 
 class Capturable:
     def __init__(self, use_graph=True, stream=None):

warp/tests/test_atomic.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 # construct kernel + test function for atomic ops on each vec/matrix type
 def make_atomic_test(type):

warp/tests/test_bool.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 TRUE_CONSTANT = wp.constant(True)
 
 
@@ -90,21 +88,23 @@ def test_bool_constant(test, device):
     test.assertTrue(compile_constant_value.numpy()[0])
 
 
-def test_bool_constant_vec(test, device):
-    vec3bool = wp.vec(length=3, dtype=wp.bool)
-    bool_selector_vec = wp.constant(vec3bool([True, False, True]))
+vec3bool = wp.vec(length=3, dtype=wp.bool)
+bool_selector_vec = wp.constant(vec3bool([True, False, True]))
+
 
-    @wp.kernel
-    def sum_from_bool_vec(sum_array: wp.array(dtype=wp.int32)):
-        i = wp.tid()
+@wp.kernel
+def sum_from_bool_vec(sum_array: wp.array(dtype=wp.int32)):
+    i = wp.tid()
 
-        if bool_selector_vec[0]:
-            sum_array[i] = sum_array[i] + 1
-        if bool_selector_vec[1]:
-            sum_array[i] = sum_array[i] + 2
-        if bool_selector_vec[2]:
-            sum_array[i] = sum_array[i] + 4
+    if bool_selector_vec[0]:
+        sum_array[i] = sum_array[i] + 1
+    if bool_selector_vec[1]:
+        sum_array[i] = sum_array[i] + 2
+    if bool_selector_vec[2]:
+        sum_array[i] = sum_array[i] + 4
 
+
+def test_bool_constant_vec(test, device):
     result_array = wp.zeros(10, dtype=wp.int32, device=device)
 
     wp.launch(sum_from_bool_vec, result_array.shape, inputs=[result_array], device=device)
@@ -112,23 +112,25 @@ def test_bool_constant_vec(test, device):
     assert_np_equal(result_array.numpy(), np.full(result_array.shape, 5))
 
 
-def test_bool_constant_mat(test, device):
-    mat22bool = wp.mat((2, 2), dtype=wp.bool)
-    bool_selector_mat = wp.constant(mat22bool([True, False, False, True]))
-
-    @wp.kernel
-    def sum_from_bool_mat(sum_array: wp.array(dtype=wp.int32)):
-        i = wp.tid()
-
-        if bool_selector_mat[0, 0]:
-            sum_array[i] = sum_array[i] + 1
-        if bool_selector_mat[0, 1]:
-            sum_array[i] = sum_array[i] + 2
-        if bool_selector_mat[1, 0]:
-            sum_array[i] = sum_array[i] + 4
-        if bool_selector_mat[1, 1]:
-            sum_array[i] = sum_array[i] + 8
+mat22bool = wp.mat((2, 2), dtype=wp.bool)
+bool_selector_mat = wp.constant(mat22bool([True, False, False, True]))
+
 
+@wp.kernel
+def sum_from_bool_mat(sum_array: wp.array(dtype=wp.int32)):
+    i = wp.tid()
+
+    if bool_selector_mat[0, 0]:
+        sum_array[i] = sum_array[i] + 1
+    if bool_selector_mat[0, 1]:
+        sum_array[i] = sum_array[i] + 2
+    if bool_selector_mat[1, 0]:
+        sum_array[i] = sum_array[i] + 4
+    if bool_selector_mat[1, 1]:
+        sum_array[i] = sum_array[i] + 8
+
+
+def test_bool_constant_mat(test, device):
     result_array = wp.zeros(10, dtype=wp.int32, device=device)
 
     wp.launch(sum_from_bool_mat, result_array.shape, inputs=[result_array], device=device)
@@ -136,8 +138,20 @@ def test_bool_constant_mat(test, device):
     assert_np_equal(result_array.numpy(), np.full(result_array.shape, 9))
 
 
+vec3bool_type = wp.types.vector(length=3, dtype=bool)
+
+
+@wp.kernel
+def test_bool_vec_anonymous_typing():
+    # Zero initialize
+    wp.expect_eq(vec3bool_type(), wp.vector(False, False, False))
+    # Scalar initialize
+    wp.expect_eq(vec3bool_type(True), wp.vector(True, True, True))
+    # Component-wise initialize
+    wp.expect_eq(vec3bool_type(True, False, True), wp.vector(True, False, True))
+
+
 def test_bool_vec_typing(test, device):
-    vec3bool_type = wp.types.vector(length=3, dtype=bool)
     # Zero initialize
     vec3bool_z = vec3bool_type()
     test.assertEqual(tuple(vec3bool_z), (False, False, False))
@@ -148,20 +162,23 @@ def test_bool_vec_typing(test, device):
     vec3bool_c = vec3bool_type(True, False, True)
     test.assertEqual(tuple(vec3bool_c), (True, False, True))
 
-    @wp.kernel
-    def test_bool_vec_anonymous_typing():
-        # Zero initialize
-        wp.expect_eq(vec3bool_type(), wp.vector(False, False, False))
-        # Scalar initialize
-        wp.expect_eq(vec3bool_type(True), wp.vector(True, True, True))
-        # Component-wise initialize
-        wp.expect_eq(vec3bool_type(True, False, True), wp.vector(True, False, True))
-
     wp.launch(test_bool_vec_anonymous_typing, (1,), inputs=[], device=device)
 
 
+mat22bool_type = wp.types.matrix((2, 2), dtype=bool)
+
+
+@wp.kernel
+def test_bool_mat_anonymous_typing():
+    # Zero initialize
+    wp.expect_eq(mat22bool_type(), wp.matrix(False, False, False, False, shape=(2, 2)))
+    # Scalar initialize
+    wp.expect_eq(mat22bool_type(True), wp.matrix(True, True, True, True, shape=(2, 2)))
+    # Component-wise initialize
+    wp.expect_eq(mat22bool_type(True, False, True, False), wp.matrix(True, False, True, False, shape=(2, 2)))
+
+
 def test_bool_mat_typing(test, device):
-    mat22bool_type = wp.types.matrix((2, 2), dtype=bool)
     # Zero initialize
     mat22bool_z = mat22bool_type()
     test.assertEqual(tuple(mat22bool_z), ((False, False), (False, False)))
@@ -172,15 +189,6 @@ def test_bool_mat_typing(test, device):
     mat22bool_c = mat22bool_type(True, False, True, False)
     test.assertEqual(tuple(mat22bool_c), ((True, False), (True, False)))
 
-    @wp.kernel
-    def test_bool_mat_anonymous_typing():
-        # Zero initialize
-        wp.expect_eq(mat22bool_type(), wp.matrix(False, False, False, False, shape=(2, 2)))
-        # Scalar initialize
-        wp.expect_eq(mat22bool_type(True), wp.matrix(True, True, True, True, shape=(2, 2)))
-        # Component-wise initialize
-        wp.expect_eq(mat22bool_type(True, False, True, False), wp.matrix(True, False, True, False, shape=(2, 2)))
-
     wp.launch(test_bool_mat_anonymous_typing, (1,), inputs=[], device=device)
 
 

warp/tests/test_builtins_resolution.py

@@ -12,8 +12,6 @@ import numpy as np
 
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 def nps(dtype, value):
     """Creates a NumPy scalar value based on the given data type."""

warp/tests/test_bvh.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def bvh_query_aabb(bvh_id: wp.uint64, lower: wp.vec3, upper: wp.vec3, bounds_intersected: wp.array(dtype=int)):

warp/tests/test_closest_point_edge_edge.py

@@ -12,7 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
 epsilon = 0.00001
 
 

warp/tests/test_codegen.py

@@ -11,10 +11,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-# wp.config.mode = "debug"
-
-wp.init()
-
 
 @wp.kernel
 def test_rename():

warp/tests/test_compile_consts.py

@@ -11,8 +11,6 @@ import warp as wp
 import warp.tests.aux_test_compile_consts_dummy
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 LOCAL_ONE = wp.constant(1)
 
 SQRT3_OVER_3 = wp.constant(0.57735026919)
@@ -27,7 +25,7 @@ class Foobar:
 
 
 @wp.kernel
-def test_constants_bool():
+def test_bool():
     if TEST_BOOL:
         expect_eq(1.0, 1.0)
     else:
@@ -35,7 +33,7 @@ def test_constants_bool():
 
 
 @wp.kernel
-def test_constants_int(a: int):
+def test_int(a: int):
     if Foobar.ONE > 0:
         a = 123 + Foobar.TWO + warp.tests.aux_test_compile_consts_dummy.MINUS_ONE
     else:
@@ -44,7 +42,7 @@ def test_constants_int(a: int):
 
 
 @wp.kernel
-def test_constants_float(x: float):
+def test_float(x: float):
     x = SQRT3_OVER_3
     for i in range(3):
         expect_eq(UNIT_VEC[i], x)
@@ -56,7 +54,7 @@ def test_constants_float(x: float):
     expect_near(wp.float32(ONE_FP16), 1.0, 1e-6)
 
 
-def test_constant_closure_capture(test, device):
+def test_closure_capture(test, device):
     def make_closure_kernel(cst):
         def closure_kernel_fn(expected: int):
             wp.expect_eq(cst, expected)
@@ -70,6 +68,124 @@ def test_constant_closure_capture(test, device):
     wp.launch(two_closure, dim=(1), inputs=[2], device=device)
 
 
+def test_hash_global_capture(test, device):
+    """Verifies that global variables are included in the module hash"""
+
+    a = 0
+    wp.launch(test_int, (1,), inputs=[a], device=device)
+
+    module_constants = wp.get_module(test_int.__module__).constants
+
+    # Ensure the expected constants and values are in the dictionary used in hashing
+    # Depending on what's been launched already, there might be additional constants present
+    test.assertEqual(module_constants["Foobar.ONE"], 1)
+    test.assertEqual(module_constants["Foobar.TWO"], 2)
+    test.assertEqual(module_constants["warp.tests.aux_test_compile_consts_dummy.MINUS_ONE"], -1)
+    test.assertEqual(module_constants["LOCAL_ONE"], 1)
+
+
+def test_hash_redefine_kernel(test, device):
+    """This test defines a second ``test_function`` so that the second launch returns the correct result."""
+
+    @wp.kernel
+    def test_function(data: wp.array(dtype=wp.float32)):
+        i = wp.tid()
+        data[i] = TEST_CONSTANT
+
+    TEST_CONSTANT = wp.constant(1.0)
+
+    test_array = wp.empty(1, dtype=wp.float32, device=device)
+    wp.launch(test_function, (1,), inputs=[test_array], device=device)
+    test.assertEqual(test_array.numpy()[0], 1.0)
+
+    module_hash_0 = wp.get_module(test_function.__module__).hash_module()
+    module_constants = wp.get_module(test_function.__module__).constants
+    test.assertEqual(module_constants["TEST_CONSTANT"], 1.0)
+
+    @wp.kernel
+    def test_function(data: wp.array(dtype=wp.float32)):
+        i = wp.tid()
+        data[i] = TEST_CONSTANT
+
+    TEST_CONSTANT = wp.constant(2.0)
+
+    wp.launch(test_function, (1,), inputs=[test_array], device=device)
+    test.assertEqual(test_array.numpy()[0], 2.0)
+
+    module_hash_1 = wp.get_module(test_function.__module__).hash_module()
+    module_constants = wp.get_module(test_function.__module__).constants
+
+    test.assertEqual(module_constants["TEST_CONSTANT"], 2.0)
+    test.assertNotEqual(module_hash_0, module_hash_1)
+
+
+def test_hash_redefine_constant_only(test, device):
+    """This test does not define a second ``test_function``, so the second launch does not invalidate the cache.
+
+    For now this is expected behavior, but we can verify that the content has is different.
+    """
+
+    @wp.kernel
+    def test_function(data: wp.array(dtype=wp.float32)):
+        i = wp.tid()
+        data[i] = TEST_CONSTANT
+
+    TEST_CONSTANT = wp.constant(1.0)
+
+    test_array = wp.empty(1, dtype=wp.float32, device=device)
+    wp.launch(test_function, (1,), inputs=[test_array], device=device)
+    test.assertEqual(test_array.numpy()[0], 1.0)
+
+    module_hash_0 = wp.get_module(test_function.__module__).hash_module()
+
+    module_constants = wp.get_module(test_function.__module__).constants
+    test.assertEqual(module_constants["TEST_CONSTANT"], 1.0)
+
+    TEST_CONSTANT = wp.constant(2.0)
+    module_hash_1 = wp.get_module(test_function.__module__).hash_module(recompute_content_hash=True)
+    module_constants = wp.get_module(test_function.__module__).constants
+    test.assertEqual(module_constants["TEST_CONSTANT"], 2.0)
+    test.assertNotEqual(module_hash_0, module_hash_1, "Module hashes should be different if TEST_CONSTANT is changed.")
+
+    TEST_CONSTANT = wp.constant(1.0)
+    module_hash_2 = wp.get_module(test_function.__module__).hash_module(recompute_content_hash=True)
+    module_constants = wp.get_module(test_function.__module__).constants
+    test.assertEqual(module_constants["TEST_CONSTANT"], 1.0)
+    test.assertEqual(module_hash_0, module_hash_2, "Module hashes should be the same if TEST_CONSTANT is the same.")
+
+
+def test_hash_shadowed_var(test, device):
+    """Tests to ensure shadowed variables are not mistakenly added to the module hash"""
+
+    TEST_CONSTANT_SHADOW_0 = wp.constant(1.0)
+    TEST_CONSTANT_SHADOW_1 = wp.constant(1.0)
+    TEST_CONSTANT_SHADOW_2 = wp.constant(1.0)
+
+    @wp.kernel
+    def test_function(data: wp.array(dtype=wp.float32)):
+        i = wp.tid()
+        TEST_CONSTANT_SHADOW_0 = 2.0
+        TEST_CONSTANT_SHADOW_1, TEST_CONSTANT_SHADOW_2 = 4.0, 8.0
+        data[i] = TEST_CONSTANT_SHADOW_0 + TEST_CONSTANT_SHADOW_1 + TEST_CONSTANT_SHADOW_2
+
+    test_array = wp.empty(1, dtype=wp.float32, device=device)
+    wp.launch(test_function, (1,), inputs=[test_array], device=device)
+    test.assertEqual(test_array.numpy()[0], 14.0)
+
+    module_hash_0 = wp.get_module(test_function.__module__).hash_module()
+    module_constants = wp.get_module(test_function.__module__).constants
+
+    test.assertFalse("TEST_CONSTANT_SHADOW_0" in module_constants, "Constant should not be in dictionary.")
+    test.assertFalse("TEST_CONSTANT_SHADOW_1" in module_constants, "Constant should not be in dictionary.")
+    test.assertFalse("TEST_CONSTANT_SHADOW_2" in module_constants, "Constant should not be in dictionary.")
+
+    TEST_CONSTANT_SHADOW_0 = wp.constant(0.0)
+    TEST_CONSTANT_SHADOW_1 = wp.constant(0.0)
+    TEST_CONSTANT_SHADOW_2 = wp.constant(0.0)
+    module_hash_1 = wp.get_module(test_function.__module__).hash_module(recompute_content_hash=True)
+    test.assertEqual(module_hash_0, module_hash_1, "Module hashes should be the same since all constants are shadowed.")
+
+
 class TestConstants(unittest.TestCase):
     def test_constant_math(self):
         # test doing math with python defined constants in *python* scope
@@ -85,11 +201,15 @@ x = 0.0
 
 devices = get_test_devices()
 
-add_kernel_test(TestConstants, test_constants_bool, dim=1, inputs=[], devices=devices)
-add_kernel_test(TestConstants, test_constants_int, dim=1, inputs=[a], devices=devices)
-add_kernel_test(TestConstants, test_constants_float, dim=1, inputs=[x], devices=devices)
+add_kernel_test(TestConstants, test_bool, dim=1, inputs=[], devices=devices)
+add_kernel_test(TestConstants, test_int, dim=1, inputs=[a], devices=devices)
+add_kernel_test(TestConstants, test_float, dim=1, inputs=[x], devices=devices)
 
-add_function_test(TestConstants, "test_constant_closure_capture", test_constant_closure_capture, devices=devices)
+add_function_test(TestConstants, "test_closure_capture", test_closure_capture, devices=devices)
+add_function_test(TestConstants, "test_hash_global_capture", test_hash_global_capture, devices=devices)
+add_function_test(TestConstants, "test_hash_redefine_kernel", test_hash_redefine_kernel, devices=devices)
+add_function_test(TestConstants, "test_hash_redefine_constant_only", test_hash_redefine_constant_only, devices=devices)
+add_function_test(TestConstants, "test_hash_shadowed_var", test_hash_shadowed_var, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_conditional.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def test_conditional_if_else():

warp/tests/test_copy.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def mul_1d(a: wp.array1d(dtype=float), s: float):

warp/tests/test_ctypes.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def add_vec2(dest: wp.array(dtype=wp.vec2), c: wp.vec2):
@@ -57,7 +55,7 @@ def test_vec2_arg(test, device, n):
     wp.launch(add_vec2, dim=n, inputs=[dest, c], device=device)
 
     # ensure type can round-trip from Python->GPU->Python
-    test.assertTrue(np.array_equal(dest.numpy(), np.tile(c, (n, 1))))
+    assert_np_equal(dest.numpy(), np.tile(c, (n, 1)))
 
 
 def test_vec2_transform(test, device, n):
@@ -67,8 +65,8 @@ def test_vec2_transform(test, device, n):
     m = np.array(((3.0, -1.0), (2.5, 4.0)))
 
     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))))
+    assert_np_equal(dest_right.numpy(), np.tile(m @ c, (n, 1)))
+    assert_np_equal(dest_left.numpy(), np.tile(c @ m, (n, 1)))
 
 
 def test_vec3_arg(test, device, n):
@@ -76,7 +74,7 @@ def test_vec3_arg(test, device, n):
     c = np.array((1.0, 2.0, 3.0))
 
     wp.launch(add_vec3, dim=n, inputs=[dest, c], device=device)
-    test.assertTrue(np.array_equal(dest.numpy(), np.tile(c, (n, 1))))
+    assert_np_equal(dest.numpy(), np.tile(c, (n, 1)))
 
 
 def test_vec3_transform(test, device, n):
@@ -86,8 +84,8 @@ def test_vec3_transform(test, device, n):
     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_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))))
+    assert_np_equal(dest_right.numpy(), np.tile(m @ c, (n, 1)))
+    assert_np_equal(dest_left.numpy(), np.tile(c @ m, (n, 1)))
 
 
 def test_transform_multiply(test, device, n):

warp/tests/test_dense.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def eval_dense_gemm(

warp/tests/test_devices.py

@@ -10,8 +10,6 @@ import unittest
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 def test_devices_get_cuda_device_functions(test, device):
     test.assertTrue(device.is_cuda)

warp/tests/test_dlpack.py

@@ -16,15 +16,13 @@ from warp.tests.unittest_utils import *
 
 N = 1024 * 1024
 
-wp.init()
-
 
 def _jax_version():
     try:
         import jax
 
         return jax.__version_info__
-    except ImportError:
+    except (ImportError, AttributeError):
         return (0, 0, 0)
 
 
@@ -315,8 +313,8 @@ def test_dlpack_warp_to_jax(test, device):
 
     test.assertEqual(a.ptr, j1.unsafe_buffer_pointer())
     test.assertEqual(a.ptr, j2.unsafe_buffer_pointer())
-    test.assertEqual(a.device, wp.device_from_jax(j1.device()))
-    test.assertEqual(a.device, wp.device_from_jax(j2.device()))
+    test.assertEqual(a.device, wp.device_from_jax(list(j1.devices())[0]))
+    test.assertEqual(a.device, wp.device_from_jax(list(j2.devices())[0]))
     test.assertEqual(a.shape, j1.shape)
     test.assertEqual(a.shape, j2.shape)
 
@@ -352,8 +350,8 @@ def test_dlpack_warp_to_jax_v2(test, device):
 
     test.assertEqual(a.ptr, j1.unsafe_buffer_pointer())
     test.assertEqual(a.ptr, j2.unsafe_buffer_pointer())
-    test.assertEqual(a.device, wp.device_from_jax(j1.device()))
-    test.assertEqual(a.device, wp.device_from_jax(j2.device()))
+    test.assertEqual(a.device, wp.device_from_jax(list(j1.devices())[0]))
+    test.assertEqual(a.device, wp.device_from_jax(list(j2.devices())[0]))
     test.assertEqual(a.shape, j1.shape)
     test.assertEqual(a.shape, j2.shape)
 
@@ -387,8 +385,8 @@ def test_dlpack_jax_to_warp(test, device):
 
         test.assertEqual(a1.ptr, j.unsafe_buffer_pointer())
         test.assertEqual(a2.ptr, j.unsafe_buffer_pointer())
-        test.assertEqual(a1.device, wp.device_from_jax(j.device()))
-        test.assertEqual(a2.device, wp.device_from_jax(j.device()))
+        test.assertEqual(a1.device, wp.device_from_jax(list(j.devices())[0]))
+        test.assertEqual(a2.device, wp.device_from_jax(list(j.devices())[0]))
         test.assertEqual(a1.shape, j.shape)
         test.assertEqual(a2.shape, j.shape)
 
@@ -423,8 +421,8 @@ def test_dlpack_jax_to_warp_v2(test, device):
 
         test.assertEqual(a1.ptr, j.unsafe_buffer_pointer())
         test.assertEqual(a2.ptr, j.unsafe_buffer_pointer())
-        test.assertEqual(a1.device, wp.device_from_jax(j.device()))
-        test.assertEqual(a2.device, wp.device_from_jax(j.device()))
+        test.assertEqual(a1.device, wp.device_from_jax(list(j.devices())[0]))
+        test.assertEqual(a2.device, wp.device_from_jax(list(j.devices())[0]))
         test.assertEqual(a1.shape, j.shape)
         test.assertEqual(a2.shape, j.shape)