warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/tests/test_utils.py

@@ -1,9 +1,17 @@
-# 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.
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import contextlib
 import inspect
@@ -79,7 +87,7 @@ def test_array_scan_error_unsupported_dtype(test, device):
 
 
 def test_radix_sort_pairs(test, device):
-    keyTypes = [int, wp.float32]
+    keyTypes = [int, wp.float32, wp.int64]
 
     for keyType in keyTypes:
         keys = wp.array((7, 2, 8, 4, 1, 6, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0), dtype=keyType, device=device)
@@ -89,18 +97,46 @@ def test_radix_sort_pairs(test, device):
         assert_np_equal(values.numpy()[:8], np.array((5, 2, 8, 4, 7, 6, 1, 3)))
 
 
-def test_radix_sort_pairs_empty(test, device):
+def test_segmented_sort_pairs(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((7, 2, 8, 4, 1, 6, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0), dtype=keyType, device=device)
+        values = wp.array((1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0), dtype=int, device=device)
+        wp.utils.segmented_sort_pairs(
+            keys,
+            values,
+            8,
+            wp.array((0, 4), dtype=int, device=device),
+            wp.array((4, 8), dtype=int, device=device),
+        )
+        assert_np_equal(keys.numpy()[:8], np.array((2, 4, 7, 8, 1, 3, 5, 6)))
+        assert_np_equal(values.numpy()[:8], np.array((2, 4, 1, 3, 5, 8, 7, 6)))
+
+
+def test_radix_sort_pairs_empty(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((), dtype=keyType, device=device)
         values = wp.array((), dtype=int, device=device)
         wp.utils.radix_sort_pairs(keys, values, 0)
 
 
-def test_radix_sort_pairs_error_insufficient_storage(test, device):
+def test_segmented_sort_pairs_empty(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((), dtype=keyType, device=device)
+        values = wp.array((), dtype=int, device=device)
+        wp.utils.segmented_sort_pairs(
+            keys, values, 0, wp.array((), dtype=int, device=device), wp.array((), dtype=int, device=device)
+        )
+
+
+def test_radix_sort_pairs_error_insufficient_storage(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((1, 2, 3), dtype=keyType, device=device)
         values = wp.array((1, 2, 3), dtype=int, device=device)
@@ -111,9 +147,28 @@ def test_radix_sort_pairs_error_insufficient_storage(test, device):
             wp.utils.radix_sort_pairs(keys, values, 3)
 
 
-def test_radix_sort_pairs_error_unsupported_dtype(test, device):
+def test_segmented_sort_pairs_error_insufficient_storage(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((1, 2, 3), dtype=keyType, device=device)
+        values = wp.array((1, 2, 3), dtype=int, device=device)
+        with test.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage must be large enough to contain 2\*count elements$",
+        ):
+            wp.utils.segmented_sort_pairs(
+                keys,
+                values,
+                3,
+                wp.array((0,), dtype=int, device=device),
+                wp.array((3,), dtype=int, device=device),
+            )
+
+
+def test_radix_sort_pairs_error_unsupported_dtype(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((1.0, 2.0, 3.0), dtype=keyType, device=device)
         values = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
@@ -124,6 +179,25 @@ def test_radix_sort_pairs_error_unsupported_dtype(test, device):
             wp.utils.radix_sort_pairs(keys, values, 1)
 
 
+def test_segmented_sort_pairs_error_unsupported_dtype(test, device):
+    keyTypes = [int, wp.float32]
+
+    for keyType in keyTypes:
+        keys = wp.array((1.0, 2.0, 3.0), dtype=keyType, device=device)
+        values = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
+        with test.assertRaisesRegex(
+            RuntimeError,
+            r"Unsupported data type$",
+        ):
+            wp.utils.segmented_sort_pairs(
+                keys,
+                values,
+                1,
+                wp.array((0,), dtype=int, device=device),
+                wp.array((3,), dtype=int, device=device),
+            )
+
+
 def test_array_sum(test, device):
     for dtype in (wp.float32, wp.float64):
         with test.subTest(dtype=dtype):
@@ -460,6 +534,20 @@ add_function_test(
     test_radix_sort_pairs_error_unsupported_dtype,
     devices=devices,
 )
+add_function_test(TestUtils, "test_segmented_sort_pairs", test_segmented_sort_pairs, devices=devices)
+add_function_test(TestUtils, "test_segmented_sort_pairs_empty", test_segmented_sort_pairs, devices=devices)
+add_function_test(
+    TestUtils,
+    "test_segmented_sort_pairs_error_insufficient_storage",
+    test_segmented_sort_pairs_error_insufficient_storage,
+    devices=devices,
+)
+add_function_test(
+    TestUtils,
+    "test_segmented_sort_pairs_error_unsupported_dtype",
+    test_segmented_sort_pairs_error_unsupported_dtype,
+    devices=devices,
+)
 add_function_test(TestUtils, "test_array_sum", test_array_sum, devices=devices)
 add_function_test(
     TestUtils, "test_array_sum_error_out_dtype_mismatch", test_array_sum_error_out_dtype_mismatch, devices=devices

warp/tests/test_vec.py

@@ -1,9 +1,17 @@
-# 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.
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import unittest
 from typing import Any
@@ -1036,7 +1044,7 @@ def test_casting_constructors(test, device, dtype, register_kernels=False):
     assert_np_equal(out, a_grad.numpy())
 
 
-def test_vec_assign(test, device, dtype, register_kernels=False):
+def test_vector_assign_inplace(test, device, dtype, register_kernels=False):
     np_type = np.dtype(dtype)
     wp_type = wp.types.np_dtype_to_warp_type[np_type]
 
@@ -1077,16 +1085,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
         g = a_vec[0] + a_vec[1]
         x[tid] = g
 
-    def vectest_in_register_overwrite(x: wp.array(dtype=vec3), a: wp.array(dtype=vec3)):
-        tid = wp.tid()
-
-        f = vec3(wp_type(0.0))
-        a_vec = a[tid]
-        f = a_vec
-        f[1] = wp_type(3.0)
-
-        x[tid] = f
-
     def vectest_component(x: wp.array(dtype=vec3), y: wp.array(dtype=wp_type)):
         i = wp.tid()
 
@@ -1098,7 +1096,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
 
     kernel_read_write_store = getkernel(vectest_read_write_store, suffix=dtype.__name__)
     kernel_in_register = getkernel(vectest_in_register, suffix=dtype.__name__)
-    kernel_in_register_overwrite = getkernel(vectest_in_register_overwrite, suffix=dtype.__name__)
     kernel_component = getkernel(vectest_component, suffix=dtype.__name__)
 
     if register_kernels:
@@ -1148,7 +1145,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
     x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
     y = wp.ones(1, dtype=wp_type, device=device, requires_grad=True)
 
-    tape = wp.Tape()
     with tape:
         wp.launch(kernel_component, dim=1, inputs=[x, y], device=device)
 
@@ -1157,20 +1153,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
     assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=np_type))
     assert_np_equal(y.grad.numpy(), np.array([6.0], dtype=np_type))
 
-    tape.reset()
-
-    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
-    a = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
-
-    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
-
-    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=np_type))
-    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
-
 
 @wp.kernel
 def test_vector_constructor_value_func():
@@ -1317,15 +1299,15 @@ def vector_augassign_kernel(
 def test_vector_augassign(test, device):
     N = 3
 
-    a = wp.zeros(N, dtype=wp.vec3, requires_grad=True)
-    b = wp.ones(N, dtype=wp.vec3, requires_grad=True)
+    a = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
+    b = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
 
-    c = wp.zeros(N, dtype=wp.vec3, requires_grad=True)
-    d = wp.ones(N, dtype=wp.vec3, requires_grad=True)
+    c = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
+    d = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
 
     tape = wp.Tape()
     with tape:
-        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d])
+        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d], device=device)
 
     tape.backward(grads={a: wp.ones_like(a), c: wp.ones_like(c)})
 
