warp-lang 1.3.3__py3-none-manylinux2014_x86_64.whl → 1.4.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_static.py

@@ -0,0 +1,412 @@
+# Copyright (c) 2024 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
+from typing import Dict, List
+
+import numpy as np
+
+import warp
+import warp as wp
+from warp.tests.unittest_utils import *
+
+global_variable = 3
+
+
+@wp.func
+def static_global_variable_func():
+    static_var = warp.static(global_variable + 2)
+    return static_var
+
+
+@wp.kernel
+def static_global_variable_kernel(results: wp.array(dtype=int)):
+    # evaluate a constant expression at codegen time
+    static_var = static_global_variable_func()
+    const_var = 3
+    # call a function at codegen time
+    static_func_result = wp.static(static_global_variable_func() + const_var)
+    results[0] = static_var
+    results[1] = static_func_result
+
+
+@wp.struct
+class StaticallyConstructableStruct:
+    mat: wp.mat33
+    vec: wp.vec3
+    i: int
+
+
+@wp.struct
+class StaticallyConstructableNestedStruct:
+    s: StaticallyConstructableStruct
+    tf: wp.transform
+    quat: wp.quat
+
+
+@wp.func
+def construct_struct(mat: wp.mat33, vec: wp.vec3, i: int):
+    s = StaticallyConstructableStruct()
+    s.mat = mat
+    s.vec = vec
+    s.i = i
+    return s
+
+
+@wp.func
+def construct_nested_struct(mat: wp.mat33, vec: wp.vec3, i: int, tf: wp.transform, quat: wp.quat):
+    n = StaticallyConstructableNestedStruct()
+    n.s = construct_struct(mat, vec, i)
+    n.tf = tf
+    n.quat = quat
+    return n
+
+
+@wp.kernel
+def construct_static_struct_kernel(results: wp.array(dtype=StaticallyConstructableStruct)):
+    static_struct = wp.static(
+        construct_struct(
+            wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0),
+            wp.vec3(1.0, 2.0, 3.0),
+            1,
+        )
+    )
+    results[0] = static_struct
+
+
+@wp.kernel
+def construct_static_nested_struct_kernel(results: wp.array(dtype=StaticallyConstructableNestedStruct)):
+    static_struct = wp.static(
+        construct_nested_struct(
+            wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0),
+            wp.vec3(1.0, 2.0, 3.0),
+            1,
+            wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), wp.pi / 2.0)),
+            wp.quat_from_axis_angle(wp.normalize(wp.vec3(1.0, 2.0, 3.0)), wp.pi / 2.0),
+        )
+    )
+    results[0] = static_struct
+
+
+def test_static_global_variable(test, device):
+    results = wp.zeros(2, dtype=int, device=device)
+    wp.launch(static_global_variable_kernel, 1, [results], device=device)
+    assert_np_equal(results.numpy(), np.array([5, 8], dtype=int))
+
+
+def test_construct_static_struct(test, device):
+    results = wp.zeros(1, dtype=StaticallyConstructableStruct, device=device)
+    wp.launch(construct_static_struct_kernel, 1, [results], device=device)
+    results = results.numpy()
+    assert_np_equal(results[0][0], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]))
+    assert_np_equal(results[0][1], np.array([1.0, 2.0, 3.0]))
+    assert_np_equal(results[0][2], 1)
+
+
+def test_construct_static_nested_struct(test, device):
+    results = wp.zeros(1, dtype=StaticallyConstructableNestedStruct, device=device)
+    wp.launch(construct_static_nested_struct_kernel, 1, [results], device=device)
+    results = results.numpy()
+
+    tf = wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), wp.pi / 2.0))
+    quat = wp.quat_from_axis_angle(wp.normalize(wp.vec3(1.0, 2.0, 3.0)), wp.pi / 2.0)
+
+    assert_np_equal(results[0][0][0], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]))
+    assert_np_equal(results[0][0][1], np.array([1.0, 2.0, 3.0]))
+    assert_np_equal(results[0][0][2], 1)
+    assert_np_equal(results[0][1], np.array(tf))
+    assert_np_equal(results[0][2], np.array(quat))
+
+
+def test_invalid_static_expression(test, device):
+    @wp.kernel
+    def invalid_kernel():
+        wp.static(1.0 / 0.0)
+
+    with test.assertRaisesRegex(
+        warp.codegen.WarpCodegenError, r"Error evaluating static expression\: float division by zero"
+    ):
+        wp.launch(invalid_kernel, 1, device=device)
+
+    @wp.kernel
+    def invalid_kernel(i: int):
+        wp.static(i * 2)
+
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"Error evaluating static expression\: name 'i' is not defined\. Make sure all variables used in the static expression are constant\.",
+    ):
+        wp.launch(invalid_kernel, 1, device=device, inputs=[3])
+
+
+def test_static_expression_return_types(test, device):
+    @wp.kernel
+    def invalid_kernel():
+        wp.static(wp.zeros(3, device=device))
+
+    with test.assertRaisesRegex(
+        warp.codegen.WarpCodegenError,
+        r"Static expression returns an unsupported value\: a Warp array cannot be created inside Warp kernels",
+    ):
+        wp.launch(invalid_kernel, 1, device=device)
+
+    @wp.struct
+    class Baz:
+        data: wp.array(dtype=int)
+        z: wp.vec3
+
+    @wp.struct
+    class Bar:
+        baz: Baz
+        y: float
+
+    @wp.struct
+    class Foo:
+        bar: Bar
+        x: int
+
+    def create_struct():
+        foo = Foo()
+        foo.bar = Bar()
+        foo.bar.baz = Baz()
+        foo.bar.baz.data = wp.zeros(3, dtype=int, device=device)
+        foo.bar.baz.z = wp.vec3(1, 2, 3)
+        foo.bar.y = 1.23
+        foo.x = 123
+        return foo
+
+    @wp.kernel
+    def invalid_kernel():
+        wp.static(create_struct())
+
+    with test.assertRaisesRegex(
+        warp.codegen.WarpCodegenError,
+        r"Static expression returns an unsupported value: the returned Warp struct contains a data type that cannot be constructed inside Warp kernels\: a Warp array cannot be created inside Warp kernels at .*?Foo\.bar\.baz",
+    ):
+        wp.launch(invalid_kernel, 1, device=device)
+
+    def function_with_no_return_value():
+        pass
+
+    @wp.kernel
+    def invalid_kernel():
+        wp.static(function_with_no_return_value())
+
+    with test.assertRaisesRegex(
+        warp.codegen.WarpCodegenError,
+        r"Static expression returns an unsupported value\: None is returned",
+    ):
+        wp.launch(invalid_kernel, 1, device=device)
+
+    class MyClass:
+        pass
+
+    @wp.kernel
+    def invalid_kernel():
+        wp.static(MyClass())
+
+    with test.assertRaisesRegex(
+        warp.codegen.WarpCodegenError,
+        r"Static expression returns an unsupported value\: value of type .*?MyClass",
+    ):
+        wp.launch(invalid_kernel, 1, device=device)
+
+
+def test_function_variable(test, device):
+    # create a function and pass it in as a static variable to the kernel
+    @wp.func
+    def func1(a: int, b: int):
+        return a + b
+
+    @wp.func
+    def func2(a: int, b: int):
+        return a - b
+
+    for func in [func1, func2]:
+        # note that this example also works without using wp.static()
+
+        @wp.kernel
+        def function_variable_kernel(results: wp.array(dtype=int)):
+            results[0] = wp.static(func)(3, 2)  # noqa: B023
+
+        results = wp.zeros(1, dtype=int, device=device)
+        # note that the kernel has to be recompiled everytime the value of func changes
+        wp.launch(function_variable_kernel, 1, [results], device=device)
+        assert_np_equal(results.numpy(), np.array([func(3, 2)], dtype=int))
+
+
+def test_function_lookup(test, device):
+    @wp.func
+    def do_add(a: float, b: float):
+        return a + b
+
+    @wp.func
+    def do_sub(a: float, b: float):
+        return a - b
+
+    @wp.func
+    def do_mul(a: float, b: float):
+        return a * b
+
+    op_handlers = {
+        "add": do_add,
+        "sub": do_sub,
+        "mul": do_mul,
+    }
+
+    inputs = wp.array([[1, 2], [3, 0]], dtype=wp.float32)
+
+    outputs = wp.empty(2, dtype=wp.float32)
+
+    for op in op_handlers.keys():
+
+        @wp.kernel
+        def operate(input: wp.array(dtype=inputs.dtype, ndim=2), output: wp.array(dtype=wp.float32)):
+            tid = wp.tid()
+            a, b = input[tid, 0], input[tid, 1]
+            # retrieve the right function to use for the captured dtype variable
+            output[tid] = wp.static(op_handlers[op])(a, b)  # noqa: B023
+
+        wp.launch(operate, dim=2, inputs=[inputs], outputs=[outputs])
+        outputs_np = outputs.numpy()
+        inputs_np = inputs.numpy()
+        for i in range(len(outputs_np)):
+            test.assertEqual(outputs_np[i], op_handlers[op](float(inputs_np[i][0]), float(inputs_np[i][1])))
+
+
+def count_ssa_occurrences(kernel: wp.Kernel, ssas: List[str]) -> Dict[str, int]:
+    # analyze the generated code
+    counts = {ssa: 0 for ssa in ssas}
+    for line in kernel.adj.blocks[0].body_forward:
+        for ssa in ssas:
+            if ssa in line:
+                counts[ssa] += 1
+    return counts
+
+
+def test_static_for_loop(test, device):
+    @wp.kernel
+    def static_loop_variable(results: wp.array(dtype=int)):
+        s = 0
+        for i in range(wp.static(static_global_variable_func())):
+            s += wp.static(i)
+        results[0] = s
+
+    wp.set_module_options(
+        options={"max_unroll": static_global_variable_func()},
+    )
+
+    results = wp.zeros(1, dtype=int, device=device)
+    wp.launch(static_loop_variable, 1, [results], device=device)
+    results = results.numpy()
+
+    s = 0
+    for i in range(wp.static(static_global_variable_func())):
+        s += wp.static(i)
+
+    test.assertEqual(results[0], s, "Static for loop has to compute the correct solution")
+
+    # analyze the generated code
+    if hasattr(static_loop_variable.adj, "blocks"):
+        counts = count_ssa_occurrences(static_loop_variable, ["add", "for"])
+
+        test.assertEqual(counts["add"], static_global_variable_func(), "Static for loop must be unrolled")
+        # there is just one occurrence of "for" in the comment referring to the original Python code
+        test.assertEqual(counts["for"], 1, "Static for loop must be unrolled")
+
+
+def test_static_if_else_elif(test, device):
+    @wp.kernel
+    def static_condition1(results: wp.array(dtype=int)):
+        if wp.static(static_global_variable_func() in {2, 3, 5}):
+            results[0] = 1
+        elif wp.static(static_global_variable_func() in {0, 1}):
+            results[0] = 2
+        else:
+            results[0] = 3
+
+    results = wp.zeros(1, dtype=int, device=device)
+    wp.launch(static_condition1, 1, [results], device=device)
+    results = results.numpy()
+    assert_np_equal(results[0], 1)
+    # TODO this needs fixing to ensure we can run these tests multiple times
+    if hasattr(static_condition1.adj, "blocks"):
+        counts = count_ssa_occurrences(static_condition1, ["if", "else"])
+
+        # if, else, elif can appear as comments but the generated code must not contain
+        # such keywords since the conditions are resolved at the time of code generation
+        assert_np_equal(counts["if"], 1)
+        assert_np_equal(counts["else"], 0)
+
+    captured_var = "hello"
+
+    @wp.kernel
+    def static_condition2(results: wp.array(dtype=int)):
+        if wp.static(captured_var == "world"):
+            results[0] = 1
+        else:
+            results[0] = 2
+
+    results = wp.zeros(1, dtype=int, device=device)
+    wp.launch(static_condition2, 1, [results], device=device)
+    results = results.numpy()
+    assert_np_equal(results[0], 2)
+    if hasattr(static_condition2.adj, "blocks"):
+        counts = count_ssa_occurrences(static_condition2, ["if", "else"])
+        assert_np_equal(counts["if"], 1)
+        assert_np_equal(counts["else"], 0)
+
+    my_list = [1, 2, 3]
+
+    @wp.kernel
+    def static_condition3(results: wp.array(dtype=int)):
+        if wp.static(len(my_list) == 0):
+            results[0] = 0
+        elif wp.static(len(my_list) == 1):
+            results[0] = 1
+        elif wp.static(len(my_list) == 2):
+            results[0] = 2
+        elif wp.static(len(my_list) == 3):
+            results[0] = 3
+
+    results = wp.zeros(1, dtype=int, device=device)
+    wp.launch(static_condition3, 1, [results], device=device)
+    results = results.numpy()
+    assert_np_equal(results[0], 3)
+    if hasattr(static_condition3.adj, "blocks"):
+        counts = count_ssa_occurrences(static_condition3, ["if", "else"])
+        assert_np_equal(counts["if"], 4)
+        assert_np_equal(counts["else"], 0)
+
+
+devices = get_test_devices()
+
+
+class TestStatic(unittest.TestCase):
+    def test_static_python_call(self):
+        # ensure wp.static() works from a Python context
+        self.assertEqual(static_global_variable_func(), 5)
+
+
+add_function_test(TestStatic, "test_static_global_variable", test_static_global_variable, devices=devices)
+add_function_test(TestStatic, "test_construct_static_struct", test_construct_static_struct, devices=devices)
+add_function_test(
+    TestStatic, "test_construct_static_nested_struct", test_construct_static_nested_struct, devices=devices
+)
+add_function_test(TestStatic, "test_function_variable", test_function_variable, devices=devices)
+add_function_test(TestStatic, "test_function_lookup", test_function_lookup, devices=devices)
+add_function_test(TestStatic, "test_invalid_static_expression", test_invalid_static_expression, devices=devices)
+add_function_test(
+    TestStatic, "test_static_expression_return_types", test_static_expression_return_types, devices=devices
+)
+add_function_test(TestStatic, "test_static_for_loop", test_static_for_loop, devices=devices)
+add_function_test(TestStatic, "test_static_if_else_elif", test_static_if_else_elif, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_streams.py

@@ -11,7 +11,7 @@ import numpy as np
 
 import warp as wp
 from warp.tests.unittest_utils import *
-from warp.utils import check_iommu
+from warp.utils import check_p2p
 
 
 @wp.kernel
@@ -334,6 +334,65 @@ def test_event_elapsed_time(test, device):
     test.assertGreater(elapsed, 0)
 
 
+def test_stream_priority_basics(test, device):
+    standard_stream = wp.Stream(device)
+    test.assertEqual(standard_stream.priority, 0, "Default priority of streams must be 0.")
+
+    # Create a high-priority stream with a priority value that is smaller than -1 (clamping expected)
+    stream_hi = wp.Stream(device, priority=-100)
+
+    # Create a low-priority stream with a priority value that is greter than 0 (clamping expected)
+    stream_lo = wp.Stream(device, priority=100)
+
+    if stream_lo.priority == stream_hi.priority:
+        test.skipTest("Device must support stream priorities.")
+
+    test.assertEqual(stream_hi.priority, -1)
+
+    test.assertEqual(stream_lo.priority, 0)
+
+    with test.assertRaises(TypeError):
+        stream_invalid_priority = wp.Stream(device, priority=0.5)
+
+
+def test_stream_priority_timings(test, device):
+    total_size = 256 * 1024 * 1024
+    each_size = 128 * 1024 * 1024
+
+    array_lo = wp.zeros(total_size, dtype=wp.float32, device=device)
+    array_hi = wp.zeros(total_size, dtype=wp.float32, device=device)
+
+    stream_lo = wp.Stream(device, 0)
+    stream_hi = wp.Stream(device, -1)
+
+    if stream_lo.priority == stream_hi.priority:
+        test.skipTest("Device must support stream priorities.")
+
+    # Create some events
+    start_lo_event = wp.Event(device, enable_timing=True)
+    start_hi_event = wp.Event(device, enable_timing=True)
+    end_lo_event = wp.Event(device, enable_timing=True)
+    end_hi_event = wp.Event(device, enable_timing=True)
+
+    wp.synchronize_device(device)
+
+    stream_lo.record_event(start_lo_event)
+    stream_hi.record_event(start_hi_event)
+
+    for copy_offset in range(0, total_size, each_size):
+        wp.copy(array_lo, array_lo, copy_offset, copy_offset, each_size, stream_lo)
+        wp.copy(array_hi, array_hi, copy_offset, copy_offset, each_size, stream_hi)
+
+    stream_lo.record_event(end_lo_event)
+    stream_hi.record_event(end_hi_event)
+
+    # get elapsed time between the two events
+    elapsed_lo = wp.get_event_elapsed_time(start_lo_event, end_lo_event)
+    elapsed_hi = wp.get_event_elapsed_time(start_hi_event, end_hi_event)
+
+    test.assertLess(elapsed_hi, elapsed_lo, "Copies on higher-priority stream should be faster.")
+
+
 devices = get_selected_cuda_test_devices()
 
 
@@ -359,7 +418,7 @@ class TestStreams(unittest.TestCase):
             cpu_stream = cpu_device.stream  # noqa: F841
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
-    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
+    @unittest.skipUnless(check_p2p(), "Peer-to-Peer transfers not supported")
     def test_stream_arg_graph_mgpu(self):
         wp.load_module(device="cuda:0")
         wp.load_module(device="cuda:1")
@@ -409,7 +468,7 @@ class TestStreams(unittest.TestCase):
             assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
-    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
+    @unittest.skipUnless(check_p2p(), "Peer-to-Peer transfers not supported")
     def test_stream_scope_graph_mgpu(self):
         wp.load_module(device="cuda:0")
         wp.load_module(device="cuda:1")
@@ -485,6 +544,8 @@ add_function_test(TestStreams, "test_stream_arg_wait_stream", test_stream_arg_wa
 add_function_test(TestStreams, "test_stream_scope_synchronize", test_stream_scope_synchronize, devices=devices)
 add_function_test(TestStreams, "test_stream_scope_wait_event", test_stream_scope_wait_event, devices=devices)
 add_function_test(TestStreams, "test_stream_scope_wait_stream", test_stream_scope_wait_stream, devices=devices)
+add_function_test(TestStreams, "test_stream_priority_basics", test_stream_priority_basics, devices=devices)
+add_function_test(TestStreams, "test_stream_priority_timings", test_stream_priority_timings, devices=devices)
 
 add_function_test(TestStreams, "test_event_synchronize", test_event_synchronize, devices=devices)
 add_function_test(TestStreams, "test_event_elapsed_time", test_event_elapsed_time, devices=devices)

warp/tests/test_struct.py

@@ -589,7 +589,7 @@ def test_dependent_module_import(c: DependentModuleImport_C):
     wp.tid()  # nop, we're just testing codegen
 
 
-def test_struct_array_content_hash(test, device):
+def test_struct_array_hash(test, device):
     # Ensure that the memory address of the struct does not affect the content hash
 
     @wp.struct
@@ -611,7 +611,7 @@ def test_struct_array_content_hash(test, device):
     def dummy_kernel(a: wp.array(dtype=ContentHashStruct)):
         i = wp.tid()
 
-    module_hash_1 = wp.get_module(dummy_kernel.__module__).hash_module(recompute_content_hash=True)
+    module_hash_1 = wp.get_module(dummy_kernel.__module__).hash_module()
 
     test.assertEqual(
         module_hash_1,
@@ -628,7 +628,7 @@ def test_struct_array_content_hash(test, device):
     def dummy_kernel(a: wp.array(dtype=ContentHashStruct)):
         i = wp.tid()
 
-    module_hash_2 = wp.get_module(dummy_kernel.__module__).hash_module(recompute_content_hash=True)
+    module_hash_2 = wp.get_module(dummy_kernel.__module__).hash_module()
 
     test.assertNotEqual(
         module_hash_2, module_hash_0, "Module hash should be different when ContentHashStruct redefined and changed."
@@ -718,7 +718,7 @@ add_kernel_test(
     devices=devices,
 )
 
-add_function_test(TestStruct, "test_struct_array_content_hash", test_struct_array_content_hash, devices=None)
+add_function_test(TestStruct, "test_struct_array_hash", test_struct_array_hash, devices=None)
 
 
 if __name__ == "__main__":

warp/tests/test_torch.py

@@ -382,6 +382,27 @@ def test_array_ctype_from_torch(test, device):
     wrap_vec_tensor_with_warp_grad(wp.transform)
 
 
+def test_cuda_array_interface(test, device):
+    # We should be able to construct Torch tensors from Warp arrays via __cuda_array_interface__ on GPU.
+    # Note that Torch does not support __array_interface__ on CPU.
+
+    torch_device = wp.device_to_torch(device)
+    n = 10
+
+    # test the types supported by both Warp and Torch
+    scalar_types = [wp.float16, wp.float32, wp.float64, wp.int8, wp.int16, wp.int32, wp.int64, wp.uint8]
+
+    for dtype in scalar_types:
+        # test round trip
+        a1 = wp.zeros(n, dtype=dtype, device=device)
+        t = torch.tensor(a1, device=torch_device)
+        a2 = wp.array(t, device=device)
+
+        assert a1.dtype == a2.dtype
+        assert a1.shape == a2.shape
+        assert a1.strides == a2.strides
+
+
 def test_to_torch(test, device):
     import torch
 
@@ -918,6 +939,9 @@ try:
             test_warp_graph_torch_stream,
             devices=torch_compatible_cuda_devices,
         )
+        add_function_test(
+            TestTorch, "test_cuda_array_interface", test_cuda_array_interface, devices=torch_compatible_cuda_devices
+        )
 
     # multi-GPU tests
     if len(torch_compatible_cuda_devices) > 1: