warp-lang 1.4.2__py3-none-manylinux2014_x86_64.whl → 1.5.1__py3-none-manylinux2014_x86_64.whl

warp/tests/test_func.py

@@ -162,7 +162,7 @@ def user_func_with_defaults(a: int = 123, b: int = 234) -> int:
 
 
 @wp.kernel
-def test_user_func_with_defaults():
+def user_func_with_defaults_kernel():
     a = user_func_with_defaults()
     wp.expect_eq(a, 357)
 
@@ -179,6 +179,25 @@ def test_user_func_with_defaults():
     wp.expect_eq(e, 234)
 
 
+def test_user_func_with_defaults(test, device):
+    wp.launch(user_func_with_defaults_kernel, dim=1, device=device)
+
+    a = user_func_with_defaults()
+    assert a == 357
+
+    b = user_func_with_defaults(111)
+    assert b == 345
+
+    c = user_func_with_defaults(111, 222)
+    assert c == 333
+
+    d = user_func_with_defaults(a=111)
+    assert d == 345
+
+    e = user_func_with_defaults(b=111)
+    assert e == 234
+
+
 @wp.func
 def user_func_return_multiple_values(a: int, b: float) -> Tuple[int, float]:
     return a + a, b * b
@@ -222,6 +241,16 @@ def test_user_func_overload_resolution(test, device):
     assert a1.numpy()[0] == 12
 
 
+@wp.func
+def user_func_return_none() -> None:
+    pass
+
+
+@wp.kernel
+def test_return_annotation_none() -> None:
+    user_func_return_none()
+
+
 devices = get_test_devices()
 
 
