warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/test_scalar_ops.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 *
+
+np_signed_int_types = [
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.byte,
+]
+
+np_unsigned_int_types = [
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+    np.ubyte,
+]
+
+np_int_types = np_signed_int_types + np_unsigned_int_types
+
+np_float_types = [np.float16, np.float32, np.float64]
+
+np_scalar_types = np_int_types + np_float_types
+
+
+def test_py_arithmetic_ops(test, device, dtype):
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    def make_scalar(value):
+        if wptype in wp._src.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return wptype._type_(value).value
+
+        return value
+
+    a = wptype(1)
+    test.assertAlmostEqual(+a, make_scalar(1))
+    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):
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    a = wptype(1)
+    test.assertAlmostEqual(wp.abs(a), 1)
+
+    if dtype in np_float_types:
+        test.assertAlmostEqual(wp.sin(a), 0.84147098480789650488, places=3)
+        test.assertAlmostEqual(wp.radians(a), 0.01745329251994329577, places=5)
+
+
+def test_py_array_ops(test, device, dtype):
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    tol = {
+        np.float16: 5.0e-3,
+        np.float32: 1.0e-6,
+        np.float64: 1.0e-8,
+    }.get(dtype, 0)
+
+    def make_array(*args):
+        if wptype in wp._src.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return np.array(tuple(wptype._type_(x).value for x in args), dtype=dtype)
+
+        return np.array(args, dtype=dtype)
+
+    arr = wp.array((1, 2, 3), dtype=wptype)
+    s = wptype(5)
+    assert_np_equal((arr + s).numpy(), make_array(6, 7, 8), tol=tol)
+    assert_np_equal((s + arr).numpy(), make_array(6, 7, 8), tol=tol)
+    assert_np_equal((arr - s).numpy(), make_array(-4, -3, -2), tol=tol)
+    assert_np_equal((s - arr).numpy(), make_array(4, 3, 2), tol=tol)
+    assert_np_equal((arr * s).numpy(), make_array(5, 10, 15), tol=tol)
+    assert_np_equal((s * arr).numpy(), make_array(5, 10, 15), tol=tol)
+    assert_np_equal((arr % s).numpy(), make_array(1, 2, 3), tol=tol)
+    assert_np_equal((s % arr).numpy(), make_array(0, 1, 2), tol=tol)
+
+    if dtype in np_float_types:
+        assert_np_equal((arr / s).numpy(), make_array(0.2, 0.4, 0.6), tol=tol)
+        assert_np_equal((s / arr).numpy(), make_array(5.0, 2.5, 1.666666667), tol=tol)
+
+
+devices = get_test_devices()
+
+
+class TestScalarOps(unittest.TestCase):
+    pass
+
+
+for dtype in np_scalar_types:
+    add_function_test(
+        TestScalarOps, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
+    )
+    add_function_test(TestScalarOps, f"test_py_math_ops_{dtype.__name__}", test_py_math_ops, devices=None, dtype=dtype)
+    add_function_test(
+        TestScalarOps, f"test_py_array_ops_{dtype.__name__}", test_py_array_ops, devices=None, dtype=dtype
+    )
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_smoothstep.py

