warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_array_reduce.py

@@ -18,8 +18,8 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp._src.utils import array_inner, array_sum
 from warp.tests.unittest_utils import *
-from warp.utils import array_inner, array_sum
 
 
 def make_test_array_sum(dtype):
@@ -28,7 +28,7 @@ def make_test_array_sum(dtype):
     def test_array_sum(test, device):
         rng = np.random.default_rng(123)
 
-        cols = wp.types.type_size(dtype)
+        cols = wp._src.types.type_size(dtype)
 
         values_np = rng.random(size=(N, cols))
         values = wp.array(values_np, device=device, dtype=dtype)
@@ -77,7 +77,7 @@ def make_test_array_inner(dtype):
     def test_array_inner(test, device):
         rng = np.random.default_rng(123)
 
-        cols = wp.types.type_size(dtype)
+        cols = wp._src.types.type_size(dtype)
 
         a_np = rng.random(size=(N, cols))
         b_np = rng.random(size=(N, cols))

warp/tests/test_atomic.py

@@ -132,14 +132,18 @@ test_atomic_mat44 = make_atomic_test(wp.mat44)
 
 
 def test_atomic_add_supported_dtypes(test, device, dtype):
+    N = 1024
     scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
 
     @wp.kernel
     def kernel(arr: wp.array(dtype=dtype)):
-        wp.atomic_add(arr, 0, dtype(scalar_type(0)))
+        wp.atomic_add(arr, 0, dtype(scalar_type(1)))
 
     arr = wp.zeros(1, dtype=dtype, device=device)
-    wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+    wp.launch(kernel, dim=N, outputs=(arr,), device=device)
+
+    size = 1 if not hasattr(dtype, "_shape_") else dtype._shape_
+    assert_np_equal(arr.numpy().flatten(), np.full(size, N, dtype=wp.types.warp_type_to_np_dtype[scalar_type]))
 
 
 def test_atomic_min_supported_dtypes(test, device, dtype):
@@ -167,8 +171,8 @@ def test_atomic_max_supported_dtypes(test, device, dtype):
 def test_atomic_add_unsupported_dtypes(test, device, dtype):
     scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
 
-    dtype_str = re.escape(wp.types.type_repr(dtype))
-    scalar_type_str = wp.types.type_repr(scalar_type)
+    dtype_str = re.escape(wp._src.types.type_repr(dtype))
+    scalar_type_str = wp._src.types.type_repr(scalar_type)
 
     @wp.kernel
     def kernel(arr: wp.array(dtype=dtype)):
@@ -188,8 +192,8 @@ def test_atomic_add_unsupported_dtypes(test, device, dtype):
 def test_atomic_min_unsupported_dtypes(test, device, dtype):
     scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
 
-    dtype_str = re.escape(wp.types.type_repr(dtype))
-    scalar_type_str = wp.types.type_repr(scalar_type)
+    dtype_str = re.escape(wp._src.types.type_repr(dtype))
+    scalar_type_str = wp._src.types.type_repr(scalar_type)
 
     @wp.kernel
     def kernel(arr: wp.array(dtype=dtype)):
@@ -209,8 +213,8 @@ def test_atomic_min_unsupported_dtypes(test, device, dtype):
 def test_atomic_max_unsupported_dtypes(test, device, dtype):
     scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
 
-    dtype_str = re.escape(wp.types.type_repr(dtype))
-    scalar_type_str = wp.types.type_repr(scalar_type)
+    dtype_str = re.escape(wp._src.types.type_repr(dtype))
+    scalar_type_str = wp._src.types.type_repr(scalar_type)
 
     @wp.kernel
     def kernel(arr: wp.array(dtype=dtype)):

