warp-lang 1.0.2__py3-none-macosx_10_13_universal2.whl → 1.1.0__py3-none-macosx_10_13_universal2.whl

warp/tests/test_atomic.py

@@ -1,141 +1,141 @@
-# 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.
-
-import unittest
-
-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):
-    def test_atomic_kernel(
-        out_add: wp.array(dtype=type),
-        out_min: wp.array(dtype=type),
-        out_max: wp.array(dtype=type),
-        val: wp.array(dtype=type),
-    ):
-        tid = wp.tid()
-
-        wp.atomic_add(out_add, 0, val[tid])
-        wp.atomic_min(out_min, 0, val[tid])
-        wp.atomic_max(out_max, 0, val[tid])
-
-    # register a custom kernel (no decorator) function
-    # this lets us register the same function definition
-    # against multiple symbols, with different arg types
-    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel")
-
-    def test_atomic(test, device):
-        n = 1024
-
-        rng = np.random.default_rng(42)
-
-        if type == wp.int32:
-            base = (rng.random(size=1, dtype=np.float32) * 100.0).astype(np.int32)
-            val = (rng.random(size=n, dtype=np.float32) * 100.0).astype(np.int32)
-
-        elif type == wp.float32:
-            base = rng.random(size=1, dtype=np.float32)
-            val = rng.random(size=n, dtype=np.float32)
-
-        else:
-            base = rng.random(size=(1, *type._shape_), dtype=float)
-            val = rng.random(size=(n, *type._shape_), dtype=float)
-
-        add_array = wp.array(base, dtype=type, device=device, requires_grad=True)
-        min_array = wp.array(base, dtype=type, device=device, requires_grad=True)
-        max_array = wp.array(base, dtype=type, device=device, requires_grad=True)
-        add_array.zero_()
-        min_array.fill_(10000)
-        max_array.fill_(-10000)
-
-        val_array = wp.array(val, dtype=type, device=device, requires_grad=True)
-
-        tape = wp.Tape()
-        with tape:
-            wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
-
-        assert_np_equal(add_array.numpy(), np.sum(val, axis=0), tol=1.0e-2)
-        assert_np_equal(min_array.numpy(), np.min(val, axis=0), tol=1.0e-2)
-        assert_np_equal(max_array.numpy(), np.max(val, axis=0), tol=1.0e-2)
-
-        if type != wp.int32:
-            add_array.grad.fill_(1)
-            tape.backward()
-            assert_np_equal(val_array.grad.numpy(), np.ones_like(val))
-            tape.zero()
-
-            min_array.grad.fill_(1)
-            tape.backward()
-            min_grad_array = np.zeros_like(val)
-            argmin = val.argmin(axis=0)
-            if val.ndim == 1:
-                min_grad_array[argmin] = 1
-            elif val.ndim == 2:
-                for i in range(val.shape[1]):
-                    min_grad_array[argmin[i], i] = 1
-            elif val.ndim == 3:
-                for i in range(val.shape[1]):
-                    for j in range(val.shape[2]):
-                        min_grad_array[argmin[i, j], i, j] = 1
-            assert_np_equal(val_array.grad.numpy(), min_grad_array)
-            tape.zero()
-
-            max_array.grad.fill_(1)
-            tape.backward()
-            max_grad_array = np.zeros_like(val)
-            argmax = val.argmax(axis=0)
-            if val.ndim == 1:
-                max_grad_array[argmax] = 1
-            elif val.ndim == 2:
-                for i in range(val.shape[1]):
-                    max_grad_array[argmax[i], i] = 1
-            elif val.ndim == 3:
-                for i in range(val.shape[1]):
-                    for j in range(val.shape[2]):
-                        max_grad_array[argmax[i, j], i, j] = 1
-            assert_np_equal(val_array.grad.numpy(), max_grad_array)
-
-    return test_atomic
-
-
-# generate test functions for atomic types
-test_atomic_int = make_atomic_test(wp.int32)
-test_atomic_float = make_atomic_test(wp.float32)
-test_atomic_vec2 = make_atomic_test(wp.vec2)
-test_atomic_vec3 = make_atomic_test(wp.vec3)
-test_atomic_vec4 = make_atomic_test(wp.vec4)
-test_atomic_mat22 = make_atomic_test(wp.mat22)
-test_atomic_mat33 = make_atomic_test(wp.mat33)
-test_atomic_mat44 = make_atomic_test(wp.mat44)
-
-
-devices = get_test_devices()
-
-
-class TestAtomic(unittest.TestCase):
-    pass
-
-
-add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
-add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
-add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
-add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
-add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
-add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
-add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
-add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
-
-
-if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
-    unittest.main(verbosity=2)
+# 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.
+
+import unittest
+
+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):
+    def test_atomic_kernel(
+        out_add: wp.array(dtype=type),
+        out_min: wp.array(dtype=type),
+        out_max: wp.array(dtype=type),
+        val: wp.array(dtype=type),
+    ):
+        tid = wp.tid()
+
+        wp.atomic_add(out_add, 0, val[tid])
+        wp.atomic_min(out_min, 0, val[tid])
+        wp.atomic_max(out_max, 0, val[tid])
+
+    # register a custom kernel (no decorator) function
+    # this lets us register the same function definition
+    # against multiple symbols, with different arg types
+    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel")
+
+    def test_atomic(test, device):
+        n = 1024
+
+        rng = np.random.default_rng(42)
+
+        if type == wp.int32:
+            base = (rng.random(size=1, dtype=np.float32) * 100.0).astype(np.int32)
+            val = (rng.random(size=n, dtype=np.float32) * 100.0).astype(np.int32)
+
+        elif type == wp.float32:
+            base = rng.random(size=1, dtype=np.float32)
+            val = rng.random(size=n, dtype=np.float32)
+
+        else:
+            base = rng.random(size=(1, *type._shape_), dtype=float)
+            val = rng.random(size=(n, *type._shape_), dtype=float)
+
+        add_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        min_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        max_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        add_array.zero_()
+        min_array.fill_(10000)
+        max_array.fill_(-10000)
+
+        val_array = wp.array(val, dtype=type, device=device, requires_grad=True)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
+
+        assert_np_equal(add_array.numpy(), np.sum(val, axis=0), tol=1.0e-2)
+        assert_np_equal(min_array.numpy(), np.min(val, axis=0), tol=1.0e-2)
+        assert_np_equal(max_array.numpy(), np.max(val, axis=0), tol=1.0e-2)
+
+        if type != wp.int32:
+            add_array.grad.fill_(1)
+            tape.backward()
+            assert_np_equal(val_array.grad.numpy(), np.ones_like(val))
+            tape.zero()
+
+            min_array.grad.fill_(1)
+            tape.backward()
+            min_grad_array = np.zeros_like(val)
+            argmin = val.argmin(axis=0)
+            if val.ndim == 1:
+                min_grad_array[argmin] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    min_grad_array[argmin[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        min_grad_array[argmin[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), min_grad_array)
+            tape.zero()
+
+            max_array.grad.fill_(1)
+            tape.backward()
+            max_grad_array = np.zeros_like(val)
+            argmax = val.argmax(axis=0)
+            if val.ndim == 1:
+                max_grad_array[argmax] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    max_grad_array[argmax[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        max_grad_array[argmax[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), max_grad_array)
+
+    return test_atomic
+
+
+# generate test functions for atomic types
+test_atomic_int = make_atomic_test(wp.int32)
+test_atomic_float = make_atomic_test(wp.float32)
+test_atomic_vec2 = make_atomic_test(wp.vec2)
+test_atomic_vec3 = make_atomic_test(wp.vec3)
+test_atomic_vec4 = make_atomic_test(wp.vec4)
+test_atomic_mat22 = make_atomic_test(wp.mat22)
+test_atomic_mat33 = make_atomic_test(wp.mat33)
+test_atomic_mat44 = make_atomic_test(wp.mat44)
+
+
+devices = get_test_devices()
+
+
+class TestAtomic(unittest.TestCase):
+    pass
+
+
+add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
+add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_bool.py

@@ -1,149 +1,204 @@
-# 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.
-
-import unittest
-
-import numpy as np
-
-import warp as wp
-from warp.tests.unittest_utils import *
-
-wp.init()
-
-TRUE_CONSTANT = wp.constant(True)
-
-
-@wp.func
-def identity_function(input_bool: wp.bool, plain_bool: bool):
-    return input_bool and plain_bool
-
-
-@wp.kernel
-def identity_test(data: wp.array(dtype=wp.bool)):
-    i = wp.tid()
-
-    data[i] = data[i] and True
-    data[i] = data[i] and wp.bool(True)
-    data[i] = data[i] and not False
-    data[i] = data[i] and not wp.bool(False)
-    data[i] = identity_function(data[i], True)
-
-    if data[i]:
-        data[i] = True
-    else:
-        data[i] = False
-
-    if not data[i]:
-        data[i] = False
-    else:
-        data[i] = True
-
-    if data[i] and True:
-        data[i] = True
-    else:
-        data[i] = False
-
-    if data[i] or False:
-        data[i] = True
-    else:
-        data[i] = False
-
-
-def test_bool_identity_ops(test, device):
-    rng = np.random.default_rng(123)
-
-    dim_x = 10
-
-    rand_np = rng.random(dim_x) > 0.5
-
-    data_array = wp.array(data=rand_np, device=device)
-
-    test.assertEqual(data_array.dtype, wp.bool)
-
-    wp.launch(identity_test, dim=data_array.shape, inputs=[data_array], device=device)
-
-    assert_np_equal(data_array.numpy(), rand_np)
-
-
-@wp.kernel
-def check_compile_constant(result: wp.array(dtype=wp.bool)):
-    if TRUE_CONSTANT:
-        result[0] = TRUE_CONSTANT
-    else:
-        result[0] = False
-
-
-def test_bool_constant(test, device):
-    compile_constant_value = wp.zeros(1, dtype=wp.bool, device=device)
-    wp.launch(check_compile_constant, 1, inputs=[compile_constant_value], device=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]))
-
-    @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
-
-    result_array = wp.zeros(10, dtype=wp.int32, device=device)
-
-    wp.launch(sum_from_bool_vec, result_array.shape, inputs=[result_array], device=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
-
-    result_array = wp.zeros(10, dtype=wp.int32, device=device)
-
-    wp.launch(sum_from_bool_mat, result_array.shape, inputs=[result_array], device=device)
-
-    assert_np_equal(result_array.numpy(), np.full(result_array.shape, 9))
-
-
-devices = get_test_devices()
-
-
-class TestBool(unittest.TestCase):
-    pass
-
-
-add_function_test(TestBool, "test_bool_identity_ops", test_bool_identity_ops, devices=devices)
-add_function_test(TestBool, "test_bool_constant", test_bool_constant, devices=devices)
-add_function_test(TestBool, "test_bool_constant_vec", test_bool_constant_vec, devices=devices)
-add_function_test(TestBool, "test_bool_constant_mat", test_bool_constant_mat, devices=devices)
-
-
-if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
-    unittest.main(verbosity=2)
+# 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.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+TRUE_CONSTANT = wp.constant(True)
+
+
+@wp.func
+def identity_function(input_bool: wp.bool, plain_bool: bool):
+    return input_bool and plain_bool
+
+
+@wp.kernel
+def identity_test(data: wp.array(dtype=wp.bool)):
+    i = wp.tid()
+
+    data[i] = data[i] and True
+    data[i] = data[i] and wp.bool(True)
+    data[i] = data[i] and not False
+    data[i] = data[i] and not wp.bool(False)
+    data[i] = identity_function(data[i], True)
+
+    if data[i]:
+        data[i] = True
+    else:
+        data[i] = False
+
+    if not data[i]:
+        data[i] = False
+    else:
+        data[i] = True
+
+    if data[i] and True:
+        data[i] = True
+    else:
+        data[i] = False
+
+    if data[i] or False:
+        data[i] = True
+    else:
+        data[i] = False
+
+    data[i] = wp.select(data[i], False, True)
+
+
+def test_bool_identity_ops(test, device):
+    rng = np.random.default_rng(123)
+
+    dim_x = 10
+
+    rand_np = rng.random(dim_x) > 0.5
+
+    data_array = wp.array(data=rand_np, device=device)
+
+    test.assertEqual(data_array.dtype, wp.bool)
+
+    wp.launch(identity_test, dim=data_array.shape, inputs=[data_array], device=device)
+
+    assert_np_equal(data_array.numpy(), rand_np)
+
+
+@wp.kernel
+def check_compile_constant(result: wp.array(dtype=wp.bool)):
+    if TRUE_CONSTANT:
+        result[0] = TRUE_CONSTANT
+    else:
+        result[0] = False
+
+
+def test_bool_constant(test, device):
+    compile_constant_value = wp.zeros(1, dtype=wp.bool, device=device)
+    wp.launch(check_compile_constant, 1, inputs=[compile_constant_value], device=device)
+    test.assertTrue(compile_constant_value.numpy()[0])
+
+    # Repeat the comparison with dtype=bool for the array
+    compile_constant_value = wp.zeros(1, dtype=bool, device=device)
+    wp.launch(check_compile_constant, 1, inputs=[compile_constant_value], device=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]))
+
+    @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
+
+    result_array = wp.zeros(10, dtype=wp.int32, device=device)
+
+    wp.launch(sum_from_bool_vec, result_array.shape, inputs=[result_array], device=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
+
+    result_array = wp.zeros(10, dtype=wp.int32, device=device)
+
+    wp.launch(sum_from_bool_mat, result_array.shape, inputs=[result_array], device=device)
+
+    assert_np_equal(result_array.numpy(), np.full(result_array.shape, 9))
+
+
+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))
+    # Scalar initialize
+    vec3bool_s = vec3bool_type(True)
+    test.assertEqual(tuple(vec3bool_s), (True, True, True))
+    # Component-wise initialize
+    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)
+
+
+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)))
+    # Scalar initialize
+    mat22bool_s = mat22bool_type(True)
+    test.assertEqual(tuple(mat22bool_s), ((True, True), (True, True)))
+    # Component-wise initialize
+    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)
+
+
+devices = get_test_devices()
+
+
+class TestBool(unittest.TestCase):
+    pass
+
+
+add_function_test(TestBool, "test_bool_identity_ops", test_bool_identity_ops, devices=devices)
+add_function_test(TestBool, "test_bool_constant", test_bool_constant, devices=devices)
+add_function_test(TestBool, "test_bool_constant_vec", test_bool_constant_vec, devices=devices)
+add_function_test(TestBool, "test_bool_constant_mat", test_bool_constant_mat, devices=devices)
+add_function_test(TestBool, "test_bool_vec_typing", test_bool_vec_typing, devices=devices)
+add_function_test(TestBool, "test_bool_mat_typing", test_bool_mat_typing, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)