warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/tests/test_math.py

@@ -0,0 +1,178 @@
+# 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 typing import Any, NamedTuple
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import add_function_test, assert_np_equal, get_test_devices
+
+
+class ScalarFloatValues(NamedTuple):
+    degrees: wp.float32 = None
+    radians: wp.float32 = None
+
+
+@wp.kernel
+def scalar_float_kernel(
+    i: int,
+    x: wp.array(dtype=wp.float32),
+    out: wp.array(dtype=wp.float32),
+):
+    if i == 0:
+        out[0] = wp.degrees(x[0])
+    elif i == 1:
+        out[0] = wp.radians(x[0])
+
+
+def test_scalar_math(test, device):
+    float_values = ScalarFloatValues(degrees=(0.123,), radians=(123.0,))
+    float_results_expected = ScalarFloatValues(degrees=7.047381, radians=2.146755)
+    adj_float_results_expected = ScalarFloatValues(degrees=57.29578, radians=0.017453)
+    for i, values in enumerate(float_values):
+        x = wp.array([values[0]], dtype=wp.float32, requires_grad=True, device=device)
+        out = wp.array([0.0], dtype=wp.float32, requires_grad=True, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(scalar_float_kernel, dim=1, inputs=[i, x, out], device=device)
+
+        assert_np_equal(out.numpy(), np.array([float_results_expected[i]]), tol=1e-6)
+
+        tape.backward(out)
+
+        assert_np_equal(tape.gradients[x].numpy(), np.array([adj_float_results_expected[i]]), tol=1e-6)
+
+
+@wp.kernel
+def test_vec_norm_kernel(vs: wp.array(dtype=Any), out: wp.array(dtype=float, ndim=2)):
+    tid = wp.tid()
+    out[tid, 0] = wp.norm_l1(vs[tid])
+    out[tid, 1] = wp.norm_l2(vs[tid])
+    out[tid, 2] = wp.norm_huber(vs[tid])
+    out[tid, 3] = wp.norm_pseudo_huber(vs[tid])
+
+
+def test_vec_norm(test, device):
+    # ground-truth implementations from SciPy
+    def huber(delta, x):
+        if x <= delta:
+            return 0.5 * x**2
+        else:
+            return delta * (x - 0.5 * delta)
+
+    def pseudo_huber(delta, x):
+        return delta**2 * (np.sqrt(1 + (x / delta) ** 2) - 1)
+
+    v0 = wp.vec3(-2.0, -1.0, -3.0)
+    v1 = wp.vec3(2.0, 1.0, 3.0)
+    v2 = wp.vec3(0.0, 0.0, 0.0)
+
+    xs = wp.array([v0, v1, v2], dtype=wp.vec3, requires_grad=True, device=device)
+    out = wp.empty((len(xs), 4), dtype=wp.float32, requires_grad=True, device=device)
+
+    wp.launch(test_vec_norm_kernel, dim=len(xs), inputs=[xs], outputs=[out], device=device)
+
+    for i, x in enumerate([v0, v1, v2]):
+        assert_np_equal(
+            out.numpy()[i],
+            np.array(
+                [
+                    np.linalg.norm(x, ord=1),
+                    np.linalg.norm(x, ord=2),
+                    huber(1.0, wp.length(x)),
+                    # note SciPy defines the Pseudo-Huber loss slightly differently
+                    pseudo_huber(1.0, wp.length(x)) + 1.0,
+                ]
+            ),
+            tol=1e-6,
+        )
+
+
+devices = get_test_devices()
+
+
+class TestMath(unittest.TestCase):
+    def test_vec_type(self):
+        vec5 = wp.vec(length=5, dtype=float)
+        v = vec5()
+        w = vec5()
+        a = vec5(1.0)
+        b = vec5(0.0, 0.0, 0.0, 0.0, 0.0)
+        c = vec5(0.0)
+
+        v[0] = 1.0
+        v.x = 0.0
+        v[1:] = [1.0, 1.0, 1.0, 1.0]
+
+        w[0] = 1.0
+        w[1:] = [0.0, 0.0, 0.0, 0.0]
+
+        self.assertEqual(v[0], w[1], "vec setter error")
+        self.assertEqual(v.x, w.y, "vec setter error")
+
+        for x in v[1:]:
+            self.assertEqual(x, 1.0, "vec slicing error")
+
+        self.assertEqual(b, c, "vec equality error")
+
+        self.assertEqual(str(v), "[0.0, 1.0, 1.0, 1.0, 1.0]", "vec to string error")
+
+    def test_mat_type(self):
+        mat55 = wp.mat(shape=(5, 5), dtype=float)
+        m1 = mat55()
+        m2 = mat55()
+
+        for i in range(5):
+            for j in range(5):
+                if i == j:
+                    m1[i, j] = 1.0
+                else:
+                    m1[i, j] = 0.0
+
+        for i in range(5):
+            m2[i] = [1.0, 1.0, 1.0, 1.0, 1.0]
+
+        a = mat55(1.0)
+        # fmt: off
+        b = mat55(
+            1.0, 0.0, 0.0, 0.0, 0.0,
+            0.0, 1.0, 0.0, 0.0, 0.0,
+            0.0, 0.0, 1.0, 0.0, 0.0,
+            0.0, 0.0, 0.0, 1.0, 0.0,
+            0.0, 0.0, 0.0, 0.0, 1.0,
+        )
+        # fmt: on
+
+        self.assertEqual(m1, b, "mat element setting error")
+        self.assertEqual(m2, a, "mat row setting error")
+        self.assertEqual(m1[0, 0], 1.0, "mat element getting error")
+        self.assertEqual(m2[0], [1.0, 1.0, 1.0, 1.0, 1.0], "mat row getting error")
+        self.assertEqual(
+            str(b),
+            "[[1.0, 0.0, 0.0, 0.0, 0.0],\n [0.0, 1.0, 0.0, 0.0, 0.0],\n [0.0, 0.0, 1.0, 0.0, 0.0],\n [0.0, 0.0, 0.0, 1.0, 0.0],\n [0.0, 0.0, 0.0, 0.0, 1.0]]",
+            "mat to string error",
+        )
+
+
+add_function_test(TestMath, "test_scalar_math", test_scalar_math, devices=devices)
+add_function_test(TestMath, "test_vec_norm", test_vec_norm, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_mlp.py

@@ -0,0 +1,282 @@
+# 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 *
+
+
+@wp.func
+def mlp_activation(z: float):
+    return wp.tanh(z)
+
+
+@wp.kernel
+def mlp_kernel(
+    weights: wp.array2d(dtype=float),
+    bias: wp.array(dtype=float),
+    x: wp.array2d(dtype=float),
+    y: wp.array2d(dtype=float),
+):
+    wp.mlp(weights, bias, mlp_activation, wp.tid(), x, y)
+
+
+@wp.kernel
+def loss_kernel(x: wp.array2d(dtype=float), loss: wp.array(dtype=float)):
+    i, j = wp.tid()
+
+    wp.atomic_add(loss, 0, x[i, j] * x[i, j])
+
+
+def test_mlp(test, device):
+    rng = np.random.default_rng(123)
+
+    m = 10
+    n = 200
+
+    batches = 20000
+
+    weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+    bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
+
+    x = wp.array(rng.random(size=(n, batches)), dtype=float, device=device)
+    y = wp.zeros(shape=(m, batches), device=device)
+
+    with wp.ScopedTimer("warp", active=False):
+        wp.launch(mlp_kernel, dim=batches, inputs=[weights, bias, x, y], device=device)
+        wp.synchronize()
+
+    # A*x + b
+    with wp.ScopedTimer("numpy", active=False):
+        expect = np.tanh(weights.numpy().reshape(m, n) @ x.numpy().reshape(-1, batches) + bias.numpy().reshape(m, 1))
+
+    result = y.numpy().reshape(-1, batches)
+
+    assert_np_equal(result, expect, tol=1.0e-6)
+
+
+def create_mlp(m, n):
+    import torch
+
+    torch.manual_seed(0)
+
+    class FeedForward(torch.nn.Module):
+        def __init__(self, input_size, hidden_size):
+            super(FeedForward, self).__init__()
+
+            self.input_size = input_size
+            self.hidden_size = hidden_size
+            self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size)
+            self.act = torch.nn.Tanh()
+
+        def forward(self, x):
+            out = self.fc1(x)
+            out = self.act(out)
+            return out
+
+    return FeedForward(m, n)
+
+
+def create_golden():
+    import torch
+
+    rng = np.random.default_rng(123)
+
+    input_size = 32
+    hidden_size = 16
+    batch_size = 64
+
+    network = create_mlp(input_size, hidden_size)
+
+    x = torch.Tensor(rng.random(size=(batch_size, input_size)))
+    x.requires_grad = True
+
+    y = network.forward(x)
+    y.retain_grad()
+
+    loss = torch.inner(y.flatten(), y.flatten())
+    loss.backward(retain_graph=True)
+
+    results = {}
+    results["weights"] = network.fc1.weight.cpu().detach().numpy()
+    results["weights_grad"] = network.fc1.weight.grad.cpu().detach().numpy()
+    results["bias"] = network.fc1.bias.cpu().detach().numpy()
+    results["bias_grad"] = network.fc1.bias.grad.cpu().detach().numpy()
+    results["x"] = x.cpu().detach().numpy()
+    results["x_grad"] = x.grad.cpu().detach().numpy()
+    results["y"] = y.cpu().detach().numpy()
+    results["y_grad"] = y.grad.cpu().detach().numpy()
+    results["loss"] = loss.cpu().detach().numpy()
+
+    np.save(os.path.join(os.path.dirname(__file__), "assets/mlp_golden.npy"), results, allow_pickle=True)
+
+
+def load_golden():
+    return np.load(os.path.join(os.path.dirname(__file__), "assets/mlp_golden.npy"), allow_pickle=True).item()
+
+
+def test_mlp_grad(test, device):
+    # uncomment to re-build golden files
+    # create_golden()
+
+    results = load_golden()
+
+    torch_weights = results["weights"]
+    torch_weights_grad = results["weights_grad"]
+    torch_bias = results["bias"]
+    torch_bias_grad = results["bias_grad"]
+    torch_x = results["x"].T
+    torch_x_grad = results["x_grad"].T
+    torch_y = results["y"].T
+    torch_y_grad = results["y_grad"].T
+    torch_loss = results["loss"].T
+
+    weights = wp.array(torch_weights, dtype=float, device=device, requires_grad=True)
+    bias = wp.array(torch_bias, dtype=float, device=device, requires_grad=True)
+
+    x = wp.array(torch_x, dtype=float, device=device, requires_grad=True)
+    y = wp.array(torch_y, dtype=float, device=device, requires_grad=True)
+    y.zero_()
+
+    loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
+
+    m = torch_weights.shape[0]
+    n = torch_weights.shape[1]
+    b = torch_x.shape[1]
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(mlp_kernel, dim=b, inputs=[weights, bias, x, y], device=device)
+        wp.launch(loss_kernel, dim=y.shape, inputs=[y, loss], device=device)
+
+    tape.backward(loss=loss)
+
+    # check forward result
+    assert_np_equal(y.numpy().reshape(-1, b), torch_y, tol=1.0e-1)
+    assert_np_equal(loss.numpy(), torch_loss, tol=1.0e-1)
+
+    # check backward result
+    assert_np_equal(tape.gradients[weights].numpy().reshape(m, n), torch_weights_grad, tol=1.0e-1)
+    assert_np_equal(tape.gradients[bias].numpy(), torch_bias_grad, tol=1.0e-1)
+    assert_np_equal(tape.gradients[x].numpy().reshape(n, b), torch_x_grad, tol=1.0e-1)
+    assert_np_equal(tape.gradients[y].numpy().reshape(m, b), torch_y_grad, tol=1.0e-1)
+
+
+def profile_mlp_torch():
+    import torch
+
+    rng = np.random.default_rng(123)
+
+    m = 128
+    n = 64
+
+    steps = 20
+
+    for i in range(steps):
+        b = 2**i
+
+        network = create_mlp(m, n)
+
+        x = torch.Tensor(rng.random(size=(b, m)))
+
+        with wp.ScopedTimer("torch_forward" + str(b)):
+            y = network.forward(x)
+            torch.cuda.synchronize()
+
+    for i in range(steps):
+        b = 2**i
+
+        network = create_mlp(m, n)
+
+        x = torch.Tensor(rng.random(size=(b, m)))
+        y = network.forward(x)
+
+        loss = torch.norm(y)
+
+        # run once to alloc all gradients
+        loss.backward(retain_graph=True)
+
+        with wp.ScopedTimer("torch-backward" + str(b)):
+            loss.backward()
+            torch.cuda.synchronize()
+
+
+def profile_mlp_warp(device):
+    rng = np.random.default_rng(123)
+
+    m = 128
+    n = 64
+
+    steps = 20
+
+    for i in range(steps):
+        b = 2**i
+
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device)
+
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device)
+        y = wp.zeros(shape=(m, b), device=device)
+
+        with wp.ScopedTimer("warp-forward" + str(b)):
+            wp.launch(mlp_kernel, dim=b, inputs=[weights, bias, x, y], device=device)
+            wp.synchronize()
+
+    for i in range(steps):
+        b = 2**i
+
+        weights = wp.array(rng.random(size=(m, n)) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+        bias = wp.array(rng.random(size=m) * 0.5 - 0.5, dtype=float, device=device, requires_grad=True)
+
+        x = wp.array(rng.random(size=(n, b)), dtype=float, device=device, requires_grad=True)
+        y = wp.zeros(shape=(m, b), device=device, requires_grad=True)
+
+        loss = wp.zeros(1, dtype=float, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(mlp_kernel, dim=b, inputs=[weights, bias, x, y], device=device)
+            wp.launch(loss_kernel, dim=y.size, inputs=[y.flatten(), loss], device=device)
+
+        # run backward once to ensure all adjoints are allocated
+        tape.backward(loss)
+        wp.synchronize()
+
+        with wp.ScopedTimer("warp-backward" + str(b)):
+            tape.backward(loss)
+            wp.synchronize()
+
+
+# profile_mlp_warp("cuda")
+# profile_mlp_torch()
+
+
+devices = get_test_devices()
+
+
+class TestMLP(unittest.TestCase):
+    pass
+
+
+add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices, check_output=False)
+add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices, check_output=False)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)

warp/tests/test_module_hashing.py

@@ -0,0 +1,258 @@
+# 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.
+
+# TODO: add more tests for kernels and generics
+
+import os
+import tempfile
+import unittest
+from importlib import util
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+FUNC_OVERLOAD_1 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.func
+def fn():
+    wp.print(17)
+
+@wp.func
+def fn(value: int):
+    wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
+"""
+
+# should be same hash as FUNC_OVERLOAD_1
+FUNC_OVERLOAD_2 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.func
+def fn():
+    wp.print(17)
+
+@wp.func
+def fn(value: int):
+    wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
+"""
+
+# should be different hash than FUNC_OVERLOAD_1 (first overload is different)
+FUNC_OVERLOAD_3 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.func
+def fn():
+    wp.print(42)
+
+@wp.func
+def fn(value: int):
+    wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
+"""
+
+# should be different hash than FUNC_OVERLOAD_1 (second overload is different)
+FUNC_OVERLOAD_4 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.func
+def fn():
+    wp.print(17)
+
+@wp.func
+def fn(value: int):
+    wp.print(value + 1)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
+"""
+
+FUNC_GENERIC_1 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+from typing import Any
+
+@wp.func
+def generic_fn(x: Any):
+    return x * x
+
+@wp.func
+def generic_fn(x: Any, y: Any):
+    return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
+"""
+
+# should be same hash as FUNC_GENERIC_1
+FUNC_GENERIC_2 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+from typing import Any
+
+@wp.func
+def generic_fn(x: Any):
+    return x * x
+
+@wp.func
+def generic_fn(x: Any, y: Any):
+    return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
+"""
+
+# should be different hash than FUNC_GENERIC_1 (first overload is different)
+FUNC_GENERIC_3 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+from typing import Any
+
+@wp.func
+def generic_fn(x: Any):
+    return x + x
+
+@wp.func
+def generic_fn(x: Any, y: Any):
+    return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
+"""
+
+# should be different hash than FUNC_GENERIC_1 (second overload is different)
+FUNC_GENERIC_4 = """# -*- coding: utf-8 -*-
+import warp as wp
+
+from typing import Any
+
+@wp.func
+def generic_fn(x: Any):
+    return x * x
+
+@wp.func
+def generic_fn(x: Any, y: Any):
+    return x + y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
+"""
+
+
+def load_code_as_module(code, name):
+    file, file_path = tempfile.mkstemp(suffix=".py")
+
+    try:
+        with os.fdopen(file, "w") as f:
+            f.write(code)
+
+        spec = util.spec_from_file_location(name, file_path)
+        module = util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    finally:
+        os.remove(file_path)
+
+    return wp.get_module(module.__name__)
+
+
+def test_function_overload_hashing(test, device):
+    m1 = load_code_as_module(FUNC_OVERLOAD_1, "func_overload_1")
+    m2 = load_code_as_module(FUNC_OVERLOAD_2, "func_overload_2")
+    m3 = load_code_as_module(FUNC_OVERLOAD_3, "func_overload_3")
+    m4 = load_code_as_module(FUNC_OVERLOAD_4, "func_overload_4")
+
+    hash1 = m1.hash_module()
+    hash2 = m2.hash_module()
+    hash3 = m3.hash_module()
+    hash4 = m4.hash_module()
+
+    test.assertEqual(hash2, hash1)
+    test.assertNotEqual(hash3, hash1)
+    test.assertNotEqual(hash4, hash1)
+
+
+def test_function_generic_overload_hashing(test, device):
+    m1 = load_code_as_module(FUNC_GENERIC_1, "func_generic_1")
+    m2 = load_code_as_module(FUNC_GENERIC_2, "func_generic_2")
+    m3 = load_code_as_module(FUNC_GENERIC_3, "func_generic_3")
+    m4 = load_code_as_module(FUNC_GENERIC_4, "func_generic_4")
+
+    hash1 = m1.hash_module()
+    hash2 = m2.hash_module()
+    hash3 = m3.hash_module()
+    hash4 = m4.hash_module()
+
+    test.assertEqual(hash2, hash1)
+    test.assertNotEqual(hash3, hash1)
+    test.assertNotEqual(hash4, hash1)
+
+
+SIMPLE_MODULE = """# -*- coding: utf-8 -*-
+import warp as wp
+
+@wp.kernel
+def k():
+    pass
+"""
+
+
+def test_module_load(test, device):
+    """Ensure that loading a module does not change its hash"""
+    m = load_code_as_module(SIMPLE_MODULE, "simple_module")
+
+    hash1 = m.hash_module()
+    m.load(device)
+    hash2 = m.hash_module()
+
+    test.assertEqual(hash1, hash2)
+
+
+class TestModuleHashing(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestModuleHashing, "test_function_overload_hashing", test_function_overload_hashing)
+add_function_test(TestModuleHashing, "test_function_generic_overload_hashing", test_function_generic_overload_hashing)
+add_function_test(TestModuleHashing, "test_module_load", test_module_load, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)