warp-lang 1.10.0__py3-none-macosx_11_0_arm64.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._src.utils import array_inner, array_sum
+from warp.tests.unittest_utils import *
+
+
+def make_test_array_sum(dtype):
+    N = 1000
+
+    def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
+        cols = wp._src.types.type_size(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._src.types.type_size(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,303 @@
+# 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")
+
+
+class TestAssertModeSwitch(unittest.TestCase):
+    """Test that switching from release mode to debug mode rebuilds the module with assertions enabled."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.config.mode
+        cls._saved_mode_module = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        # Don't set any mode initially - use whatever the default is
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.config.mode = cls._saved_mode
+        wp.set_module_options({"mode": cls._saved_mode_module})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_switch_to_debug_mode(self):
+        """Test that switching from release mode to debug mode rebuilds the module with assertions enabled."""
+        with wp.ScopedDevice("cpu"):
+            # Create an array that will trigger an assertion
+            input_array = wp.zeros(1, dtype=int)
+
+            # In default mode, this should not assert
+            capture = StdErrCapture()
+            capture.begin()
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+            output = capture.end()
+
+            # Should not have any assertion output in release mode
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr in release mode, got {output}")
+
+            # Now switch to debug mode and have it compile a new kernel
+            wp.config.mode = "debug"
+
+            @wp.kernel
+            def expect_ones_debug(a: wp.array(dtype=int)):
+                i = wp.tid()
+                assert a[i] == 1
+
+            # In debug mode, this should assert
+            capture = StdErrCapture()
+            capture.begin()
+            wp.launch(expect_ones_debug, input_array.shape, inputs=[input_array])
+            output = capture.end()
+
+            # Should have assertion output in debug mode
+            # 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")
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)