warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/tests/test_array_reduce.py

@@ -0,0 +1,156 @@
+# 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 unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+from warp.utils import array_inner, array_sum
+
+
+def make_test_array_sum(dtype):
+    N = 1000
+
+    def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
+        cols = wp.types.type_length(dtype)
+
+        values_np = rng.random(size=(N, cols))
+        values = wp.array(values_np, device=device, dtype=dtype)
+
+        vsum = array_sum(values)
+        ref_vsum = values_np.sum(axis=0)
+
+        assert_np_equal(vsum / N, ref_vsum / N, 0.0001)
+
+    return test_array_sum
+
+
+def make_test_array_sum_axis(dtype):
+    I = 5
+    J = 10
+    K = 2
+
+    N = I * J * K
+
+    def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
+        values_np = rng.random(size=(I, J, K))
+        values = wp.array(values_np, shape=(I, J, K), device=device, dtype=dtype)
+
+        for axis in range(3):
+            vsum = array_sum(values, axis=axis)
+            ref_vsum = values_np.sum(axis=axis)
+
+            assert_np_equal(vsum.numpy() / N, ref_vsum / N, 0.0001)
+
+    return test_array_sum
+
+
+def test_array_sum_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    assert_np_equal(array_sum(values), np.zeros(2))
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_sum(values, axis=0).numpy(), np.zeros((1, 3)))
+
+
+def make_test_array_inner(dtype):
+    N = 1000
+
+    def test_array_inner(test, device):
+        rng = np.random.default_rng(123)
+
+        cols = wp.types.type_length(dtype)
+
+        a_np = rng.random(size=(N, cols))
+        b_np = rng.random(size=(N, cols))
+
+        a = wp.array(a_np, device=device, dtype=dtype)
+        b = wp.array(b_np, device=device, dtype=dtype)
+
+        ab = array_inner(a, b)
+        ref_ab = np.dot(a_np.flatten(), b_np.flatten())
+
+        test.assertAlmostEqual(ab / N, ref_ab / N, places=5)
+
+    return test_array_inner
+
+
+def make_test_array_inner_axis(dtype):
+    I = 5
+    J = 10
+    K = 2
+
+    N = I * J * K
+
+    def test_array_inner(test, device):
+        rng = np.random.default_rng(123)
+
+        a_np = rng.random(size=(I, J, K))
+        b_np = rng.random(size=(I, J, K))
+
+        a = wp.array(a_np, shape=(I, J, K), device=device, dtype=dtype)
+        b = wp.array(b_np, shape=(I, J, K), device=device, dtype=dtype)
+
+        ab = array_inner(a, b, axis=0)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [1, 2])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+        ab = array_inner(a, b, axis=1)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [0, 2])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+        ab = array_inner(a, b, axis=2)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [0, 1])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+    return test_array_inner
+
+
+def test_array_inner_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    test.assertEqual(array_inner(values, values), 0.0)
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_inner(values, values, axis=0).numpy(), np.zeros((1, 3)))
+
+
+devices = get_test_devices()
+
+
+class TestArrayReduce(unittest.TestCase):
+    pass
+
+
+add_function_test(TestArrayReduce, "test_array_sum_double", make_test_array_sum(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_vec3", make_test_array_sum(wp.vec3), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_axis_float", make_test_array_sum_axis(wp.float32), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_empty", test_array_sum_empty, devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_double", make_test_array_inner(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_vec3", make_test_array_inner(wp.vec3), devices=devices)
+add_function_test(
+    TestArrayReduce, "test_array_inner_axis_float", make_test_array_inner_axis(wp.float32), devices=devices
+)
+add_function_test(TestArrayReduce, "test_array_inner_empty", test_array_inner_empty, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_assert.py

@@ -0,0 +1,250 @@
+# 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
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def expect_ones(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1
+
+
+@wp.kernel
+def expect_ones_with_msg(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_compound(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] > 0 and a[i] < 2
+
+
+@wp.func
+def expect_ones_function(value: int):
+    assert value == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_call_function(a: wp.array(dtype=int)):
+    i = wp.tid()
+    expect_ones_function(a[i])
+
+
+class TestAssertRelease(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in release mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.config.mode = "release"
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+
+# NOTE: Failed assertions on CUDA devices leaves the CUDA context in an unrecoverable state,
+# so we currently do not test them.
+class TestAssertDebug(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in debug mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.set_module_options({"mode": "debug"})
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1")
+
+    def test_basic_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1.*Array element must be 1")
+
+    def test_compound_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] > 0 and a\[i\] < 2")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert value == 1.*Array element must be 1")
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_atomic.py

@@ -0,0 +1,153 @@
+# 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
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+# 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)
+
+        elif type == wp.float64:
+            base = rng.random(size=1, dtype=np.float64)
+            val = rng.random(size=n, dtype=np.float64)
+
+        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_double = make_atomic_test(wp.float64)
+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_double", test_atomic_double, 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.clear_kernel_cache()
+    unittest.main(verbosity=2)