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

warp/tests/test_print.py

@@ -20,6 +20,30 @@ def test_print_kernel():
     wp.printf("this is an int %d\n", 123)
 
 
+@wp.kernel
+def test_print_numeric_kernel(value: int):
+    # signed ints
+    wp.print(wp.int8(value))
+    wp.print(wp.int16(value))
+    wp.print(wp.int32(value))
+    wp.print(wp.int64(value))
+    # unsigned ints
+    wp.print(wp.uint8(value))
+    wp.print(wp.uint16(value))
+    wp.print(wp.uint32(value))
+    wp.print(wp.uint64(value))
+    # floats
+    wp.print(wp.float16(value))
+    wp.print(wp.float32(value))
+    wp.print(wp.float64(value))
+
+
+@wp.kernel
+def test_print_boolean_kernel(value: wp.bool):
+    wp.print(value)
+    wp.print(not value)
+
+
 def test_print(test, device):
     wp.load_module(device=device)
     capture = StdOutCapture()
@@ -39,12 +63,77 @@ def test_print(test, device):
         )
 
 
+def test_print_numeric(test, device):
+    wp.load_module(device=device)
+
+    capture = StdOutCapture()
+    capture.begin()
+    wp.launch(kernel=test_print_numeric_kernel, dim=1, inputs=[17], device=device)
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(
+            s,
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}"
+            rf"17{os.linesep}",
+        )
+
+    capture = StdOutCapture()
+    capture.begin()
+    wp.launch(kernel=test_print_numeric_kernel, dim=1, inputs=[-1], device=device)
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(
+            s,
+            rf"-1{os.linesep}"
+            rf"-1{os.linesep}"
+            rf"-1{os.linesep}"
+            rf"-1{os.linesep}"
+            rf"255{os.linesep}"
+            rf"65535{os.linesep}"
+            rf"4294967295{os.linesep}"
+            rf"18446744073709551615{os.linesep}"
+            rf"-1{os.linesep}"
+            rf"-1{os.linesep}"
+            rf"-1{os.linesep}",
+        )
+
+
+def test_print_boolean(test, device):
+    wp.load_module(device=device)
+    capture = StdOutCapture()
+    capture.begin()
+    wp.launch(kernel=test_print_boolean_kernel, dim=1, inputs=[True], device=device)
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(s, rf"True{os.linesep}False{os.linesep}")
+
+
 class TestPrint(unittest.TestCase):
     pass
 
 
 devices = get_test_devices()
 add_function_test(TestPrint, "test_print", test_print, devices=devices, check_output=False)
+add_function_test(TestPrint, "test_print_numeric", test_print_numeric, devices=devices, check_output=False)
+add_function_test(TestPrint, "test_print_boolean", test_print_boolean, devices=devices, check_output=False)
 
 
 if __name__ == "__main__":

warp/tests/test_quat.py

@@ -1886,6 +1886,110 @@ def test_quat_identity(test, device, dtype, register_kernels=False):
 ############################################################
 
 
+def test_quat_assign(test, device, dtype, register_kernels=False):
+    np_type = np.dtype(dtype)
+    wp_type = wp.types.np_dtype_to_warp_type[np_type]
+
+    quat = wp.types.quaternion(dtype=wp_type)
+
+    def quattest_read_write_store(x: wp.array(dtype=wp_type), a: wp.array(dtype=quat)):
+        tid = wp.tid()
+
+        t = a[tid]
+        t[0] = x[tid]
+        a[tid] = t
+
+    def quattest_in_register(x: wp.array(dtype=wp_type), a: wp.array(dtype=quat)):
+        tid = wp.tid()
+
+        g = wp_type(0.0)
+        q = a[tid]
+        g = q[0] + wp_type(2.0) * q[1] + wp_type(3.0) * q[2] + wp_type(4.0) * q[3]
+        x[tid] = g
+
+    def quattest_in_register_overwrite(x: wp.array(dtype=quat), a: wp.array(dtype=quat)):
+        tid = wp.tid()
+
+        f = quat()
+        a_quat = a[tid]
+        f = a_quat
+        f[1] = wp_type(3.0)
+
+        x[tid] = f
+
+    def quattest_component(x: wp.array(dtype=quat), y: wp.array(dtype=wp_type)):
+        i = wp.tid()
+
+        a = quat()
+        a.x = wp_type(1.0) * y[i]
+        a.y = wp_type(2.0) * y[i]
+        a.z = wp_type(3.0) * y[i]
+        a.w = wp_type(4.0) * y[i]
+        x[i] = a
+
+    kernel_read_write_store = getkernel(quattest_read_write_store, suffix=dtype.__name__)
+    kernel_in_register = getkernel(quattest_in_register, suffix=dtype.__name__)
+    kernel_in_register_overwrite = getkernel(quattest_in_register_overwrite, suffix=dtype.__name__)
+    kernel_component = getkernel(quattest_component, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
+    x = wp.full(1, value=2.0, dtype=wp_type, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_read_write_store, dim=1, inputs=[x, a], device=device)
+
+    tape.backward(grads={a: wp.ones_like(a, requires_grad=False)})
+
+    assert_np_equal(a.numpy(), np.array([[2.0, 1.0, 1.0, 1.0]], dtype=np_type))
+    assert_np_equal(x.grad.numpy(), np.array([1.0], dtype=np_type))
+
+    tape.reset()
+
+    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
+    x = wp.zeros(1, dtype=wp_type, device=device, requires_grad=True)
+
+    with tape:
+        wp.launch(kernel_in_register, dim=1, inputs=[x, a], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([10.0], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros(1, dtype=quat, requires_grad=True)
+    y = wp.ones(1, dtype=wp_type, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_component, dim=1, inputs=[x, y])
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np_type))
+    assert_np_equal(y.grad.numpy(), np.array([10.0], dtype=np_type))
+
+    x = wp.zeros(1, dtype=quat, device=device, requires_grad=True)
+    a = wp.ones(1, dtype=quat, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0, 1.0]], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0, 1.0]], dtype=np_type))
+
+
+############################################################
+
+
 def test_quat_euler_conversion(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
     N = 3
@@ -2088,6 +2192,13 @@ for dtype in np_float_types:
         devices=devices,
         dtype=dtype,
     )
