warp-lang 1.0.1__py3-none-manylinux2014_aarch64.whl → 1.1.0__py3-none-manylinux2014_aarch64.whl

warp/tests/test_options.py

@@ -1,125 +1,123 @@
-# 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 contextlib
-import io
-
-import warp as wp
-from warp.tests.unittest_utils import *
-
-wp.init()
-
-from warp.context import runtime  # noqa: E402
-
-
-@wp.kernel
-def scale(
-    x: wp.array(dtype=float),
-    y: wp.array(dtype=float),
-):
-    y[0] = x[0] ** 2.0
-
-
-@wp.kernel(enable_backward=True)
-def scale_1(
-    x: wp.array(dtype=float),
-    y: wp.array(dtype=float),
-):
-    y[0] = x[0] ** 2.0
-
-
-@wp.kernel(enable_backward=False)
-def scale_2(
-    x: wp.array(dtype=float),
-    y: wp.array(dtype=float),
-):
-    y[0] = x[0] ** 2.0
-
-
-def test_options_1(test, device):
-    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
-    y = wp.zeros_like(x)
-
-    wp.set_module_options({"enable_backward": False})
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(scale, dim=1, inputs=[x, y], device=device)
-
-    with contextlib.redirect_stdout(io.StringIO()) as f:
-        tape.backward(y)
-
-    expected = f"Warp UserWarning: Running the tape backwards may produce incorrect gradients because recorded kernel {scale.key} is defined in a module with the option 'enable_backward=False' set.\n"
-
-    assert f.getvalue() == expected
-    assert_np_equal(tape.gradients[x].numpy(), np.array(0.0))
-
-
-def test_options_2(test, device):
-    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
-    y = wp.zeros_like(x)
-
-    wp.set_module_options({"enable_backward": True})
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(scale, dim=1, inputs=[x, y], device=device)
-
-    tape.backward(y)
-    assert_np_equal(tape.gradients[x].numpy(), np.array(6.0))
-
-
-def test_options_3(test, device):
-    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
-    y = wp.zeros_like(x)
-
-    wp.set_module_options({"enable_backward": False})
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(scale_1, dim=1, inputs=[x, y], device=device)
-
-    tape.backward(y)
-    assert_np_equal(tape.gradients[x].numpy(), np.array(6.0))
-
-
-def test_options_4(test, device):
-    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
-    y = wp.zeros_like(x)
-
-    wp.set_module_options({"enable_backward": True})
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(scale_2, dim=1, inputs=[x, y], device=device)
-
-    with contextlib.redirect_stdout(io.StringIO()) as f:
-        tape.backward(y)
-
-    expected = f"Warp UserWarning: Running the tape backwards may produce incorrect gradients because recorded kernel {scale_2.key} is configured with the option 'enable_backward=False'.\n"
-
-    assert f.getvalue() == expected
-    assert_np_equal(tape.gradients[x].numpy(), np.array(0.0))
-
-
-devices = get_test_devices()
-
-
-class TestOptions(unittest.TestCase):
-    pass
-
-
-add_function_test(TestOptions, "test_options_1", test_options_1, devices=devices)
-add_function_test(TestOptions, "test_options_2", test_options_2, devices=devices)
-add_function_test(TestOptions, "test_options_3", test_options_3, devices=devices)
-add_function_test(TestOptions, "test_options_4", test_options_4, 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 contextlib
+import io
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+@wp.kernel
+def scale(
+    x: wp.array(dtype=float),
+    y: wp.array(dtype=float),
+):
+    y[0] = x[0] ** 2.0
+
+
+@wp.kernel(enable_backward=True)
+def scale_1(
+    x: wp.array(dtype=float),
+    y: wp.array(dtype=float),
+):
+    y[0] = x[0] ** 2.0
+
+
+@wp.kernel(enable_backward=False)
+def scale_2(
+    x: wp.array(dtype=float),
+    y: wp.array(dtype=float),
+):
+    y[0] = x[0] ** 2.0
+
+
+def test_options_1(test, device):
+    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
+    y = wp.zeros_like(x)
+
+    wp.set_module_options({"enable_backward": False})
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scale, dim=1, inputs=[x, y], device=device)
+
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        tape.backward(y)
+
+    expected = f"Warp UserWarning: Running the tape backwards may produce incorrect gradients because recorded kernel {scale.key} is defined in a module with the option 'enable_backward=False' set.\n"
+
+    assert f.getvalue() == expected
+    assert_np_equal(tape.gradients[x].numpy(), np.array(0.0))
+
+
+def test_options_2(test, device):
+    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
+    y = wp.zeros_like(x)
+
+    wp.set_module_options({"enable_backward": True})
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scale, dim=1, inputs=[x, y], device=device)
+
+    tape.backward(y)
+    assert_np_equal(tape.gradients[x].numpy(), np.array(6.0))
+
+
+def test_options_3(test, device):
+    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
+    y = wp.zeros_like(x)
+
+    wp.set_module_options({"enable_backward": False})
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scale_1, dim=1, inputs=[x, y], device=device)
+
+    tape.backward(y)
+    assert_np_equal(tape.gradients[x].numpy(), np.array(6.0))
+
+
+def test_options_4(test, device):
+    x = wp.array([3.0], dtype=float, requires_grad=True, device=device)
+    y = wp.zeros_like(x)
+
+    wp.set_module_options({"enable_backward": True})
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(scale_2, dim=1, inputs=[x, y], device=device)
+
+    with contextlib.redirect_stdout(io.StringIO()) as f:
+        tape.backward(y)
+
+    expected = f"Warp UserWarning: Running the tape backwards may produce incorrect gradients because recorded kernel {scale_2.key} is configured with the option 'enable_backward=False'.\n"
+
+    assert f.getvalue() == expected
+    assert_np_equal(tape.gradients[x].numpy(), np.array(0.0))
+
+
+devices = get_test_devices()
+
+
+class TestOptions(unittest.TestCase):
+    pass
+
+
+add_function_test(TestOptions, "test_options_1", test_options_1, devices=devices)
+add_function_test(TestOptions, "test_options_2", test_options_2, devices=devices)
+add_function_test(TestOptions, "test_options_3", test_options_3, devices=devices)
+add_function_test(TestOptions, "test_options_4", test_options_4, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_peer.py

@@ -1,137 +1,133 @@
-# 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 warp as wp
-from warp.tests.unittest_utils import *
-
-wp.init()
-
-
-def get_device_pair_with_peer_access_support():
-    devices = wp.get_cuda_devices()
-    for target_device in devices:
-        for peer_device in devices:
-            if target_device != peer_device:
-                if wp.is_peer_access_supported(target_device, peer_device):
-                    return (target_device, peer_device)
-    return None
-
-
-def get_device_pair_without_peer_access_support():
-    devices = wp.get_cuda_devices()
-    for target_device in devices:
-        for peer_device in devices:
-            if target_device != peer_device:
-                if not wp.is_peer_access_supported(target_device, peer_device):
-                    return (target_device, peer_device)
-    return None
-
-
-def test_peer_access_self(test, device):
-
-    device = wp.get_device(device)
-
-    assert device.is_cuda
-
-    # device can access self
-    can_access = wp.is_peer_access_supported(device, device)
-    test.assertTrue(can_access)
-
-    # setting peer access to self is a no-op
-    wp.set_peer_access_enabled(device, device, True)
-    wp.set_peer_access_enabled(device, device, False)
-
-    # should always be enabled
-    enabled = wp.is_peer_access_enabled(device, device)        
-    test.assertTrue(enabled)
-
-
-@unittest.skipUnless(get_device_pair_with_peer_access_support(), "Requires devices with peer access support")
-def test_peer_access(test, _):
-
-    target_device, peer_device = get_device_pair_with_peer_access_support()
-
-    was_enabled = wp.is_peer_access_enabled(target_device, peer_device)
-
-    if was_enabled:
-        # try disabling
-        wp.set_peer_access_enabled(target_device, peer_device, False)
-        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
-        test.assertFalse(is_enabled)
-
-        # try re-enabling
-        wp.set_peer_access_enabled(target_device, peer_device, True)
-        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
-        test.assertTrue(is_enabled)
-    else:
-        # try enabling
-        wp.set_peer_access_enabled(target_device, peer_device, True)
-        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
-        test.assertTrue(is_enabled)
-
-        # try re-disabling
-        wp.set_peer_access_enabled(target_device, peer_device, False)
-        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
-        test.assertFalse(is_enabled)
-
-
-@unittest.skipUnless(get_device_pair_without_peer_access_support(), "Requires devices without peer access support")
-def test_peer_access_exceptions_unsupported(test, _):
-
-    # get a CUDA device pair without peer access support
-    target_device, peer_device = get_device_pair_without_peer_access_support()
-
-    # querying is ok, but must return False
-    test.assertFalse(wp.is_peer_access_enabled(target_device, peer_device))
-
-    # enabling should raise RuntimeError
-    with test.assertRaises(RuntimeError):
-        wp.set_peer_access_enabled(target_device, peer_device, True)
-
-    # disabling should not raise an error
-    wp.set_peer_access_enabled(target_device, peer_device, False)
-
-
-@unittest.skipUnless(wp.is_cpu_available() and wp.is_cuda_available(), "Requires both CUDA and CPU devices")
-def test_peer_access_exceptions_cpu(test, _):
-
-    # querying is ok, but must return False
-    test.assertFalse(wp.is_peer_access_enabled("cuda:0", "cpu"))
-    test.assertFalse(wp.is_peer_access_enabled("cpu", "cuda:0"))
-
-    # enabling should raise ValueError
-    with test.assertRaises(ValueError):
-        wp.set_peer_access_enabled("cpu", "cuda:0", True)
-    with test.assertRaises(ValueError):
-        wp.set_peer_access_enabled("cuda:0", "cpu", True)
-
-    # disabling should not raise an error
-    wp.set_peer_access_enabled("cpu", "cuda:0", False)
-    wp.set_peer_access_enabled("cuda:0", "cpu", False)
-
-
-class TestPeer(unittest.TestCase):
-    pass
-
-
-cuda_test_devices = get_cuda_test_devices()
-
-add_function_test(TestPeer, "test_peer_access_self", test_peer_access_self, devices=cuda_test_devices)
-
-# peer access tests
-add_function_test(TestPeer, "test_peer_access", test_peer_access)
-
-# peer access exceptions
-add_function_test(TestPeer, "test_peer_access_exceptions_unsupported", test_peer_access_exceptions_unsupported)
-add_function_test(TestPeer, "test_peer_access_exceptions_cpu", test_peer_access_exceptions_cpu)
-
-
-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 warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+def get_device_pair_with_peer_access_support():
+    devices = wp.get_cuda_devices()
+    for target_device in devices:
+        for peer_device in devices:
+            if target_device != peer_device:
+                if wp.is_peer_access_supported(target_device, peer_device):
+                    return (target_device, peer_device)
+    return None
+
+
+def get_device_pair_without_peer_access_support():
+    devices = wp.get_cuda_devices()
+    for target_device in devices:
+        for peer_device in devices:
+            if target_device != peer_device:
+                if not wp.is_peer_access_supported(target_device, peer_device):
+                    return (target_device, peer_device)
+    return None
+
+
+def test_peer_access_self(test, device):
+    device = wp.get_device(device)
+
+    assert device.is_cuda
+
+    # device can access self
+    can_access = wp.is_peer_access_supported(device, device)
+    test.assertTrue(can_access)
+
+    # setting peer access to self is a no-op
+    wp.set_peer_access_enabled(device, device, True)
+    wp.set_peer_access_enabled(device, device, False)
+
+    # should always be enabled
+    enabled = wp.is_peer_access_enabled(device, device)
+    test.assertTrue(enabled)
+
+
+@unittest.skipUnless(get_device_pair_with_peer_access_support(), "Requires devices with peer access support")
+def test_peer_access(test, _):
+    target_device, peer_device = get_device_pair_with_peer_access_support()
+
+    was_enabled = wp.is_peer_access_enabled(target_device, peer_device)
+
+    if was_enabled:
+        # try disabling
+        wp.set_peer_access_enabled(target_device, peer_device, False)
+        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
+        test.assertFalse(is_enabled)
+
+        # try re-enabling
+        wp.set_peer_access_enabled(target_device, peer_device, True)
+        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
+        test.assertTrue(is_enabled)
+    else:
+        # try enabling
+        wp.set_peer_access_enabled(target_device, peer_device, True)
+        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
+        test.assertTrue(is_enabled)
+
+        # try re-disabling
+        wp.set_peer_access_enabled(target_device, peer_device, False)
+        is_enabled = wp.is_peer_access_enabled(target_device, peer_device)
+        test.assertFalse(is_enabled)
+
+
+@unittest.skipUnless(get_device_pair_without_peer_access_support(), "Requires devices without peer access support")
+def test_peer_access_exceptions_unsupported(test, _):
+    # get a CUDA device pair without peer access support
+    target_device, peer_device = get_device_pair_without_peer_access_support()
+
+    # querying is ok, but must return False
+    test.assertFalse(wp.is_peer_access_enabled(target_device, peer_device))
+
+    # enabling should raise RuntimeError
+    with test.assertRaises(RuntimeError):
+        wp.set_peer_access_enabled(target_device, peer_device, True)
+
+    # disabling should not raise an error
+    wp.set_peer_access_enabled(target_device, peer_device, False)
+
+
+@unittest.skipUnless(wp.is_cpu_available() and wp.is_cuda_available(), "Requires both CUDA and CPU devices")
+def test_peer_access_exceptions_cpu(test, _):
+    # querying is ok, but must return False
+    test.assertFalse(wp.is_peer_access_enabled("cuda:0", "cpu"))
+    test.assertFalse(wp.is_peer_access_enabled("cpu", "cuda:0"))
+
+    # enabling should raise ValueError
+    with test.assertRaises(ValueError):
+        wp.set_peer_access_enabled("cpu", "cuda:0", True)
+    with test.assertRaises(ValueError):
+        wp.set_peer_access_enabled("cuda:0", "cpu", True)
+
+    # disabling should not raise an error
+    wp.set_peer_access_enabled("cpu", "cuda:0", False)
+    wp.set_peer_access_enabled("cuda:0", "cpu", False)
+
+
+class TestPeer(unittest.TestCase):
+    pass
+
+
+cuda_test_devices = get_cuda_test_devices()
+
+add_function_test(TestPeer, "test_peer_access_self", test_peer_access_self, devices=cuda_test_devices)
+
+# peer access tests
+add_function_test(TestPeer, "test_peer_access", test_peer_access)
+
+# peer access exceptions
+add_function_test(TestPeer, "test_peer_access_exceptions_unsupported", test_peer_access_exceptions_unsupported)
+add_function_test(TestPeer, "test_peer_access_exceptions_cpu", test_peer_access_exceptions_cpu)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)