warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/test_fixedarray.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 *
+
+devices = get_test_devices()
+
+
+@wp.kernel
+def test_zeros():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_arg_func(arr: wp.array(ndim=2, dtype=int)):
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+
+@wp.kernel
+def test_func_arg():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    test_func_arg_func(arr)
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_return_func():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    return arr
+
+
+@wp.kernel
+def test_func_return():
+    arr = test_func_return_func()
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_return_annotation_func() -> wp.fixedarray(shape=(2, 3), dtype=int):
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    return arr
+
+
+@wp.kernel
+def test_func_return_annotation():
+    arr = test_func_return_annotation_func()
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+def test_error_invalid_func_return_annotation(test, device):
+    @wp.func
+    def func() -> wp.array(ndim=2, dtype=int):
+        arr = wp.zeros(shape=(2, 3), dtype=int)
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                arr[i][j] = i * arr.shape[1] + j
+
+        return arr
+
+    @wp.kernel
+    def kernel():
+        arr = func()
+
+    with test.assertRaisesRegex(
+        wp._src.codegen.WarpCodegenError,
+        r"The function `func` returns a fixed-size array whereas it has its return type annotated as `Array\[int32\]`.$",
+    ):
+        wp.launch(kernel, 1, device=device)
+
+
+def test_error_runtime_shape(test, device):
+    @wp.kernel
+    def kernel():
+        tid = wp.tid()
+        wp.zeros(shape=(tid,), dtype=int)
+
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"the `shape` argument must be specified as a constant when zero-initializing an array$",
+    ):
+        wp.launch(kernel, 1, device=device)
+
+
+@wp.kernel
+def test_capture_if_kernel():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+def test_capture_if(test, device):
+    if (
+        not wp.get_device(device).is_cuda
+        or wp._src.context.runtime.toolkit_version < (12, 4)
+        or wp._src.context.runtime.driver_version < (12, 4)
+    ):
+        return
+
+    def foo():
+        wp.launch(test_capture_if_kernel, dim=512, block_dim=128, device=device)
+
+    cond = wp.ones(1, dtype=wp.int32, device=device)
+    with wp.ScopedCapture(device=device) as capture:
+        wp.capture_if(condition=cond, on_true=foo)
+
+    wp.capture_launch(capture.graph)
+
+
+@wp.struct
+class test_func_struct_MyStruct:
+    offset: int
+    dist: float
+
+
+@wp.func
+def test_func_struct_func():
+    arr = wp.zeros(shape=(2, 3), dtype=test_func_struct_MyStruct)
+    count = float(arr.shape[0] * arr.shape[1] - 1)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j].offset = i * arr.shape[1] + j
+            arr[i][j].dist = float(arr[i][j].offset) / count
+
+    return arr
+
+
+@wp.kernel
+def test_func_struct():
+    arr = test_func_struct_func()
+
+    wp.expect_eq(arr[0][0].offset, 0)
+    wp.expect_near(arr[0][0].dist, 0.0)
+    wp.expect_eq(arr[0][1].offset, 1)
+    wp.expect_near(arr[0][1].dist, 0.2)
+    wp.expect_eq(arr[0][2].offset, 2)
+    wp.expect_near(arr[0][2].dist, 0.4)
+    wp.expect_eq(arr[1][0].offset, 3)
+    wp.expect_near(arr[1][0].dist, 0.6)
+    wp.expect_eq(arr[1][1].offset, 4)
+    wp.expect_near(arr[1][1].dist, 0.8)
+    wp.expect_eq(arr[1][2].offset, 5)
+    wp.expect_near(arr[1][2].dist, 1.0)
+
+
+class TestFixedArray(unittest.TestCase):
+    pass
+
+
+add_kernel_test(TestFixedArray, kernel=test_zeros, name="test_zeros", dim=1, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_arg, name="test_func_arg", dim=1, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_return, name="test_func_return", dim=1, devices=devices)
+add_kernel_test(
+    TestFixedArray, kernel=test_func_return_annotation, name="test_func_return_annotation", dim=1, devices=devices
+)
+add_function_test(
+    TestFixedArray,
+    "test_error_invalid_func_return_annotation",
+    test_error_invalid_func_return_annotation,
+    devices=devices,
+)
+add_function_test(TestFixedArray, "test_error_runtime_shape", test_error_runtime_shape, devices=devices)
+add_function_test(TestFixedArray, "test_capture_if", test_capture_if, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_struct, name="test_func_struct", dim=1, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_fp16.py

@@ -0,0 +1,136 @@
+# 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.tests.unittest_utils import *
+
+
+@wp.kernel
+def load_store_half(f32: wp.array(dtype=wp.float32), f16: wp.array(dtype=wp.float16)):
+    tid = wp.tid()
+
+    # check conversion from f32->f16
+    a = wp.float16(f32[tid])
+    b = f16[tid]
+
+    wp.expect_eq(a, b)
+
+    # check stores
+    f16[tid] = a
+
+
+def test_fp16_conversion(test, device):
+    s = [1.0, 2.0, 3.0, -3.14159]
+
+    np_f32 = np.array(s, dtype=np.float32)
+    np_f16 = np.array(s, dtype=np.float16)
+
+    wp_f32 = wp.array(s, dtype=wp.float32, device=device)
+    wp_f16 = wp.array(s, dtype=wp.float16, device=device)
+
+    assert_np_equal(np_f32, wp_f32.numpy())
+    assert_np_equal(np_f16, wp_f16.numpy())
+
+    wp.launch(load_store_half, dim=len(s), inputs=[wp_f32, wp_f16], device=device)
+
+    # check that stores worked
+    assert_np_equal(np_f16, wp_f16.numpy())
+
+
+@wp.kernel
+def value_load_store_half(f16_value: wp.float16, f16_array: wp.array(dtype=wp.float16)):
+    wp.expect_eq(f16_value, f16_array[0])
+
+    # check stores
+    f16_array[0] = f16_value
+
+
+def test_fp16_kernel_parameter(test, device):
+    """Test the ability to pass in fp16 into kernels as parameters"""
+
+    s = [1.0, 2.0, 3.0, -3.14159]
+
+    for test_val in s:
+        np_f16 = np.array([test_val], dtype=np.float16)
+        wp_f16 = wp.array([test_val], dtype=wp.float16, device=device)
+
+        wp.launch(value_load_store_half, (1,), inputs=[wp.float16(test_val), wp_f16], device=device)
+
+        # check that stores worked
+        assert_np_equal(np_f16, wp_f16.numpy())
+
+        # Do the same thing but pass in test_val as a Python float to test automatic conversion
+        wp_f16 = wp.array([test_val], dtype=wp.float16, device=device)
+        wp.launch(value_load_store_half, (1,), inputs=[test_val, wp_f16], device=device)
+        assert_np_equal(np_f16, wp_f16.numpy())
+
+
+@wp.kernel
+def mul_half(input: wp.array(dtype=wp.float16), output: wp.array(dtype=wp.float16)):
+    tid = wp.tid()
+
+    # convert to compute type fp32
+    x = wp.float(input[tid]) * 2.0
+
+    # store back as fp16
+    output[tid] = wp.float16(x)
+
+
+def test_fp16_grad(test, device):
+    rng = np.random.default_rng(123)
+
+    # checks that gradients are correctly propagated for
+    # fp16 arrays, even when intermediate calculations
+    # are performed in e.g.: fp32
+
+    s = rng.random(size=15).astype(np.float16)
+
+    input = wp.array(s, dtype=wp.float16, device=device, requires_grad=True)
+    output = wp.zeros_like(input)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(mul_half, dim=len(s), inputs=[input, output], device=device)
+
+    ones = wp.array(np.ones(len(output)), dtype=wp.float16, device=device)
+
+    tape.backward(grads={output: ones})
+
+    assert_np_equal(input.grad.numpy(), np.ones(len(s)) * 2.0)
+
+
+class TestFp16(unittest.TestCase):
+    pass
+
+
+devices = []
+if wp.is_cpu_available():
+    devices.append("cpu")
+for cuda_device in get_selected_cuda_test_devices():
+    if cuda_device.arch >= 70:
+        devices.append(cuda_device)
+
+add_function_test(TestFp16, "test_fp16_conversion", test_fp16_conversion, devices=devices)
+add_function_test(TestFp16, "test_fp16_grad", test_fp16_grad, devices=devices)
+add_function_test(TestFp16, "test_fp16_kernel_parameter", test_fp16_kernel_parameter, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)