@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 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
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@dataclass
+class TestData:
+    a: Any
+    b: Any
+    t: float
+    expected: Any
+    expected_adj_a: Any = None
+    expected_adj_b: Any = None
+    expected_adj_t: float = None
+
+    def check_backwards(self):
+        return self.expected_adj_a is not None and self.expected_adj_b is not None and self.expected_adj_t is not None
+
+
+TEST_DATA = {
+    wp.float32: (
+        TestData(a=1.0, b=2.0, t=1.5, expected=0.5, expected_adj_a=-0.75, expected_adj_b=-0.75, expected_adj_t=1.5),
+        TestData(
+            a=-1.0,
+            b=2.0,
+            t=-0.25,
+            expected=0.15625,
+            expected_adj_a=-0.28125,
+            expected_adj_b=-0.09375,
+            expected_adj_t=0.375,
+        ),
+        TestData(a=0.0, b=1.0, t=9.9, expected=1.0, expected_adj_a=0.0, expected_adj_b=0.0, expected_adj_t=0.0),
+        TestData(a=0.0, b=1.0, t=-9.9, expected=0.0, expected_adj_a=0.0, expected_adj_b=0.0, expected_adj_t=0.0),
+    ),
+}
+
+
+def test_smoothstep(test, device):
+    def make_kernel_fn(data_type):
+        def fn(
+            a: wp.array(dtype=data_type),
+            b: wp.array(dtype=data_type),
+            t: wp.array(dtype=float),
+            out: wp.array(dtype=data_type),
+        ):
+            out[0] = wp.smoothstep(a[0], b[0], t[0])
+
+        return fn
+
+    for data_type, test_data_set in TEST_DATA.items():
+        kernel_fn = make_kernel_fn(data_type)
+        kernel = wp.Kernel(
+            func=kernel_fn,
+            key=f"test_smoothstep{data_type.__name__}_kernel",
+        )
+
+        for test_data in test_data_set:
+            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._src.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
+
+            with wp.Tape() as tape:
+                wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
+
+            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)
+                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()
+
+
+class TestSmoothstep(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSmoothstep, "test_smoothstep", test_smoothstep, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_snippet.py

@@ -0,0 +1,318 @@
+# 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 *
+
+
+def test_basic(test, device):
+    snippet = """
+    out[tid] = a * x[tid] + y[tid];
+    """
+    adj_snippet = """
+    adj_a += x[tid] * adj_out[tid];
+    adj_x[tid] += a * adj_out[tid];
+    adj_y[tid] += adj_out[tid];
+    """
+
+    @wp.func_native(snippet, adj_snippet)
+    def saxpy(
+        a: wp.float32,
+        x: wp.array(dtype=wp.float32),
+        y: wp.array(dtype=wp.float32),
+        out: wp.array(dtype=wp.float32),
+        tid: int,
+    ):  # fmt: skip
+        ...
+
+    @wp.kernel
+    def saxpy_cu(
+        a: wp.float32, x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32), out: wp.array(dtype=wp.float32)
+    ):
+        tid = wp.tid()
+        saxpy(a, x, y, out, tid)
+
+    @wp.kernel
+    def saxpy_py(
+        a: wp.float32, x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32), out: wp.array(dtype=wp.float32)
+    ):
+        tid = wp.tid()
+        out[tid] = a * x[tid] + y[tid]
+
+    N = 128
+
+    a1 = 2.0
+    x1 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device, requires_grad=True)
+    y1 = wp.zeros_like(x1)
+    out1 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device)
+    adj_out1 = wp.array(np.ones(N, dtype=np.float32), dtype=wp.float32, device=device)
+
+    a2 = 2.0
+    x2 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device, requires_grad=True)
+    y2 = wp.zeros_like(x2)
+    out2 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device)
+    adj_out2 = wp.array(np.ones(N, dtype=np.float32), dtype=wp.float32, device=device)
+
+    tape = wp.Tape()
+
+    with tape:
+        wp.launch(kernel=saxpy_cu, dim=N, inputs=[a1, x1, y1], outputs=[out1], device=device)
+        wp.launch(kernel=saxpy_py, dim=N, inputs=[a2, x2, y2], outputs=[out2], device=device)
+
+    tape.backward(grads={out1: adj_out1, out2: adj_out2})
+
+    # test forward snippet
+    assert_np_equal(out1.numpy(), out2.numpy())
+
+    # test backward snippet
+    assert_np_equal(x1.grad.numpy(), a1 * np.ones(N, dtype=np.float32))
+    assert_np_equal(x1.grad.numpy(), x2.grad.numpy())
+
+    assert_np_equal(y1.grad.numpy(), np.ones(N, dtype=np.float32))
+    assert_np_equal(y1.grad.numpy(), y2.grad.numpy())
+
+
+def test_shared_memory(test, device):
+    snippet = """
+        __shared__ int s[128];
+
+        s[tid] = d[tid];
+        __syncthreads();
+        d[tid] = s[N - tid - 1];
+        """
+
+    @wp.func_native(snippet)
+    def reverse(d: wp.array(dtype=int), N: int, tid: int):
+        """Reverse the array d in place using shared memory."""
+        return
+
+    @wp.kernel
+    def reverse_kernel(d: wp.array(dtype=int), N: int):
+        tid = wp.tid()
+        reverse(d, N, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=int), dtype=int, device=device)
+    y = np.arange(127, -1, -1, dtype=int)
+
+    wp.launch(kernel=reverse_kernel, dim=N, inputs=[x, N], device=device)
+
+    assert_np_equal(x.numpy(), y)
+    assert reverse.__doc__ == "Reverse the array d in place using shared memory."
+
+
+def test_cpu_snippet(test, device):
+    snippet = """
+    int inc = 1;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):  # fmt: skip
+        ...
+
+    @wp.kernel
+    def increment(x: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.int32)):
+        tid = wp.tid()
+        increment_snippet(x, out, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=np.int32), dtype=wp.int32, device=device)
+    out = wp.zeros(N, dtype=wp.int32, device=device)
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
+
+
+def test_custom_replay_grad(test, device):
+    num_threads = 16
+    counter = wp.zeros(1, dtype=wp.int32, device=device)
+    thread_ids = wp.zeros(num_threads, dtype=wp.int32, device=device)
+    inputs = wp.array(np.arange(num_threads, dtype=np.float32), device=device, requires_grad=True)
+    outputs = wp.zeros_like(inputs)
+
+    snippet = """
+        int next_index = atomicAdd(counter, 1);
+        thread_values[tid] = next_index;
+        """
+    replay_snippet = ""
+
+    @wp.func_native(snippet, replay_snippet=replay_snippet)
+    def reversible_increment(counter: wp.array(dtype=int), thread_values: wp.array(dtype=int), tid: int):  # fmt: skip
+        ...
+
+    @wp.kernel
+    def run_atomic_add(
+        input: wp.array(dtype=float),
+        counter: wp.array(dtype=int),
+        thread_values: wp.array(dtype=int),
+        output: wp.array(dtype=float),
+    ):
+        tid = wp.tid()
+        reversible_increment(counter, thread_values, tid)
+        idx = thread_values[tid]
+        output[idx] = input[idx] ** 2.0
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(
+            run_atomic_add, dim=num_threads, inputs=[inputs, counter, thread_ids], outputs=[outputs], device=device
+        )
+
+    tape.backward(grads={outputs: wp.array(np.ones(num_threads, dtype=np.float32), device=device)})
+    assert_np_equal(inputs.grad.numpy(), 2.0 * inputs.numpy(), tol=1e-4)
+
+
+def test_replay_simplification(test, device):
+    num_threads = 8
+    x = wp.array(1.0 + np.arange(num_threads, dtype=np.float32), device=device, requires_grad=True)
+    y = wp.zeros_like(x)
+    z = wp.zeros_like(x)
+
+    snippet = "y[tid] = powf(x[tid], 2.0);"
+    replay_snippet = "y[tid] = x[tid];"
+    adj_snippet = "adj_x[tid] += 2.0 * adj_y[tid];"
+
+    @wp.func_native(snippet, adj_snippet=adj_snippet, replay_snippet=replay_snippet)
+    def square(x: wp.array(dtype=float), y: wp.array(dtype=float), tid: int):  # fmt: skip
+        ...
+
+    @wp.kernel
+    def log_square_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+        tid = wp.tid()
+        square(x, y, tid)
+        z[tid] = wp.log(y[tid])
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(log_square_kernel, dim=num_threads, inputs=[x, y], outputs=[z], device=device)
+
+    tape.backward(grads={z: wp.array(np.ones(num_threads, dtype=np.float32), device=device)})
+    assert_np_equal(x.grad.numpy(), 2.0 / (1.0 + np.arange(num_threads)), tol=1e-6)
+
+
+def test_recompile_snippet(test, device):
+    snippet = """
+    int inc = 1;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):  # fmt: skip
+        ...
+
+    @wp.kernel
+    def increment(x: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.int32)):
+        tid = wp.tid()
+        increment_snippet(x, out, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=np.int32), dtype=wp.int32, device=device)
+    out = wp.zeros(N, dtype=wp.int32, device=device)
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
+
+    snippet = """
+    int inc = 2;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):  # fmt: skip
+        ...
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), 1 + np.arange(1, N + 1, 1, dtype=np.int32))
+
+
+def test_return_type(test, device):
+    snippet = """
+        float sq = x * x;
+        return sq;
+        """
+    adj_snippet = """
+        adj_x += 2 * x * adj_ret;
+        """
+
+    # check python built-in return type compilation
+    @wp.func_native(snippet, adj_snippet)
+    def square(x: float) -> float: ...
+
+    # check warp built-in return type compilation
+    @wp.func_native(snippet, adj_snippet)
+    def square(x: wp.float32) -> wp.float32: ...
+
+    @wp.kernel
+    def square_kernel(i: wp.array(dtype=float), o: wp.array(dtype=float)):
+        tid = wp.tid()
+        x = i[tid]
+        o[tid] = square(x)
+
+    N = 5
+    x = wp.array(np.arange(N, dtype=float), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros_like(x)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel=square_kernel, dim=N, inputs=[x, y], device=device)
+
+    y.grad = wp.ones(N, dtype=float, device=device)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.array([0.0, 1.0, 4.0, 9.0, 16.0]))
+    assert_np_equal(x.grad.numpy(), np.array([0.0, 2.0, 4.0, 6.0, 8.0]))
+
+
+class TestSnippets(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSnippets, "test_basic", test_basic, devices=get_selected_cuda_test_devices())
+add_function_test(TestSnippets, "test_shared_memory", test_shared_memory, devices=get_selected_cuda_test_devices())
+add_function_test(TestSnippets, "test_cpu_snippet", test_cpu_snippet, devices=["cpu"])
+add_function_test(
+    TestSnippets, "test_custom_replay_grad", test_custom_replay_grad, devices=get_selected_cuda_test_devices()
+)
+add_function_test(
+    TestSnippets, "test_replay_simplification", test_replay_simplification, devices=get_selected_cuda_test_devices()
+)
+add_function_test(
+    TestSnippets, "test_recompile_snippet", test_recompile_snippet, devices=get_selected_cuda_test_devices()
+)
+add_function_test(TestSnippets, "test_return_type", test_return_type, devices=get_selected_cuda_test_devices())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)