warp-lang 1.4.1__py3-none-manylinux2014_x86_64.whl → 1.5.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_func.py

@@ -7,7 +7,7 @@
 
 import math
 import unittest
-from typing import Tuple
+from typing import Any, Tuple
 
 import numpy as np
 
@@ -191,6 +191,47 @@ def test_user_func_return_multiple_values():
     wp.expect_eq(b, 54756.0)
 
 
+@wp.func
+def user_func_overload(
+    b: wp.array(dtype=Any),
+    i: int,
+):
+    return b[i] * 2.0
+
+
+@wp.kernel
+def user_func_overload_resolution_kernel(
+    a: wp.array(dtype=Any),
+    b: wp.array(dtype=Any),
+):
+    i = wp.tid()
+    a[i] = user_func_overload(b, i)
+
+
+def test_user_func_overload_resolution(test, device):
+    a0 = wp.array((1, 2, 3), dtype=wp.vec3)
+    b0 = wp.array((2, 3, 4), dtype=wp.vec3)
+
+    a1 = wp.array((5,), dtype=float)
+    b1 = wp.array((6,), dtype=float)
+
+    wp.launch(user_func_overload_resolution_kernel, a0.shape, (a0, b0))
+    wp.launch(user_func_overload_resolution_kernel, a1.shape, (a1, b1))
+
+    assert_np_equal(a0.numpy()[0], (4, 6, 8))
+    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()
 
 
@@ -375,6 +416,12 @@ add_kernel_test(
     dim=1,
     devices=devices,
 )
+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_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_mesh_query_point.py

@@ -769,20 +769,21 @@ def point_query_aabb_and_closest(
 
 @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_set_mesh_points(test, device):
+    rng = np.random.default_rng(123)
+
     vs, fs = load_mesh()
 
     vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
-    velocities1_np = np.random.randn(vertices1.shape[0], 3)
+    velocities1_np = rng.standard_normal(size=(vertices1.shape[0], 3))
     velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
 
     faces = wp.array(fs, dtype=wp.int32, device=device)
     mesh = wp.Mesh(vertices1, faces, velocities=velocities1)
     fs_2D = faces.reshape((-1, 3))
-    np.random.seed(12345)
     n = 1000
     query_radius = 0.2
 
-    pts1 = wp.array(np.random.randn(n, 3), dtype=wp.vec3, device=device)
+    pts1 = wp.array(rng.standard_normal(size=(n, 3)), dtype=wp.vec3, device=device)
 
     query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
     query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)
@@ -804,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_paddle.py

@@ -7,8 +7,6 @@
 
 import unittest
 
-import numpy as np
-
 import warp as wp
 from warp.tests.unittest_utils import *
 
@@ -444,7 +442,7 @@ def test_from_paddle_slices(test, device):
     assert a.ptr == t.data_ptr()
     assert a.is_contiguous
     assert a.shape == tuple(t.shape)
-    assert_np_equal(a.numpy(), t.cpu().numpy())
+    assert_np_equal(a.numpy(), t.numpy())
 
     # 1D slice with non-contiguous stride
     t_base = paddle.arange(10, dtype=paddle.float32).to(device=paddle_device)