warp/tests/test_atomic_bitwise.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def test_atomic_bitwise_kernel(a: wp.array(dtype=wp.uint32), b: wp.array(dtype=wp.uint32), op_type: int):
+    i = wp.tid()
+    word_idx = i // 32
+    bit_idx = i % 32
+    bit_mask = wp.uint32(1) << wp.uint32(bit_idx)
+    if op_type == 0:
+        a[word_idx] &= (b[word_idx] & bit_mask) | ~bit_mask
+    elif op_type == 1:
+        a[word_idx] |= b[word_idx] & bit_mask
+    elif op_type == 2:
+        a[word_idx] ^= b[word_idx] & bit_mask
+    elif op_type == 3:
+        wp.atomic_and(a, word_idx, (b[word_idx] & bit_mask) | ~bit_mask)
+    elif op_type == 4:
+        wp.atomic_or(a, word_idx, b[word_idx] & bit_mask)
+    elif op_type == 5:
+        wp.atomic_xor(a, word_idx, b[word_idx] & bit_mask)
+
+
+def test_atomic_bitwise_scalar(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=n, dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=n, dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        and_op_array = wp.array(a, dtype=wp.uint32, device=device)
+        or_op_array = wp.array(a, dtype=wp.uint32, device=device)
+        xor_op_array = wp.array(a, dtype=wp.uint32, device=device)
+        atomic_and_array = wp.array(a, dtype=wp.uint32, device=device)
+        atomic_or_array = wp.array(a, dtype=wp.uint32, device=device)
+        atomic_xor_array = wp.array(a, dtype=wp.uint32, device=device)
+
+        target_array = wp.array(b, dtype=wp.uint32, device=device)
+
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[and_op_array, target_array, 0])
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[or_op_array, target_array, 1])
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[xor_op_array, target_array, 2])
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[atomic_and_array, target_array, 3])
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[atomic_or_array, target_array, 4])
+        wp.launch(test_atomic_bitwise_kernel, dim=n * 32, inputs=[atomic_xor_array, target_array, 5])
+
+        assert_np_equal(and_op_array.numpy(), expected_and)
+        assert_np_equal(or_op_array.numpy(), expected_or)
+        assert_np_equal(xor_op_array.numpy(), expected_xor)
+        assert_np_equal(atomic_and_array.numpy(), expected_and)
+        assert_np_equal(atomic_or_array.numpy(), expected_or)
+        assert_np_equal(atomic_xor_array.numpy(), expected_xor)
+
+
+@wp.kernel
+def test_atomic_bitwise_vector_kernel(a: wp.array(dtype=wp.vec3ui), b: wp.array(dtype=wp.vec3ui), op_type: int):
+    i = wp.tid()
+    word_idx = i // 32
+    bit_idx = i % 32
+    bit_mask = wp.vec3ui(wp.uint32(1)) << wp.vec3ui(wp.uint32(bit_idx))
+    if op_type == 0:
+        a[word_idx] &= (b[word_idx] & bit_mask) | ~bit_mask
+    elif op_type == 1:
+        a[word_idx] |= b[word_idx] & bit_mask
+    elif op_type == 2:
+        a[word_idx] ^= b[word_idx] & bit_mask
+    elif op_type == 3:
+        wp.atomic_and(a, word_idx, (b[word_idx] & bit_mask) | ~bit_mask)
+    elif op_type == 4:
+        wp.atomic_or(a, word_idx, b[word_idx] & bit_mask)
+    elif op_type == 5:
+        wp.atomic_xor(a, word_idx, b[word_idx] & bit_mask)
+
+
+def test_atomic_bitwise_vector(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3), dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3), dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        and_op_array = wp.array(a, dtype=wp.vec3ui, device=device)
+        or_op_array = wp.array(a, dtype=wp.vec3ui, device=device)
+        xor_op_array = wp.array(a, dtype=wp.vec3ui, device=device)
+        atomic_and_array = wp.array(a, dtype=wp.vec3ui, device=device)
+        atomic_or_array = wp.array(a, dtype=wp.vec3ui, device=device)
+        atomic_xor_array = wp.array(a, dtype=wp.vec3ui, device=device)
+
+        target_array = wp.array(b, dtype=wp.vec3ui, device=device)
+
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[and_op_array, target_array, 0])
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[or_op_array, target_array, 1])
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[xor_op_array, target_array, 2])
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[atomic_and_array, target_array, 3])
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[atomic_or_array, target_array, 4])
+        wp.launch(test_atomic_bitwise_vector_kernel, dim=n * 32, inputs=[atomic_xor_array, target_array, 5])
+
+        assert_np_equal(and_op_array.numpy(), expected_and)
+        assert_np_equal(or_op_array.numpy(), expected_or)
+        assert_np_equal(xor_op_array.numpy(), expected_xor)
+        assert_np_equal(atomic_and_array.numpy(), expected_and)
+        assert_np_equal(atomic_or_array.numpy(), expected_or)
+        assert_np_equal(atomic_xor_array.numpy(), expected_xor)
+
+
+mat33ui = wp._src.types.matrix(shape=(3, 3), dtype=wp.uint32)
+
+
+@wp.kernel
+def test_atomic_bitwise_matrix_kernel(a: wp.array(dtype=mat33ui), b: wp.array(dtype=mat33ui), op_type: int):
+    i = wp.tid()
+    word_idx = i // 32
+    bit_idx = i % 32
+    bit_mask = mat33ui(wp.uint32(1)) << mat33ui(wp.uint32(bit_idx))
+    if op_type == 0:
+        a[word_idx] &= (b[word_idx] & bit_mask) | ~bit_mask
+    elif op_type == 1:
+        a[word_idx] |= b[word_idx] & bit_mask
+    elif op_type == 2:
+        a[word_idx] ^= b[word_idx] & bit_mask
+    elif op_type == 3:
+        wp.atomic_and(a, word_idx, (b[word_idx] & bit_mask) | ~bit_mask)
+    elif op_type == 4:
+        wp.atomic_or(a, word_idx, b[word_idx] & bit_mask)
+    elif op_type == 5:
+        wp.atomic_xor(a, word_idx, b[word_idx] & bit_mask)
+
+
+def test_atomic_bitwise_matrix(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3, 3), dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3, 3), dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        and_op_array = wp.array(a, dtype=mat33ui, device=device)
+        or_op_array = wp.array(a, dtype=mat33ui, device=device)
+        xor_op_array = wp.array(a, dtype=mat33ui, device=device)
+        atomic_and_array = wp.array(a, dtype=mat33ui, device=device)
+        atomic_or_array = wp.array(a, dtype=mat33ui, device=device)
+        atomic_xor_array = wp.array(a, dtype=mat33ui, device=device)
+
+        target_array = wp.array(b, dtype=mat33ui, device=device)
+
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[and_op_array, target_array, 0])
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[or_op_array, target_array, 1])
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[xor_op_array, target_array, 2])
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[atomic_and_array, target_array, 3])
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[atomic_or_array, target_array, 4])
+        wp.launch(test_atomic_bitwise_matrix_kernel, dim=n * 32, inputs=[atomic_xor_array, target_array, 5])
+
+        assert_np_equal(and_op_array.numpy(), expected_and)
+        assert_np_equal(or_op_array.numpy(), expected_or)
+        assert_np_equal(xor_op_array.numpy(), expected_xor)
+        assert_np_equal(atomic_and_array.numpy(), expected_and)
+        assert_np_equal(atomic_or_array.numpy(), expected_or)
+        assert_np_equal(atomic_xor_array.numpy(), expected_xor)
+
+
+devices = get_test_devices()
+
+
+class TestAtomicBitwise(unittest.TestCase):
+    pass
+
+
+add_function_test(TestAtomicBitwise, "test_atomic_bitwise_scalar", test_atomic_bitwise_scalar, devices=devices)
+add_function_test(TestAtomicBitwise, "test_atomic_bitwise_vector", test_atomic_bitwise_vector, devices=devices)
+add_function_test(TestAtomicBitwise, "test_atomic_bitwise_matrix", test_atomic_bitwise_matrix, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_atomic_cas.py

@@ -30,7 +30,7 @@ def getkernel(func, suffix=""):
 
 
 def test_atomic_cas(test, device, dtype, register_kernels=False):
-    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    warp_type = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.array([0], dtype=warp_type, device=device)
@@ -84,7 +84,7 @@ def test_atomic_cas(test, device, dtype, register_kernels=False):
 
 
 def test_atomic_cas_2d(test, device, dtype, register_kernels=False):
-    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    warp_type = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
@@ -138,7 +138,7 @@ def test_atomic_cas_2d(test, device, dtype, register_kernels=False):
 
 
 def test_atomic_cas_3d(test, device, dtype, register_kernels=False):
-    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    warp_type = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
@@ -228,7 +228,7 @@ def create_spinlock_test_4d(dtype):
 
 
 def test_atomic_cas_4d(test, device, dtype, register_kernels=False):
-    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    warp_type = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.zeros(shape=(1, 1, 1, 1), dtype=warp_type, device=device)

warp/tests/test_bool.py

@@ -146,7 +146,7 @@ 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)
+vec3bool_type = wp._src.types.vector(length=3, dtype=bool)
 
 
 @wp.kernel
@@ -173,7 +173,7 @@ def test_bool_vec_typing(test, device):
     wp.launch(test_bool_vec_anonymous_typing, (1,), inputs=[], device=device)
 
 
-mat22bool_type = wp.types.matrix((2, 2), dtype=bool)
+mat22bool_type = wp._src.types.matrix((2, 2), dtype=bool)
 
 
 @wp.kernel