warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_rounding.py

@@ -0,0 +1,185 @@
+# 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
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+compare_to_numpy = False
+print_results = False
+
+
+@wp.kernel
+def test_kernel(
+    x: wp.array(dtype=float),
+    x_round: wp.array(dtype=float),
+    x_rint: wp.array(dtype=float),
+    x_trunc: wp.array(dtype=float),
+    x_cast: wp.array(dtype=float),
+    x_floor: wp.array(dtype=float),
+    x_ceil: wp.array(dtype=float),
+    x_frac: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    x_round[tid] = wp.round(x[tid])
+    x_rint[tid] = wp.rint(x[tid])
+    x_trunc[tid] = wp.trunc(x[tid])
+    x_cast[tid] = float(int(x[tid]))
+    x_floor[tid] = wp.floor(x[tid])
+    x_ceil[tid] = wp.ceil(x[tid])
+    x_frac[tid] = wp.frac(x[tid])
+
+
+def test_rounding(test, device):
+    nx = np.array(
+        [
+            4.9,
+            4.5,
+            4.1,
+            3.9,
+            3.5,
+            3.1,
+            2.9,
+            2.5,
+            2.1,
+            1.9,
+            1.5,
+            1.1,
+            0.9,
+            0.5,
+            0.1,
+            -0.1,
+            -0.5,
+            -0.9,
+            -1.1,
+            -1.5,
+            -1.9,
+            -2.1,
+            -2.5,
+            -2.9,
+            -3.1,
+            -3.5,
+            -3.9,
+            -4.1,
+            -4.5,
+            -4.9,
+        ],
+        dtype=np.float32,
+    )
+
+    x = wp.array(nx, device=device)
+    N = len(x)
+
+    x_round = wp.empty(N, dtype=float, device=device)
+    x_rint = wp.empty(N, dtype=float, device=device)
+    x_trunc = wp.empty(N, dtype=float, device=device)
+    x_cast = wp.empty(N, dtype=float, device=device)
+    x_floor = wp.empty(N, dtype=float, device=device)
+    x_ceil = wp.empty(N, dtype=float, device=device)
+    x_frac = wp.empty(N, dtype=float, device=device)
+
+    wp.launch(
+        kernel=test_kernel, dim=N, inputs=[x, x_round, x_rint, x_trunc, x_cast, x_floor, x_ceil, x_frac], device=device
+    )
+
+    wp.synchronize()
+
+    nx_round = x_round.numpy().reshape(N)
+    nx_rint = x_rint.numpy().reshape(N)
+    nx_trunc = x_trunc.numpy().reshape(N)
+    nx_cast = x_cast.numpy().reshape(N)
+    nx_floor = x_floor.numpy().reshape(N)
+    nx_ceil = x_ceil.numpy().reshape(N)
+    nx_frac = x_frac.numpy().reshape(N)
+
+    tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_cast, nx_floor, nx_ceil, nx_frac], axis=1)
+
+    golden = np.array(
+        [
+            [4.9, 5.0, 5.0, 4.0, 4.0, 4.0, 5.0, 0.9],
+            [4.5, 5.0, 4.0, 4.0, 4.0, 4.0, 5.0, 0.5],
+            [4.1, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 0.1],
+            [3.9, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0, 0.9],
+            [3.5, 4.0, 4.0, 3.0, 3.0, 3.0, 4.0, 0.5],
+            [3.1, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 0.1],
+            [2.9, 3.0, 3.0, 2.0, 2.0, 2.0, 3.0, 0.9],
+            [2.5, 3.0, 2.0, 2.0, 2.0, 2.0, 3.0, 0.5],
+            [2.1, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0, 0.1],
+            [1.9, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0, 0.9],
+            [1.5, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0, 0.5],
+            [1.1, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 0.1],
+            [0.9, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.9],
+            [0.5, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.5],
+            [0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.1],
+            [-0.1, -0.0, -0.0, -0.0, 0.0, -1.0, -0.0, -0.1],
+            [-0.5, -1.0, -0.0, -0.0, 0.0, -1.0, -0.0, -0.5],
+            [-0.9, -1.0, -1.0, -0.0, 0.0, -1.0, -0.0, -0.9],
+            [-1.1, -1.0, -1.0, -1.0, -1.0, -2.0, -1.0, -0.1],
+            [-1.5, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0, -0.5],
+            [-1.9, -2.0, -2.0, -1.0, -1.0, -2.0, -1.0, -0.9],
+            [-2.1, -2.0, -2.0, -2.0, -2.0, -3.0, -2.0, -0.1],
+            [-2.5, -3.0, -2.0, -2.0, -2.0, -3.0, -2.0, -0.5],
+            [-2.9, -3.0, -3.0, -2.0, -2.0, -3.0, -2.0, -0.9],
+            [-3.1, -3.0, -3.0, -3.0, -3.0, -4.0, -3.0, -0.1],
+            [-3.5, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0, -0.5],
+            [-3.9, -4.0, -4.0, -3.0, -3.0, -4.0, -3.0, -0.9],
+            [-4.1, -4.0, -4.0, -4.0, -4.0, -5.0, -4.0, -0.1],
+            [-4.5, -5.0, -4.0, -4.0, -4.0, -5.0, -4.0, -0.5],
+            [-4.9, -5.0, -5.0, -4.0, -4.0, -5.0, -4.0, -0.9],
+        ],
+        dtype=np.float32,
+    )
+
+    assert_np_equal(tab, golden, tol=1e-6)
+
+    if print_results:
+        np.set_printoptions(formatter={"float": lambda x: "{:6.1f}".format(x).replace(".0", ".")})
+
+        print("----------------------------------------------")
+        print("   %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "cast", "floor", "ceil"))
+        print(tab)
+        print("----------------------------------------------")
+
+    if compare_to_numpy:
+        nx_round = np.round(nx)
+        nx_rint = np.rint(nx)
+        nx_trunc = np.trunc(nx)
+        nx_fix = np.fix(nx)
+        nx_floor = np.floor(nx)
+        nx_ceil = np.ceil(nx)
+        nx_frac = np.modf(nx)[0]
+
+        tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_fix, nx_floor, nx_ceil, nx_frac], axis=1)
+        print("   %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "fix", "floor", "ceil"))
+        print(tab)
+        print("----------------------------------------------")
+
+
+class TestRounding(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestRounding, "test_rounding", test_rounding, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_runlength_encode.py

@@ -0,0 +1,196 @@
+# 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
+from functools import partial
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+from warp.utils import runlength_encode
+
+
+def test_runlength_encode_int(test, device, n):
+    rng = np.random.default_rng(123)
+
+    values_np = np.sort(rng.integers(-10, high=10, size=n, dtype=int))
+
+    unique_values_np, unique_counts_np = np.unique(values_np, return_counts=True)
+
+    values = wp.array(values_np, device=device, dtype=int)
+
+    unique_values = wp.empty_like(values)
+    unique_counts = wp.empty_like(values)
+
+    run_count = runlength_encode(values, unique_values, unique_counts)
+
+    test.assertEqual(run_count, len(unique_values_np))
+    assert_np_equal(unique_values.numpy()[:run_count], unique_values_np[:run_count])
+    assert_np_equal(unique_counts.numpy()[:run_count], unique_counts_np[:run_count])
+
+
+def test_runlength_encode_error_insufficient_storage(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(1, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=int, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Output array storage sizes must be at least equal to value_count$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+    values = wp.zeros(123, dtype=int, device="cpu")
+    run_values = wp.empty(123, dtype=int, device="cpu")
+    run_lengths = wp.empty(1, dtype=int, device="cpu")
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Output array storage sizes must be at least equal to value_count$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_dtypes_mismatch(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=float, device=device)
+    run_lengths = wp.empty_like(values, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"values and run_values data types do not match$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_run_length_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty(123, dtype=int, device=device)
+    run_lengths = wp.empty(123, dtype=float, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"run_lengths array must be of type int32$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+def test_runlength_encode_error_run_count_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=int, device=device)
+    run_values = wp.empty_like(values, device=device)
+    run_lengths = wp.empty_like(values, device=device)
+    run_count = wp.empty(shape=(1,), dtype=float, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"run_count array must be of type int32$",
+    ):
+        runlength_encode(values, run_values, run_lengths, run_count=run_count)
+
+
+def test_runlength_encode_error_unsupported_dtype(test, device):
+    values = wp.zeros(123, dtype=float, device=device)
+    run_values = wp.empty(123, dtype=float, device=device)
+    run_lengths = wp.empty(123, dtype=int, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"Unsupported data type$",
+    ):
+        runlength_encode(values, run_values, run_lengths)
+
+
+devices = get_test_devices()
+
+
+class TestRunlengthEncode(unittest.TestCase):
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_devices_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cuda:0")
+        run_lengths = wp.empty_like(values, device="cpu")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage devices do not match$",
+        ):
+            runlength_encode(values, run_values, run_lengths)
+
+    @unittest.skipUnless(wp.is_cuda_available(), "Requires CUDA")
+    def test_runlength_encode_error_run_count_device_mismatch(self):
+        values = wp.zeros(123, dtype=int, device="cpu")
+        run_values = wp.empty_like(values, device="cpu")
+        run_lengths = wp.empty_like(values, device="cpu")
+        run_count = wp.empty(shape=(1,), dtype=int, device="cuda:0")
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"run_count storage device does not match other arrays$",
+        ):
+            runlength_encode(values, run_values, run_lengths, run_count=run_count)
+
+
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_int", partial(test_runlength_encode_int, n=100), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode, "test_runlength_encode_empty", partial(test_runlength_encode_int, n=0), devices=devices
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_insufficient_storage",
+    test_runlength_encode_error_insufficient_storage,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_dtypes_mismatch",
+    test_runlength_encode_error_dtypes_mismatch,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_length_unsupported_dtype",
+    test_runlength_encode_error_run_length_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_run_count_unsupported_dtype",
+    test_runlength_encode_error_run_count_unsupported_dtype,
+    devices=devices,
+)
+add_function_test(
+    TestRunlengthEncode,
+    "test_runlength_encode_error_unsupported_dtype",
+    test_runlength_encode_error_unsupported_dtype,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_scalar_ops.py

@@ -0,0 +1,105 @@
+# 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.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    def make_scalar(value):
+        if wptype in wp.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.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    def make_scalar(value):
+        if wptype in wp.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return wptype._type_(value).value
+
+        return value
+
+    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)
+
+
+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)
+
+
+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.types.type_length(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)