warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/test_tape.py

@@ -0,0 +1,242 @@
+# 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.kernel
+def mul_constant(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    y[tid] = x[tid] * 2.0
+
+
+@wp.struct
+class Multiplicands:
+    x: wp.array(dtype=float)
+    y: wp.array(dtype=float)
+
+
+@wp.kernel
+def mul_variable(mutiplicands: Multiplicands, z: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    z[tid] = mutiplicands.x[tid] * mutiplicands.y[tid]
+
+
+@wp.kernel
+def dot_product(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+    tid = wp.tid()
+
+    wp.atomic_add(z, 0, x[tid] * y[tid])
+
+
+def test_tape_mul_constant(test, device):
+    dim = 8
+    iters = 16
+    tape = wp.Tape()
+
+    # record onto tape
+    with tape:
+        # input data
+        x0 = wp.array(np.zeros(dim), dtype=wp.float32, device=device, requires_grad=True)
+        x = x0
+
+        for _i in range(iters):
+            y = wp.empty_like(x, requires_grad=True)
+            wp.launch(kernel=mul_constant, dim=dim, inputs=[x], outputs=[y], device=device)
+            x = y
+
+    # loss = wp.sum(x)
+    x.grad = wp.array(np.ones(dim), device=device, dtype=wp.float32)
+
+    # run backward
+    tape.backward()
+
+    # grad = 2.0^iters
+    assert_np_equal(tape.gradients[x0].numpy(), np.ones(dim) * (2**iters))
+
+
+def test_tape_mul_variable(test, device):
+    dim = 8
+    tape = wp.Tape()
+
+    # record onto tape
+    with tape:
+        # input data (Note: We're intentionally testing structs in tapes here)
+        multiplicands = Multiplicands()
+        multiplicands.x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
+        multiplicands.y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
+        z = wp.zeros_like(multiplicands.x)
+
+        wp.launch(kernel=mul_variable, dim=dim, inputs=[multiplicands], outputs=[z], device=device)
+
+    # loss = wp.sum(x)
+    z.grad = wp.array(np.ones(dim), device=device, dtype=wp.float32)
+
+    # run backward
+    tape.backward()
+
+    # grad_x=y, grad_y=x
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy())
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy())
+
+    # run backward again with different incoming gradient
+    # should accumulate the same gradients again onto output
+    # so gradients = 2.0*prev
+    tape.backward()
+
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy() * 2.0)
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy() * 2.0)
+
+    # Clear launches and zero out the gradients
+    tape.reset()
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), np.zeros_like(tape.gradients[multiplicands].x.numpy()))
+    test.assertFalse(tape.launches)
+
+
+def test_tape_dot_product(test, device):
+    dim = 8
+    tape = wp.Tape()
+
+    # record onto tape
+    with tape:
+        # input data
+        x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
+        y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
+        z = wp.zeros(n=1, dtype=wp.float32, device=device, requires_grad=True)
+
+        wp.launch(kernel=dot_product, dim=dim, inputs=[x, y], outputs=[z], device=device)
+
+    # scalar loss
+    tape.backward(loss=z)
+
+    # grad_x=y, grad_y=x
+    assert_np_equal(tape.gradients[x].numpy(), y.numpy())
+    assert_np_equal(tape.gradients[y].numpy(), x.numpy())
+
+
+@wp.kernel
+def assign_chain_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+    tid = wp.tid()
+    y[tid] = x[tid]
+    z[tid] = y[tid]
+
+
+def test_tape_zero_multiple_outputs(test, device):
+    x = wp.array(np.arange(3), dtype=float, device=device, requires_grad=True)
+    y = wp.zeros_like(x)
+    z = wp.zeros_like(x)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(assign_chain_kernel, dim=3, inputs=[x, y, z], device=device)
+
+    tape.backward(grads={y: wp.ones_like(x)})
+    assert_np_equal(x.grad.numpy(), np.ones(3, dtype=float))
+    tape.zero()
+
+    tape.backward(grads={z: wp.ones_like(x)})
+    assert_np_equal(x.grad.numpy(), np.ones(3, dtype=float))
+
+
+@wp.struct
+class NestedStruct:
+    arr: wp.array(dtype=float)
+
+
+@wp.struct
+class WrapperStruct:
+    nested: NestedStruct
+
+
+@wp.kernel
+def nested_loss_kernel(wrapper: WrapperStruct, loss: wp.array(dtype=float)):
+    i = wp.tid()
+    wp.atomic_add(loss, 0, wrapper.nested.arr[i])
+
+
+def test_tape_nested_struct(test, device):
+    wrapper = WrapperStruct()
+    wrapper.nested = NestedStruct()
+    wrapper.nested.arr = wp.ones(shape=(1,), dtype=float, requires_grad=True, device=device)
+
+    loss = wp.zeros(shape=(1,), dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(nested_loss_kernel, dim=1, inputs=(wrapper, loss), device=device)
+
+    assert_np_equal(loss.numpy(), [1.0])
+
+    tape.backward(loss)
+    assert_np_equal(wrapper.nested.arr.grad.numpy(), [1.0])
+
+    tape.zero()
+
+    assert_np_equal(wrapper.nested.arr.grad.numpy(), [0.0])
+
+
+def test_tape_visualize(test, device):
+    dim = 8
+    tape = wp.Tape()
+
+    # record onto tape
+    with tape:
+        # input data
+        x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
+        y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
+        z = wp.zeros(n=1, dtype=wp.float32, device=device, requires_grad=True)
+
+        tape.record_scope_begin("my loop")
+        for _ in range(16):
+            wp.launch(kernel=dot_product, dim=dim, inputs=[x, y], outputs=[z], device=device)
+        tape.record_scope_end()
+
+    # generate GraphViz diagram code
+    dot_code = tape.visualize(simplify_graph=True)
+
+    assert "repeated 16x" in dot_code
+    assert "my loop" in dot_code
+    assert dot_code.count("dot_product") == 1
+
+
+devices = get_test_devices()
+
+
+class TestTape(unittest.TestCase):
+    def test_tape_no_nested_tapes(self):
+        with self.assertRaises(RuntimeError):
+            with wp.Tape():
+                with wp.Tape():
+                    pass
+
+
+add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
+add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
+add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
+add_function_test(TestTape, "test_tape_zero_multiple_outputs", test_tape_zero_multiple_outputs, devices=devices)
+add_function_test(TestTape, "test_tape_nested_struct", test_tape_nested_struct, devices=devices)
+add_function_test(TestTape, "test_tape_visualize", test_tape_visualize, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_transient_module.py

@@ -0,0 +1,93 @@
+# 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 os
+import tempfile
+import unittest
+from importlib import util
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+CODE = """# -*- coding: utf-8 -*-
+
+import warp as wp
+
+@wp.struct
+class Data:
+    x: wp.array(dtype=int)
+
+@wp.func
+def increment(x: int):
+    # This shouldn't be picked up.
+    return x + 123
+
+@wp.func
+def increment(x: int):
+    return x + 1
+
+@wp.kernel
+def compute(data: Data):
+    data.x[0] = increment(data.x[0])
+"""
+
+
+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 module
+
+
+def test_transient_module(test, device):
+    module = load_code_as_module(CODE, "")
+    # Loading it a second time shouldn't be an issue.
+    module = load_code_as_module(CODE, "")
+
+    assert len(module.compute.module.structs) == 1
+    assert len(module.compute.module.functions) == 1
+
+    data = module.Data()
+    data.x = wp.array([123], dtype=int, device=device)
+
+    wp.set_module_options({"foo": "bar"}, module=module)
+    assert wp.get_module_options(module=module).get("foo") == "bar"
+    assert module.compute.module.options.get("foo") == "bar"
+
+    wp.launch(module.compute, dim=1, inputs=[data], device=device)
+    assert_np_equal(data.x.numpy(), np.array([124]))
+
+
+devices = get_test_devices()
+
+
+class TestTransientModule(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTransientModule, "test_transient_module", test_transient_module, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_triangle_closest_point.py

@@ -0,0 +1,192 @@
+# 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
+
+from warp.tests.unittest_utils import *
+
+
+@wp.func
+def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
+    ab = b - a
+    ac = c - a
+    ap = p - a
+
+    d1 = wp.dot(ab, ap)
+    d2 = wp.dot(ac, ap)
+
+    if d1 <= 0.0 and d2 <= 0.0:
+        return wp.vec3(1.0, 0.0, 0.0)
+
+    bp = p - b
+    d3 = wp.dot(ab, bp)
+    d4 = wp.dot(ac, bp)
+
+    if d3 >= 0.0 and d4 <= d3:
+        return wp.vec3(0.0, 1.0, 0.0)
+
+    vc = d1 * d4 - d3 * d2
+    v = d1 / (d1 - d3)
+    if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
+        return wp.vec3(1.0 - v, v, 0.0)
+
+    cp = p - c
+    d5 = wp.dot(ab, cp)
+    d6 = wp.dot(ac, cp)
+
+    if d6 >= 0.0 and d5 <= d6:
+        return wp.vec3(0.0, 0.0, 1.0)
+
+    vb = d5 * d2 - d1 * d6
+    w = d2 / (d2 - d6)
+    if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
+        return wp.vec3(1.0 - w, 0.0, w)
+
+    va = d3 * d6 - d5 * d4
+    w = (d4 - d3) / ((d4 - d3) + (d5 - d6))
+    if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
+        return wp.vec3(0.0, 1.0 - w, w)
+
+    denom = 1.0 / (va + vb + vc)
+    v = vb * denom
+    w = vc * denom
+
+    return wp.vec3(1.0 - v - w, v, w)
+
+
+# a-b is the edge where the closest point is located at
+@wp.func
+def check_edge_feasible_region(p: wp.vec3, a: wp.vec3, b: wp.vec3, c: wp.vec3, eps: float):
+    ap = p - a
+    bp = p - b
+    ab = b - a
+
+    if wp.dot(ap, ab) < -eps:
+        return False
+
+    if wp.dot(bp, ab) > eps:
+        return False
+
+    ab_sqr_norm = wp.dot(ab, ab)
+    if ab_sqr_norm < eps:
+        return False
+
+    t = wp.dot(ab, c - a) / ab_sqr_norm
+
+    perpendicular_foot = a + t * ab
+
+    if wp.dot(c - perpendicular_foot, p - perpendicular_foot) > eps:
+        return False
+
+    return True
+
+
+# closest point is a
+@wp.func
+def check_vertex_feasible_region(p: wp.vec3, a: wp.vec3, b: wp.vec3, c: wp.vec3, eps: float):
+    ap = p - a
+    ba = a - b
+    ca = a - c
+
+    if wp.dot(ap, ba) < -eps:
+        return False
+
+    if wp.dot(p, ca) < -eps:
+        return False
+
+    return True
+
+
+@wp.kernel
+def test_triangle_closest_point_kernel(tri: wp.array(dtype=wp.vec3), passed: wp.array(dtype=wp.bool)):
+    state = wp.uint32(wp.rand_init(wp.int32(123), wp.int32(0)))
+    eps = 1e-5
+
+    a = tri[0]
+    b = tri[1]
+    c = tri[2]
+
+    for _i in range(1000):
+        l = wp.float32(0.0)
+        while l < eps:
+            p = wp.vec3(wp.randn(state), wp.randn(state), wp.randn(state))
+            l = wp.length(p)
+
+        # project to a sphere with r=2
+        p = 2.0 * p / l
+
+        bary = triangle_closest_point_barycentric(tri[0], tri[1], tri[2], p)
+
+        for dim in range(3):
+            v1_index = (dim + 1) % 3
+            v2_index = (dim + 2) % 3
+            v1 = tri[v1_index]
+            v2 = tri[v2_index]
+            v3 = tri[dim]
+
+            # on edge
+            if bary[dim] == 0.0 and bary[v1_index] != 0.0 and bary[v2_index] != 0.0:
+                if not check_edge_feasible_region(p, v1, v2, v3, eps):
+                    passed[0] = False
+                    return
+
+                # p-closest_p must be perpendicular to v1-v2
+                closest_p = a * bary[0] + b * bary[1] + c * bary[2]
+                e = v1 - v2
+                err = wp.dot(e, closest_p - p)
+                if wp.abs(err) > eps:
+                    passed[0] = False
+                    return
+
+            if bary[v1_index] == 0.0 and bary[v2_index] == 0.0:
+                if not check_vertex_feasible_region(p, v3, v1, v2, eps):
+                    passed[0] = False
+                    return
+
+            if bary[dim] != 0.0 and bary[v1_index] != 0.0 and bary[v2_index] != 0.0:
+                closest_p = a * bary[0] + b * bary[1] + c * bary[2]
+                e1 = v1 - v2
+                e2 = v1 - v3
+                if wp.abs(wp.dot(e1, closest_p - p)) > eps or wp.abs(wp.dot(e2, closest_p - p)) > eps:
+                    passed[0] = False
+                    return
+
+
+def test_triangle_closest_point(test, device):
+    passed = wp.array([True], dtype=wp.bool, device=device)
+
+    a = wp.vec3(1.0, 0.0, 0.0)
+    b = wp.vec3(0.0, 0.0, 0.0)
+    c = wp.vec3(0.0, 1.0, 0.0)
+
+    tri = wp.array([a, b, c], dtype=wp.vec3, device=device)
+    wp.launch(test_triangle_closest_point_kernel, dim=1, inputs=[tri, passed], device=device)
+    passed = passed.numpy()
+
+    test.assertTrue(passed.all())
+
+
+devices = get_test_devices()
+
+
+class TestTriangleClosestPoint(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTriangleClosestPoint, "test_triangle_closest_point", test_triangle_closest_point, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)