warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_large.py

@@ -0,0 +1,141 @@
+# 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 math
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+@wp.kernel
+def conditional_sum(result: wp.array(dtype=wp.uint64)):
+    i, j, k = wp.tid()
+
+    if i == 0:
+        wp.atomic_add(result, 0, wp.uint64(1))
+
+
+def test_large_launch_large_kernel(test, device):
+    """Test tid() on kernel launch of 2**33 threads.
+
+    The function conditional sum will add 1 to result for every thread that has an i index of 0.
+    Due to the size of the grid, this test is not run on CPUs
+    """
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+
+    large_dim_length = 2**16
+    half_result = large_dim_length * large_dim_length
+
+    wp.launch(kernel=conditional_sum, dim=[2, large_dim_length, large_dim_length], inputs=[test_result], device=device)
+    test.assertEqual(test_result.numpy()[0], half_result)
+
+
+@wp.kernel
+def count_elements(result: wp.array(dtype=wp.uint64)):
+    wp.atomic_add(result, 0, wp.uint64(1))
+
+
+def test_large_launch_max_blocks(test, device):
+    # Loop over 1000x1x1 elements using a grid of 256 threads
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+    wp.launch(count_elements, (1000,), inputs=[test_result], max_blocks=1, device=device)
+    test.assertEqual(test_result.numpy()[0], 1000)
+
+    # Loop over 2x10x10 elements using a grid of 256 threads, using the tid() index to count half the elements
+    test_result.zero_()
+    wp.launch(
+        conditional_sum,
+        (
+            2,
+            50,
+            10,
+        ),
+        inputs=[test_result],
+        max_blocks=1,
+        device=device,
+    )
+    test.assertEqual(test_result.numpy()[0], 500)
+
+
+def test_large_launch_very_large_kernel(test, device):
+    """Due to the size of the grid, this test is not run on CPUs"""
+
+    # Dim is chosen to be larger than the maximum CUDA one-dimensional grid size (total threads)
+    dim = (2**31 - 1) * 256 + 1
+    test_result = wp.zeros(shape=(1,), dtype=wp.uint64, device=device)
+    wp.launch(count_elements, (dim,), inputs=[test_result], device=device)
+    test.assertEqual(test_result.numpy()[0], dim)
+
+
+def test_large_arrays_slow(test, device):
+    # The goal of this test is to use arrays just large enough to know
+    # if there's a flaw in handling arrays with more than 2**31-1 elements
+    # Unfortunately, it takes a long time to run so it won't be run automatically
+    # without changes to support how frequently a test may be run
+    total_elements = 2**31 + 8
+
+    # 1-D to 4-D arrays: test zero_, fill_, then zero_ for scalar data types:
+    for total_dims in range(1, 5):
+        dim_x = math.ceil(total_elements ** (1 / total_dims))
+        shape_tuple = tuple([dim_x] * total_dims)
+
+        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+            a1 = wp.zeros(shape_tuple, dtype=wptype, device=device)
+            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+            a1.fill_(127)
+            assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
+
+            a1.zero_()
+            assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+
+def test_large_arrays_fast(test, device):
+    # A truncated version of test_large_arrays_slow meant to catch basic errors
+    total_elements = 2**31 + 8
+
+    nptype = np.dtype(np.int8)
+    wptype = wp.types.np_dtype_to_warp_type[nptype]
+
+    a1 = wp.zeros((total_elements,), dtype=wptype, device=device)
+    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+    a1.fill_(127)
+    assert_np_equal(a1.numpy(), 127 * np.ones_like(a1.numpy()))
+
+    a1.zero_()
+    assert_np_equal(a1.numpy(), np.zeros_like(a1.numpy()))
+
+
+devices = get_test_devices()
+
+
+class TestLarge(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestLarge, "test_large_launch_large_kernel", test_large_launch_large_kernel, devices=get_unique_cuda_test_devices()
+)
+
+add_function_test(TestLarge, "test_large_launch_max_blocks", test_large_launch_max_blocks, devices=devices)
+add_function_test(
+    TestLarge,
+    "test_large_launch_very_large_kernel",
+    test_large_launch_very_large_kernel,
+    devices=get_unique_cuda_test_devices(),
+)
+
+add_function_test(TestLarge, "test_large_arrays_fast", test_large_arrays_fast, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_launch.py

@@ -5,14 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -74,20 +72,258 @@ def test4d(test, device):
     wp.launch(kernel4d, dim=a.shape, inputs=[wp.array(a, dtype=int, device=device)], device=device)
 
 
-def register(parent):
-    devices = get_test_devices()
+@wp.struct
+class Params:
+    a: wp.array(dtype=int)
+    i: int
+    f: float
+
+
+@wp.kernel
+def kernel_cmd(params: Params, i: int, f: float, v: wp.vec3, m: wp.mat33, out: wp.array(dtype=int)):
+    tid = wp.tid()
+
+    wp.expect_eq(params.i, i)
+    wp.expect_eq(params.f, f)
+
+    wp.expect_eq(i, int(f))
+
+    wp.expect_eq(v[0], f)
+    wp.expect_eq(v[1], f)
+    wp.expect_eq(v[2], f)
+
+    wp.expect_eq(m[0, 0], f)
+    wp.expect_eq(m[1, 1], f)
+    wp.expect_eq(m[2, 2], f)
+
+    out[tid] = tid + i
+
+
+def test_launch_cmd(test, device):
+    n = 1
+
+    ref = np.arange(0, n)
+    out = wp.zeros(n, dtype=int, device=device)
+
+    params = Params()
+    params.i = 1
+    params.f = 1.0
+
+    v = wp.vec3(params.f, params.f, params.f)
+
+    m = wp.mat33(params.f, 0.0, 0.0, 0.0, params.f, 0.0, 0.0, 0.0, params.f)
+
+    # standard launch
+    wp.launch(kernel_cmd, dim=n, inputs=[params, params.i, params.f, v, m, out], device=device)
+
+    assert_np_equal(out.numpy(), ref + params.i)
+
+    # cmd launch
+    out.zero_()
+
+    cmd = wp.launch(kernel_cmd, dim=n, inputs=[params, params.i, params.f, v, m, out], device=device, record_cmd=True)
+
+    cmd.launch()
+
+    assert_np_equal(out.numpy(), ref + params.i)
+
+
+def test_launch_cmd_set_param(test, device):
+    n = 1
+
+    ref = np.arange(0, n)
+
+    params = Params()
+    v = wp.vec3()
+    m = wp.mat33()
+
+    cmd = wp.launch(kernel_cmd, dim=n, inputs=[params, 0, 0.0, v, m, None], device=device, record_cmd=True)
+
+    # cmd param modification
+    out = wp.zeros(n, dtype=int, device=device)
+
+    params.i = 13
+    params.f = 13.0
+
+    v = wp.vec3(params.f, params.f, params.f)
+
+    m = wp.mat33(params.f, 0.0, 0.0, 0.0, params.f, 0.0, 0.0, 0.0, params.f)
+
+    cmd.set_param_at_index(0, params)
+    cmd.set_param_at_index(1, params.i)
+    cmd.set_param_at_index(2, params.f)
+    cmd.set_param_at_index(3, v)
+    cmd.set_param_at_index(4, m)
+    cmd.set_param_by_name("out", out)
+
+    cmd.launch()
+
+    assert_np_equal(out.numpy(), ref + params.i)
+
+    # test changing params after launch directly
+    # because we now cache the ctypes object inside the wp.struct
+    # instance  the command buffer will be automatically updated
+    params.i = 14
+    params.f = 14.0
+
+    v = wp.vec3(params.f, params.f, params.f)
+
+    m = wp.mat33(params.f, 0.0, 0.0, 0.0, params.f, 0.0, 0.0, 0.0, params.f)
+
+    # this is the line we explicitly leave out to
+    # ensure that param changes are reflected in the launch
+    # launch.set_param_at_index(0, params)
+
+    cmd.set_param_at_index(1, params.i)
+    cmd.set_param_at_index(2, params.f)
+    cmd.set_param_at_index(3, v)
+    cmd.set_param_at_index(4, m)
+    cmd.set_param_by_name("out", out)
+
+    cmd.launch()
+
+    assert_np_equal(out.numpy(), ref + params.i)
+
+
+def test_launch_cmd_set_ctype(test, device):
+    n = 1
+
+    ref = np.arange(0, n)
+
+    params = Params()
+    v = wp.vec3()
+    m = wp.mat33()
+
+    cmd = wp.launch(kernel_cmd, dim=n, inputs=[params, 0, 0.0, v, m, None], device=device, record_cmd=True)
+
+    # cmd param modification
+    out = wp.zeros(n, dtype=int, device=device)
+
+    # cmd param modification
+    out.zero_()
+
+    params.i = 13
+    params.f = 13.0
+
+    v = wp.vec3(params.f, params.f, params.f)
+
+    m = wp.mat33(params.f, 0.0, 0.0, 0.0, params.f, 0.0, 0.0, 0.0, params.f)
+
+    cmd.set_param_at_index_from_ctype(0, params.__ctype__())
+    cmd.set_param_at_index_from_ctype(1, params.i)
+    cmd.set_param_at_index_from_ctype(2, params.f)
+    cmd.set_param_at_index_from_ctype(3, v)
+    cmd.set_param_at_index_from_ctype(4, m)
+    cmd.set_param_by_name_from_ctype("out", out.__ctype__())
+
+    cmd.launch()
+
+    assert_np_equal(out.numpy(), ref + params.i)
+
+
+@wp.kernel
+def arange(out: wp.array(dtype=int)):
+    tid = wp.tid()
+    out[tid] = tid
+
+
+def test_launch_cmd_set_dim(test, device):
+    n = 10
+
+    ref = np.arange(0, n, dtype=int)
+    out = wp.zeros(n, dtype=int, device=device)
+
+    cmd = wp.launch(arange, dim=n, inputs=[out], device=device, record_cmd=True)
+
+    cmd.set_dim(5)
+    cmd.launch()
+
+    # check first half the array is filled while rest is still zero
+    assert_np_equal(out.numpy()[0:5], ref[0:5])
+    assert_np_equal(out.numpy()[5:], np.zeros(5))
+
+    out.zero_()
+
+    cmd.set_dim(10)
+    cmd.launch()
+
+    # check the whole array was filled
+    assert_np_equal(out.numpy(), ref)
+
+
+def test_launch_cmd_empty(test, device):
+    n = 10
+
+    ref = np.arange(0, n, dtype=int)
+    out = wp.zeros(n, dtype=int, device=device)
+
+    cmd = wp.Launch(arange, device)
+    cmd.set_dim(5)
+    cmd.set_param_by_name("out", out)
+
+    cmd.launch()
+
+    # check first half the array is filled while rest is still zero
+    assert_np_equal(out.numpy()[0:5], ref[0:5])
+    assert_np_equal(out.numpy()[5:], np.zeros(5))
+
+    out.zero_()
+
+    cmd.set_dim(10)
+    cmd.launch()
+
+    # check the whole array was filled
+    assert_np_equal(out.numpy(), ref)
+
+
+@wp.kernel
+def kernel_mul(
+    values: wp.array(dtype=int),
+    coeff: int,
+    out: wp.array(dtype=int),
+):
+    tid = wp.tid()
+    out[tid] = values[tid] * coeff
+
+
+def test_launch_tuple_args(test, device):
+    values = wp.array(np.arange(0, 4), dtype=int, device=device)
+    coeff = 3
+    out = wp.empty_like(values)
+
+    wp.launch(
+        kernel_mul,
+        dim=len(values),
+        inputs=(
+            values,
+            coeff,
+        ),
+        outputs=(out,),
+        device=device,
+    )
+
+    assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
+
+
+devices = get_test_devices()
+
+
+class TestLaunch(unittest.TestCase):
+    pass
 
-    class TestLaunch(parent):
-        pass
 
-    add_function_test(TestLaunch, "test_1d_launch", test1d, devices=devices)
-    add_function_test(TestLaunch, "test_2d_launch", test2d, devices=devices)
-    add_function_test(TestLaunch, "test_3d_launch", test3d, devices=devices)
-    add_function_test(TestLaunch, "test_4d_launch", test4d, devices=devices)
+add_function_test(TestLaunch, "test_launch_1d", test1d, devices=devices)
+add_function_test(TestLaunch, "test_launch_2d", test2d, devices=devices)
+add_function_test(TestLaunch, "test_launch_3d", test3d, devices=devices)
+add_function_test(TestLaunch, "test_launch_4d", test4d, devices=devices)
 
-    return TestLaunch
+add_function_test(TestLaunch, "test_launch_cmd", test_launch_cmd, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_param, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_lerp.py

@@ -5,14 +5,16 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
 from dataclasses import dataclass
 from typing import Any
-import unittest
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
+
+wp.init()
 
 
 @dataclass
@@ -162,8 +164,6 @@ TEST_DATA = {
     ),
 }
 
-wp.init()
-
 
 def test_lerp(test, device):
     def make_kernel_fn(data_type):
@@ -179,84 +179,83 @@ def test_lerp(test, device):
 
     for data_type in TEST_DATA:
         kernel_fn = make_kernel_fn(data_type)
-        module = wp.get_module(kernel_fn.__module__)
         kernel = wp.Kernel(
             func=kernel_fn,
             key=f"test_lerp_{data_type.__name__}_kernel",
-            module=module,
         )
 
-        for test_data in TEST_DATA[data_type]:
-            a = wp.array(
-                [test_data.a],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            b = wp.array(
-                [test_data.b],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            t = wp.array(
-                [test_data.t],
-                dtype=float,
-                device=device,
-                requires_grad=True,
-            )
-            out = wp.array(
-                [0] * wp.types.type_length(data_type),
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-
-            tape = wp.Tape()
-            with tape:
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[a, b, t, out],
+        with test.subTest(data_type=data_type):
+            for test_data in TEST_DATA[data_type]:
+                a = wp.array(
+                    [test_data.a],
+                    dtype=data_type,
                     device=device,
+                    requires_grad=True,
                 )
-
-            assert_np_equal(
-                out.numpy(),
-                np.array([test_data.expected]),
-                tol=1e-6,
-            )
-
-            if test_data.check_backwards():
-                tape.backward(out)
-
-                assert_np_equal(
-                    tape.gradients[a].numpy(),
-                    np.array([test_data.expected_adj_a]),
-                    tol=1e-6,
+                b = wp.array(
+                    [test_data.b],
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
-                assert_np_equal(
-                    tape.gradients[b].numpy(),
-                    np.array([test_data.expected_adj_b]),
-                    tol=1e-6,
+                t = wp.array(
+                    [test_data.t],
+                    dtype=float,
+                    device=device,
+                    requires_grad=True,
+                )
+                out = wp.array(
+                    [0] * wp.types.type_length(data_type),
+                    dtype=data_type,
+                    device=device,
+                    requires_grad=True,
                 )
+
+                tape = wp.Tape()
+                with tape:
+                    wp.launch(
+                        kernel,
+                        dim=1,
+                        inputs=[a, b, t, out],
+                        device=device,
+                    )
+
                 assert_np_equal(
-                    tape.gradients[t].numpy(),
-                    np.array([test_data.expected_adj_t]),
+                    out.numpy(),
+                    np.array([test_data.expected]),
                     tol=1e-6,
                 )
 
+                if test_data.check_backwards():
+                    tape.backward(out)
+
+                    assert_np_equal(
+                        tape.gradients[a].numpy(),
+                        np.array([test_data.expected_adj_a]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[b].numpy(),
+                        np.array([test_data.expected_adj_b]),
+                        tol=1e-6,
+                    )
+                    assert_np_equal(
+                        tape.gradients[t].numpy(),
+                        np.array([test_data.expected_adj_t]),
+                        tol=1e-6,
+                    )
+
+
+devices = get_test_devices()
+
 
-def register(parent):
-    devices = get_test_devices()
+class TestLerp(unittest.TestCase):
+    pass
 
-    class TestLerp(parent):
-        pass
 
-    add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
-    return TestLerp
+add_function_test(TestLerp, "test_lerp", test_lerp, devices=devices)
 
 
 if __name__ == "__main__":
-    _ = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)