@@ -1338,6 +1320,38 @@ def test_vector_augassign(test, device):
     assert_np_equal(d.grad.numpy(), -wp.ones_like(d).numpy())
 
 
+def test_vector_assign_copy(test, device):
+    saved_enable_vector_component_overwrites_setting = wp.config.enable_vector_component_overwrites
+    try:
+        wp.config.enable_vector_component_overwrites = True
+
+        @wp.kernel
+        def vec_in_register_overwrite(x: wp.array(dtype=wp.vec3), a: wp.array(dtype=wp.vec3)):
+            tid = wp.tid()
+
+            f = wp.vec3(0.0)
+            a_vec = a[tid]
+            f = a_vec
+            f[1] = 3.0
+
+            x[tid] = f
+
+        x = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
+        a = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(vec_in_register_overwrite, dim=1, inputs=[x, a], device=device)
+
+        tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+        assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=float))
+        assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
+
+    finally:
+        wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
+
+
 devices = get_test_devices()
 
 
@@ -1406,8 +1420,8 @@ for dtype in np_float_types:
     )
     add_function_test_register_kernel(
         TestVec,
-        f"test_vec_assign_{dtype.__name__}",
-        test_vec_assign,
+        f"test_vector_assign_inplace_{dtype.__name__}",
+        test_vector_assign_inplace,
         devices=devices,
         dtype=dtype,
     )
@@ -1460,6 +1474,12 @@ add_function_test(
     test_vector_augassign,
     devices=devices,
 )
+add_function_test(
+    TestVec,
+    "test_vector_assign_copy",
+    test_vector_assign_copy,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":

warp/tests/test_vec_lite.py

@@ -1,9 +1,17 @@
-# 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.
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import unittest
 

warp/tests/test_vec_scalar_ops.py

@@ -1,9 +1,17 @@
-# 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.
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import unittest
 

warp/tests/test_verify_fp.py

@@ -1,9 +1,17 @@
-# 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.
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import unittest