@@ -396,9 +425,7 @@ add_function_test(TestFunc, func=test_func_closure_capture, name="test_func_clos
 add_function_test(TestFunc, func=test_multi_valued_func, name="test_multi_valued_func", devices=devices)
 add_kernel_test(TestFunc, kernel=test_func_defaults, name="test_func_defaults", dim=1, devices=devices)
 add_kernel_test(TestFunc, kernel=test_builtin_shadowing, name="test_builtin_shadowing", dim=1, devices=devices)
-add_kernel_test(
-    TestFunc, kernel=test_user_func_with_defaults, name="test_user_func_with_defaults", dim=1, devices=devices
-)
+add_function_test(TestFunc, func=test_user_func_with_defaults, name="test_user_func_with_defaults", devices=devices)
 add_kernel_test(
     TestFunc,
     kernel=test_user_func_return_multiple_values,
@@ -409,6 +436,9 @@ add_kernel_test(
 add_function_test(
     TestFunc, func=test_user_func_overload_resolution, name="test_user_func_overload_resolution", devices=devices
 )
+add_kernel_test(
+    TestFunc, kernel=test_return_annotation_none, name="test_return_annotation_none", dim=1, devices=devices
+)
 
 
 if __name__ == "__main__":

warp/tests/test_generics.py

@@ -522,6 +522,57 @@ def test_type_attribute_error(test, device):
         )
 
 
+@wp.func
+def vec_int_annotation_func(v: wp.vec(3, wp.Int)) -> wp.Int:
+    return v[0] + v[1] + v[2]
+
+
+@wp.func
+def vec_float_annotation_func(v: wp.vec(3, wp.Float)) -> wp.Float:
+    return v[0] + v[1] + v[2]
+
+
+@wp.func
+def vec_scalar_annotation_func(v: wp.vec(3, wp.Scalar)) -> wp.Scalar:
+    return v[0] + v[1] + v[2]
+
+
+@wp.func
+def mat_int_annotation_func(m: wp.mat((2, 2), wp.Int)) -> wp.Int:
+    return m[0, 0] + m[0, 1] + m[1, 0] + m[1, 1]
+
+
+@wp.func
+def mat_float_annotation_func(m: wp.mat((2, 2), wp.Float)) -> wp.Float:
+    return m[0, 0] + m[0, 1] + m[1, 0] + m[1, 1]
+
+
+@wp.func
+def mat_scalar_annotation_func(m: wp.mat((2, 2), wp.Scalar)) -> wp.Scalar:
+    return m[0, 0] + m[0, 1] + m[1, 0] + m[1, 1]
+
+
+mat22s = wp.mat((2, 2), wp.int16)
+mat22d = wp.mat((2, 2), wp.float64)
+
+
+@wp.kernel
+def test_annotations_kernel():
+    vi16 = wp.vec3s(wp.int16(1), wp.int16(2), wp.int16(3))
+    vf64 = wp.vec3d(wp.float64(1), wp.float64(2), wp.float64(3))
+    wp.expect_eq(vec_int_annotation_func(vi16), wp.int16(6))
+    wp.expect_eq(vec_float_annotation_func(vf64), wp.float64(6))
+    wp.expect_eq(vec_scalar_annotation_func(vi16), wp.int16(6))
+    wp.expect_eq(vec_scalar_annotation_func(vf64), wp.float64(6))
+
+    mi16 = mat22s(wp.int16(1), wp.int16(2), wp.int16(3), wp.int16(4))
+    mf64 = mat22d(wp.float64(1), wp.float64(2), wp.float64(3), wp.float64(4))
+    wp.expect_eq(mat_int_annotation_func(mi16), wp.int16(10))
+    wp.expect_eq(mat_float_annotation_func(mf64), wp.float64(10))
+    wp.expect_eq(mat_scalar_annotation_func(mi16), wp.int16(10))
+    wp.expect_eq(mat_scalar_annotation_func(mf64), wp.float64(10))
+
+
 class TestGenerics(unittest.TestCase):
     pass
 
@@ -590,6 +641,7 @@ add_kernel_test(
 )
 add_function_test(TestGenerics, "test_type_operator_misspell", test_type_operator_misspell, devices=devices)
 add_function_test(TestGenerics, "test_type_attribute_error", test_type_attribute_error, devices=devices)
+add_kernel_test(TestGenerics, name="test_annotations_kernel", kernel=test_annotations_kernel, dim=1, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_iter.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2024 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.kernel
+def reversed_kernel(
+    start: wp.int32,
+    end: wp.int32,
+    step: wp.int32,
+    out_count: wp.array(dtype=wp.int32),
+    out_values: wp.array(dtype=wp.int32),
+):
+    count = wp.int32(0)
+    for i in reversed(range(start, end, step)):
+        out_values[count] = i
+        count += 1
+
+    out_count[0] = count
+
+
+def test_reversed(test, device):
+    count = wp.empty(1, dtype=wp.int32)
+    values = wp.empty(32, dtype=wp.int32)
+
+    start, end, step = (-2, 8, 3)
+    wp.launch(
+        reversed_kernel,
+        dim=1,
+        inputs=(start, end, step),
+        outputs=(count, values),
+    )
+    expected = tuple(reversed(range(start, end, step)))
+    assert count.numpy()[0] == len(expected)
+    assert_np_equal(values.numpy()[: len(expected)], expected)
+
+    start, end, step = (9, -3, -2)
+    wp.launch(
+        reversed_kernel,
+        dim=1,
+        inputs=(start, end, step),
+        outputs=(count, values),
+    )
+    expected = tuple(reversed(range(start, end, step)))
+    assert count.numpy()[0] == len(expected)
+    assert_np_equal(values.numpy()[: len(expected)], expected)
+
+
+devices = get_test_devices()
+
+
+class TestIter(unittest.TestCase):
+    pass
+
+
+add_function_test(TestIter, "test_reversed", test_reversed, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_lerp.py

@@ -31,66 +31,14 @@ class TestData:
 
 TEST_DATA = {
     wp.float32: (
-        TestData(
-            a=1.0,
-            b=5.0,
-            t=0.75,
-            expected=4.0,
-            expected_adj_a=0.25,
-            expected_adj_b=0.75,
-            expected_adj_t=4.0,
-        ),
-        TestData(
-            a=-2.0,
-            b=5.0,
-            t=0.25,
-            expected=-0.25,
-            expected_adj_a=0.75,
-            expected_adj_b=0.25,
-            expected_adj_t=7.0,
-        ),
-        TestData(
-            a=1.23,
-            b=2.34,
-            t=0.5,
-            expected=1.785,
-            expected_adj_a=0.5,
-            expected_adj_b=0.5,
-            expected_adj_t=1.11,
-        ),
-    ),
-    wp.vec2: (
-        TestData(
-            a=[1, 2],
-            b=[3, 4],
-            t=0.5,
-            expected=[2, 3],
-        ),
-    ),
-    wp.vec3: (
-        TestData(
-            a=[1, 2, 3],
-            b=[3, 4, 5],
-            t=0.5,
-            expected=[2, 3, 4],
-        ),
-    ),
-    wp.vec4: (
-        TestData(
-            a=[1, 2, 3, 4],
-            b=[3, 4, 5, 6],
-            t=0.5,
-            expected=[2, 3, 4, 5],
-        ),
-    ),
-    wp.mat22: (
-        TestData(
-            a=[[1, 2], [2, 1]],
-            b=[[3, 4], [4, 3]],
-            t=0.5,
-            expected=[[2, 3], [3, 2]],
-        ),
+        TestData(a=1.0, b=5.0, t=0.75, expected=4.0, expected_adj_a=0.25, expected_adj_b=0.75, expected_adj_t=4.0),
+        TestData(a=-2.0, b=5.0, t=0.25, expected=-0.25, expected_adj_a=0.75, expected_adj_b=0.25, expected_adj_t=7.0),
+        TestData(a=1.23, b=2.34, t=0.5, expected=1.785, expected_adj_a=0.5, expected_adj_b=0.5, expected_adj_t=1.11),
     ),
+    wp.vec2: (TestData(a=[1, 2], b=[3, 4], t=0.5, expected=[2, 3]),),
+    wp.vec3: (TestData(a=[1, 2, 3], b=[3, 4, 5], t=0.5, expected=[2, 3, 4]),),
+    wp.vec4: (TestData(a=[1, 2, 3, 4], b=[3, 4, 5, 6], t=0.5, expected=[2, 3, 4, 5]),),
+    wp.mat22: (TestData(a=[[1, 2], [2, 1]], b=[[3, 4], [4, 3]], t=0.5, expected=[[2, 3], [3, 2]]),),
     wp.mat33: (
         TestData(
             a=[[1, 2, 3], [3, 1, 2], [2, 3, 1]],
@@ -107,30 +55,9 @@ TEST_DATA = {
             expected=[[2, 3, 4, 5], [5, 2, 3, 4], [4, 5, 2, 3], [3, 4, 5, 2]],
         ),
     ),
-    wp.quat: (
-        TestData(
-            a=[1, 2, 3, 4],
-            b=[3, 4, 5, 6],
-            t=0.5,
-            expected=[2, 3, 4, 5],
-        ),
-    ),
-    wp.transform: (
-        TestData(
-            a=[1, 2, 3, 4, 5, 6, 7],
-            b=[3, 4, 5, 6, 7, 8, 9],
-            t=0.5,
-            expected=[2, 3, 4, 5, 6, 7, 8],
-        ),
-    ),
-    wp.spatial_vector: (
-        TestData(
-            a=[1, 2, 3, 4, 5, 6],
-            b=[3, 4, 5, 6, 7, 8],
-            t=0.5,
-            expected=[2, 3, 4, 5, 6, 7],
-        ),
-    ),
+    wp.quat: (TestData(a=[1, 2, 3, 4], b=[3, 4, 5, 6], t=0.5, expected=[2, 3, 4, 5]),),
+    wp.transform: (TestData(a=[1, 2, 3, 4, 5, 6, 7], b=[3, 4, 5, 6, 7, 8, 9], t=0.5, expected=[2, 3, 4, 5, 6, 7, 8]),),
+    wp.spatial_vector: (TestData(a=[1, 2, 3, 4, 5, 6], b=[3, 4, 5, 6, 7, 8], t=0.5, expected=[2, 3, 4, 5, 6, 7]),),
     wp.spatial_matrix: (
         TestData(
             a=[
@@ -175,12 +102,12 @@ def test_lerp(test, device):
 
         return fn
 
-    for data_type in TEST_DATA:
+    for data_type, test_data_set in TEST_DATA.items():
         kernel_fn = make_kernel_fn(data_type)
         kernel = wp.Kernel(func=kernel_fn, key=f"test_lerp_{data_type.__name__}_kernel")
 
         with test.subTest(data_type=data_type):
-            for test_data in TEST_DATA[data_type]:
+            for test_data in test_data_set:
                 a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
                 b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
                 t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
@@ -188,8 +115,7 @@ def test_lerp(test, device):
                     [0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True
                 )
 
-                tape = wp.Tape()
-                with tape:
+                with wp.Tape() as tape:
                     wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
 
                 assert_np_equal(out.numpy(), np.array([test_data.expected]), tol=1e-6)

warp/tests/test_mat_scalar_ops.py

@@ -1501,7 +1501,7 @@ def test_matmat_multiplication(test, device, dtype, register_kernels=False):
     tol = {
         np.float16: 2.0e-2,
         np.float32: 5.0e-6,
-        np.float64: 1.0e-8,
+        np.float64: 5.0e-7,
     }.get(dtype, 0)
 
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]

warp/tests/test_matmul.py

@@ -5,6 +5,7 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import itertools
 import unittest
 from typing import Any
 
@@ -105,19 +106,15 @@ class gemm_test_bed_runner:
         assert_np_equal(C.grad.numpy(), adj_C_np)
 
     def run(self):
-        Ms = [64, 128, 256]
-        Ns = [64, 128, 256]
-        Ks = [64, 128, 256]
+        Ms = [16, 32, 64]
+        Ns = [16, 32, 64]
+        Ks = [16, 32, 64]
         batch_counts = [1, 4]
         betas = [0.0, 1.0]
         alpha = 1.0
 
-        for batch_count in batch_counts:
-            for m in Ms:
-                for n in Ns:
-                    for k in Ks:
-                        for beta in betas:
-                            self.run_and_verify(m, n, k, batch_count, alpha, beta)
+        for batch_count, m, n, k, beta in itertools.product(batch_counts, Ms, Ns, Ks, betas):
+            self.run_and_verify(m, n, k, batch_count, alpha, beta)
 
 
 class gemm_test_bed_runner_transpose:

warp/tests/test_matmul_lite.py

@@ -102,19 +102,14 @@ class gemm_test_bed_runner:
         assert_np_equal(C.grad.numpy(), adj_C_np)
 
     def run(self):
-        Ms = [8]
-        Ns = [16]
-        Ks = [32]
-        batch_counts = [1]
-        betas = [1.0]
+        m = 8
+        n = 16
+        k = 32
+        batch_count = 1
+        beta = 1.0
         alpha = 1.0
 
-        for batch_count in batch_counts:
-            for m in Ms:
-                for n in Ns:
-                    for k in Ks:
-                        for beta in betas:
-                            self.run_and_verify(m, n, k, batch_count, alpha, beta)
+        self.run_and_verify(m, n, k, batch_count, alpha, beta)
 
 
 class gemm_test_bed_runner_transpose:

warp/tests/test_mesh_query_point.py

@@ -805,7 +805,7 @@ def test_set_mesh_points(test, device):
         device=device,
     )
 
-    shift = np.random.randn(3)
+    shift = rng.standard_normal(size=3)
 
     vs_higher = vs + shift
     vertices2 = wp.array(vs_higher, dtype=wp.vec3, device=device)

warp/tests/test_module_hashing.py

@@ -214,12 +214,35 @@ def test_function_generic_overload_hashing(test, device):
     test.assertNotEqual(hash4, hash1)
 
 
+SIMPLE_MODULE = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.kernel
+def k():
+    pass
+"""
+
+
+def test_module_load(test, device):
+    """Ensure that loading a module does not change its hash"""
+    m = load_code_as_module(SIMPLE_MODULE, "simple_module")
+
+    hash1 = m.hash_module()
+    m.load(device)
+    hash2 = m.hash_module()
+
+    test.assertEqual(hash1, hash2)
+
+
 class TestModuleHashing(unittest.TestCase):
     pass
 
 
+devices = get_test_devices()
+
 add_function_test(TestModuleHashing, "test_function_overload_hashing", test_function_overload_hashing)
 add_function_test(TestModuleHashing, "test_function_generic_overload_hashing", test_function_generic_overload_hashing)
+add_function_test(TestModuleHashing, "test_module_load", test_module_load, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_overwrite.py

@@ -507,6 +507,50 @@ def test_kernel_read_func_write(test, device):
         wp.config.verify_autograd_array_access = saved_verify_autograd_array_access_setting
 
 
+@wp.func
+def atomic_func(
+    a: wp.array(dtype=wp.int32),
+    b: wp.array(dtype=wp.int32),
+    c: wp.array(dtype=wp.int32),
+    d: wp.array(dtype=wp.int32),
+    i: int,
+):
+    wp.atomic_add(a, i, 1)
+    wp.atomic_sub(b, i, 1)
+    wp.atomic_min(c, i, 1)
+    wp.atomic_max(d, i, 3)
+
+
+@wp.kernel(enable_backward=False)
+def atomic_kernel(
+    a: wp.array(dtype=wp.int32), b: wp.array(dtype=wp.int32), c: wp.array(dtype=wp.int32), d: wp.array(dtype=wp.int32)
+):
+    i = wp.tid()
+    atomic_func(a, b, c, d, i)
+
+
+# atomic operations should mark arrays as WRITE
+def test_atomic_operations(test, device):
+    saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
+    try:
+        wp.config.verify_autograd_array_access = True
+
+        a = wp.array((1, 2, 3), dtype=wp.int32, device=device)
+        b = wp.array((1, 2, 3), dtype=wp.int32, device=device)
+        c = wp.array((1, 2, 3), dtype=wp.int32, device=device)
+        d = wp.array((1, 2, 3), dtype=wp.int32, device=device)
+
+        wp.launch(atomic_kernel, dim=a.shape, inputs=(a, b, c, d), device=device)
+
+        test.assertEqual(atomic_kernel.adj.args[0].is_write, True)
+        test.assertEqual(atomic_kernel.adj.args[1].is_write, True)
+        test.assertEqual(atomic_kernel.adj.args[2].is_write, True)
+        test.assertEqual(atomic_kernel.adj.args[3].is_write, True)
+
+    finally:
+        wp.config.verify_autograd_array_access = saved_verify_autograd_array_access_setting
+
+
 class TestOverwrite(unittest.TestCase):
     pass
 
@@ -535,6 +579,7 @@ add_function_test(TestOverwrite, "test_reset", test_reset, devices=devices)
 add_function_test(TestOverwrite, "test_copy", test_copy, devices=devices)
 add_function_test(TestOverwrite, "test_matmul", test_matmul, devices=devices)
 add_function_test(TestOverwrite, "test_batched_matmul", test_batched_matmul, devices=devices)
+add_function_test(TestOverwrite, "test_atomic_operations", test_atomic_operations, devices=devices)
 
 # Some warning are only issued during codegen, and codegen only runs on cuda_0 in the MGPU case.
 cuda_device = get_cuda_test_devices(mode="basic")