@@ -456,7 +454,7 @@ def test_from_paddle_slices(test, device):
     # copy contents to contiguous array
     a_contiguous = wp.empty_like(a)
     wp.launch(copy1d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # 2D slices (non-contiguous)
     t_base = paddle.arange(24, dtype=paddle.float32).to(device=paddle_device).reshape((4, 6))
@@ -468,7 +466,7 @@ def test_from_paddle_slices(test, device):
     # copy contents to contiguous array
     a_contiguous = wp.empty_like(a)
     wp.launch(copy2d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # 3D slices (non-contiguous)
     t_base = paddle.arange(36, dtype=paddle.float32).to(device=paddle_device).reshape((4, 3, 3))
@@ -480,7 +478,7 @@ def test_from_paddle_slices(test, device):
     # copy contents to contiguous array
     a_contiguous = wp.empty_like(a)
     wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # 2D slices of vec3 (inner contiguous, outer non-contiguous)
     t_base = paddle.arange(150, dtype=paddle.float32).to(device=paddle_device).reshape((10, 5, 3))
@@ -492,7 +490,7 @@ def test_from_paddle_slices(test, device):
     # copy contents to contiguous array
     a_contiguous = wp.empty_like(a)
     wp.launch(copy2d_vec3_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # 2D slices of mat22 (inner contiguous, outer non-contiguous)
     t_base = paddle.arange(200, dtype=paddle.float32).to(device=paddle_device).reshape((10, 5, 2, 2))
@@ -504,7 +502,7 @@ def test_from_paddle_slices(test, device):
     # copy contents to contiguous array
     a_contiguous = wp.empty_like(a)
     wp.launch(copy2d_mat22_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
 
 def test_from_paddle_zero_strides(test, device):
@@ -522,7 +520,7 @@ def test_from_paddle_zero_strides(test, device):
     assert a.shape == tuple(t.shape)
     a_contiguous = wp.empty_like(a)
     wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # expand middle dimension
     t = t_base.unsqueeze(1).expand([-1, 3, -1])
@@ -532,7 +530,7 @@ def test_from_paddle_zero_strides(test, device):
     assert a.shape == tuple(t.shape)
     a_contiguous = wp.empty_like(a)
     wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
     # expand innermost dimension
     t = t_base.unsqueeze(2).expand([-1, -1, 3])
@@ -542,77 +540,7 @@ def test_from_paddle_zero_strides(test, device):
     assert a.shape == tuple(t.shape)
     a_contiguous = wp.empty_like(a)
     wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
-    assert_np_equal(a_contiguous.numpy(), t.cpu().numpy())
-
-
-def test_paddle_mgpu_from_paddle(test, device):
-    import paddle
-
-    n = 32
-
-    t0 = paddle.arange(0, n, 1, dtype=paddle.int32).to(device="gpu:0")
-    t1 = paddle.arange(0, n * 2, 2, dtype=paddle.int32).to(device="gpu:1")
-
-    a0 = wp.from_paddle(t0, dtype=wp.int32)
-    a1 = wp.from_paddle(t1, dtype=wp.int32)
-
-    assert a0.device == "gpu:0"
-    assert a1.device == "gpu:1"
-
-    expected0 = np.arange(0, n, 1)
-    expected1 = np.arange(0, n * 2, 2)
-
-    assert_np_equal(a0.numpy(), expected0)
-    assert_np_equal(a1.numpy(), expected1)
-
-
-def test_paddle_mgpu_to_paddle(test, device):
-    n = 32
-
-    with wp.ScopedDevice("gpu:0"):
-        a0 = wp.empty(n, dtype=wp.int32)
-        wp.launch(arange, dim=a0.size, inputs=[0, 1, a0])
-
-    with wp.ScopedDevice("gpu:1"):
-        a1 = wp.empty(n, dtype=wp.int32)
-        wp.launch(arange, dim=a1.size, inputs=[0, 2, a1])
-
-    t0 = wp.to_paddle(a0)
-    t1 = wp.to_paddle(a1)
-
-    assert str(t0.device) == "gpu:0"
-    assert str(t1.device) == "gpu:1"
-
-    expected0 = np.arange(0, n, 1, dtype=np.int32)
-    expected1 = np.arange(0, n * 2, 2, dtype=np.int32)
-
-    assert_np_equal(t0.cpu().numpy(), expected0)
-    assert_np_equal(t1.cpu().numpy(), expected1)
-
-
-def test_paddle_mgpu_interop(test, device):
-    import paddle
-
-    n = 1024 * 1024
-
-    with paddle.cuda.device(0):
-        t0 = paddle.arange(n, dtype=paddle.float32).to(device="gpu")
-        a0 = wp.from_paddle(t0)
-        wp.launch(inc, dim=a0.size, inputs=[a0], stream=wp.stream_from_paddle())
-
-    with paddle.cuda.device(1):
-        t1 = paddle.arange(n, dtype=paddle.float32).to(device="gpu")
-        a1 = wp.from_paddle(t1)
-        wp.launch(inc, dim=a1.size, inputs=[a1], stream=wp.stream_from_paddle())
-
-    assert a0.device == "gpu:0"
-    assert a1.device == "gpu:1"
-
-    expected = np.arange(n, dtype=int) + 1
-
-    # ensure the paddle tensors were modified by warp
-    assert_np_equal(t0.cpu().numpy(), expected)
-    assert_np_equal(t1.cpu().numpy(), expected)
+    assert_np_equal(a_contiguous.numpy(), t.numpy())
 
 
 def test_paddle_autograd(test, device):
@@ -624,6 +552,9 @@ def test_paddle_autograd(test, device):
     class TestFunc(paddle.autograd.PyLayer):
         @staticmethod
         def forward(ctx, x):
+            # ensure Paddle operations complete before running Warp
+            wp.synchronize_device()
+
             # allocate output array
             y = paddle.empty_like(x)
 
@@ -632,10 +563,16 @@ def test_paddle_autograd(test, device):
 
             wp.launch(kernel=op_kernel, dim=len(x), inputs=[wp.from_paddle(x)], outputs=[wp.from_paddle(y)])
 
+            # ensure Warp operations complete before returning data to Paddle
+            wp.synchronize_device()
+
             return y
 
         @staticmethod
         def backward(ctx, adj_y):
+            # ensure Paddle operations complete before running Warp
+            wp.synchronize_device()
+
             # adjoints should be allocated as zero initialized
             adj_x = paddle.zeros_like(ctx.x).contiguous()
             adj_y = adj_y.contiguous()
@@ -655,6 +592,9 @@ def test_paddle_autograd(test, device):
                 adjoint=True,
             )
 
+            # ensure Warp operations complete before returning data to Paddle
+            wp.synchronize_device()
+
             return adj_x
 
     # run autograd on given device
@@ -691,7 +631,7 @@ def test_warp_graph_warp_stream(test, device):
     paddle_stream = wp.stream_to_paddle(device)
 
     # capture graph
-    with wp.ScopedDevice(device), paddle.device.stream(paddle_stream):
+    with wp.ScopedDevice(device), paddle.device.stream_guard(paddle.device.Stream(paddle_stream)):
         wp.capture_begin(force_module_load=False)
         try:
             t += 1.0
@@ -837,11 +777,11 @@ try:
     #         devices=paddle_compatible_cuda_devices,
     #     )
 
-    # multi-GPU tests
-    if len(paddle_compatible_cuda_devices) > 1:
-        add_function_test(TestPaddle, "test_paddle_mgpu_from_paddle", test_paddle_mgpu_from_paddle)
-        add_function_test(TestPaddle, "test_paddle_mgpu_to_paddle", test_paddle_mgpu_to_paddle)
-        add_function_test(TestPaddle, "test_paddle_mgpu_interop", test_paddle_mgpu_interop)
+    # multi-GPU not supported yet.
+    # if len(paddle_compatible_cuda_devices) > 1:
+    #     add_function_test(TestPaddle, "test_paddle_mgpu_from_paddle", test_paddle_mgpu_from_paddle)
+    #     add_function_test(TestPaddle, "test_paddle_mgpu_to_paddle", test_paddle_mgpu_to_paddle)
+    #     add_function_test(TestPaddle, "test_paddle_mgpu_interop", test_paddle_mgpu_interop)
 
 except Exception as e:
     print(f"Skipping Paddle tests due to exception: {e}")