+    add_function_test_register_kernel(
+        TestQuat,
+        f"test_quat_assign_{dtype.__name__}",
+        test_quat_assign,
+        devices=devices,
+        dtype=dtype,
+    )
     add_function_test(
         TestQuat, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
     )

warp/tests/test_reload.py

@@ -64,6 +64,32 @@ def test_redefine(test, device):
     assert_np_equal(np.arange(0, n, 1) * 2.0, y.numpy())
 
 
+def test_redefine_command(test, device):
+    # Test whether executable modules are retained correctly.
+    # A module can have multiple executable versions in use if something
+    # still holds a reference to them.
+
+    @wp.kernel
+    def k(value: int):
+        wp.expect_eq(value, 17)
+
+    # record a command
+    cmd1 = wp.launch(k, dim=1, inputs=[17], device=device, record_cmd=True)
+    cmd1.launch()
+
+    # redefine the kernel, triggering module reload
+    @wp.kernel
+    def k(value: int):
+        wp.expect_eq(value, 42)
+
+    # re-record the command
+    cmd2 = wp.launch(k, dim=1, inputs=[42], device=device, record_cmd=True)
+    cmd2.launch()
+
+    # previous command should still work
+    cmd1.launch()
+
+
 square_two = """import warp as wp
 
 
@@ -79,6 +105,7 @@ def kern(expect: float):
 
 def run(expect, device):
     wp.launch(kern, dim=1, inputs=[expect], device=device)
+    wp.synchronize_device(device)
 """
 
 square_four = """import warp as wp
@@ -96,6 +123,7 @@ def kern(expect: float):
 
 def run(expect, device):
     wp.launch(kern, dim=1, inputs=[expect], device=device)
+    wp.synchronize_device(device)
 """
 
 
@@ -222,9 +250,11 @@ class TestReload(unittest.TestCase):
 
 
 add_function_test(TestReload, "test_redefine", test_redefine, devices=devices)
+add_function_test(TestReload, "test_redefine_command", test_redefine_command, devices=devices)
 add_function_test(TestReload, "test_reload", test_reload, devices=devices)
 add_function_test(TestReload, "test_reload_class", test_reload_class, devices=devices)
-add_function_test(TestReload, "test_reload_references", test_reload_references, devices=devices)
+# TODO: test_reload_references sometimes has issues running on cuda:1
+add_function_test(TestReload, "test_reload_references", test_reload_references, devices=get_test_devices("basic"))
 add_function_test(
     TestReload, "test_graph_launch_after_module_reload", test_graph_launch_after_module_reload, devices=cuda_devices
 )

warp/tests/test_scalar_ops.py

@@ -50,12 +50,14 @@ def test_py_arithmetic_ops(test, device, dtype):
     test.assertAlmostEqual(-a, make_scalar(-1))
     test.assertAlmostEqual(a + wptype(5), make_scalar(6))
     test.assertAlmostEqual(a - wptype(5), make_scalar(-4))
+    test.assertAlmostEqual(a % wptype(2), make_scalar(1))
 
     a = wptype(2)
     test.assertAlmostEqual(a * wptype(2), make_scalar(4))
     test.assertAlmostEqual(wptype(2) * a, make_scalar(4))
     test.assertAlmostEqual(a / wptype(2), make_scalar(1))
     test.assertAlmostEqual(wptype(24) / a, make_scalar(12))
+    test.assertAlmostEqual(a % wptype(2), make_scalar(0))
 
 
 def test_py_math_ops(test, device, dtype):