warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/tests/test_rand.py

@@ -0,0 +1,339 @@
+# 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 test_kernel(
+    kernel_seed: int,
+    int_a: wp.array(dtype=int),
+    int_ab: wp.array(dtype=int),
+    uint_a: wp.array(dtype=wp.uint32),
+    uint_ab: wp.array(dtype=wp.uint32),
+    float_01: wp.array(dtype=float),
+    float_ab: wp.array(dtype=float),
+):
+    tid = wp.tid()
+
+    state = wp.rand_init(kernel_seed, tid)
+
+    int_a[tid] = wp.randi(state)
+    int_ab[tid] = wp.randi(state, 0, 100)
+    uint_a[tid] = wp.randu(state)
+    uint_ab[tid] = wp.randu(state, wp.uint32(0), wp.uint32(100))
+    float_01[tid] = wp.randf(state)
+    float_ab[tid] = wp.randf(state, 0.0, 100.0)
+
+
+def test_rand(test, device):
+    N = 10
+
+    int_a = wp.zeros(N, dtype=int, device=device)
+    int_ab = wp.zeros(N, dtype=int, device=device)
+
+    uint_a = wp.zeros(N, dtype=wp.uint32, device=device)
+    uint_ab = wp.zeros(N, dtype=wp.uint32, device=device)
+
+    float_01 = wp.zeros(N, dtype=float, device=device)
+    float_ab = wp.zeros(N, dtype=float, device=device)
+
+    seed = 42
+
+    wp.launch(
+        kernel=test_kernel,
+        dim=N,
+        inputs=[seed, int_a, int_ab, uint_a, uint_ab, float_01, float_ab],
+        outputs=[],
+        device=device,
+    )
+
+    int_a_true = np.array(
+        [
+            -575632308,
+            59537738,
+            1898992239,
+            442961864,
+            -1069147335,
+            -478445524,
+            1803659809,
+            2122909397,
+            -1888556360,
+            334603718,
+        ]
+    )
+    int_ab_true = np.array([46, 58, 46, 83, 85, 39, 72, 99, 18, 41])
+    uint_a_true = np.array(
+        [
+            3133687854,
+            3702303309,
+            1235698096,
+            3516599792,
+            800302729,
+            2620462179,
+            2423739693,
+            3024873594,
+            2783682377,
+            1188846332,
+        ]
+    )
+    uint_ab_true = np.array([6, 55, 2, 92, 55, 93, 65, 23, 48, 0])
+    float_01_true = np.array(
+        [
+            0.8265858,
+            0.5874614,
+            0.1508659,
+            0.9498008,
+            0.02531803,
+            0.8520948,
+            0.0001185536,
+            0.4855958,
+            0.06277305,
+            0.2214079,
+        ]
+    )
+    float_ab_true = np.array(
+        [79.84678, 76.362206, 32.135242, 99.70866, 70.45863, 20.6523, 45.164482, 55.583008, 76.60291, 35.36277]
+    )
+
+    assert_np_equal(int_a.numpy(), int_a_true)
+    assert_np_equal(int_ab.numpy(), int_ab_true)
+
+    assert_np_equal(uint_a.numpy(), uint_a_true)
+    assert_np_equal(uint_ab.numpy(), uint_ab_true)
+
+    assert_np_equal(float_01.numpy(), float_01_true, 1e-04)
+    assert_np_equal(float_ab.numpy(), float_ab_true, 1e-04)
+
+
+@wp.kernel
+def randn_kernel(
+    x: wp.array(dtype=float),
+):
+    tid = wp.tid()
+    r = wp.rand_init(tid)
+    x[tid] = wp.randn(r)
+
+
+def test_randn(test, device):
+    N = 100000000
+
+    samples = wp.zeros(N, dtype=float, device=device)
+
+    wp.launch(randn_kernel, inputs=[samples], dim=N, device=device)
+
+    randn_samples = samples.numpy()
+
+    test.assertFalse(np.any(np.isinf(randn_samples)))
+    test.assertFalse(np.any(np.isnan(randn_samples)))
+
+    randn_true = np.array(
+        [
+            -1.8213255,
+            0.27881497,
+            -1.1122388,
+            0.5936895,
+            0.04976363,
+            0.69087356,
+            0.2935363,
+            0.8405019,
+            -0.8436684,
+            0.53108305,
+        ]
+    )
+
+    err = np.max(np.abs(randn_samples[0:10] - randn_true))
+    test.assertTrue(err < 1e-04)
+
+
+@wp.kernel
+def sample_cdf_kernel(kernel_seed: int, cdf: wp.array(dtype=float), samples: wp.array(dtype=int)):
+    tid = wp.tid()
+    state = wp.rand_init(kernel_seed, tid)
+
+    samples[tid] = wp.sample_cdf(state, cdf)
+
+
+def test_sample_cdf(test, device):
+    seed = 42
+    cdf = np.arange(0.0, 1.0, 0.01, dtype=float)
+    cdf = cdf * cdf
+    cdf = wp.array(cdf, dtype=float, device=device)
+    num_samples = 1000
+    samples = wp.zeros(num_samples, dtype=int, device=device)
+
+    wp.launch(kernel=sample_cdf_kernel, dim=num_samples, inputs=[seed, cdf, samples], device=device)
+
+    # histogram should be linear
+    # plt.hist(samples.numpy())
+    # plt.show()
+
+
+@wp.kernel
+def sampling_kernel(
+    kernel_seed: int,
+    triangle_samples: wp.array(dtype=wp.vec2),
+    square_samples: wp.array(dtype=wp.vec2),
+    ring_samples: wp.array(dtype=wp.vec2),
+    disk_samples: wp.array(dtype=wp.vec2),
+    sphere_surface_samples: wp.array(dtype=wp.vec3),
+    sphere_samples: wp.array(dtype=wp.vec3),
+    hemisphere_surface_samples: wp.array(dtype=wp.vec3),
+    hemisphere_samples: wp.array(dtype=wp.vec3),
+    cube_samples: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    state = wp.rand_init(kernel_seed, tid)
+
+    triangle_samples[tid] = wp.sample_triangle(state)
+    ring_samples[tid] = wp.sample_unit_ring(state)
+    disk_samples[tid] = wp.sample_unit_disk(state)
+    sphere_surface_samples[tid] = wp.sample_unit_sphere_surface(state)
+    sphere_samples[tid] = wp.sample_unit_sphere(state)
+    hemisphere_surface_samples[tid] = wp.sample_unit_hemisphere_surface(state)
+    hemisphere_samples[tid] = wp.sample_unit_hemisphere(state)
+    square_samples[tid] = wp.sample_unit_square(state)
+    cube_samples[tid] = wp.sample_unit_cube(state)
+
+
+def test_sampling_methods(test, device):
+    seed = 42
+    num_samples = 100
+
+    triangle_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
+    square_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
+    ring_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
+    disk_samples = wp.zeros(num_samples, dtype=wp.vec2, device=device)
+    sphere_surface_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
+    sphere_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
+    hemisphere_surface_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
+    hemisphere_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
+    cube_samples = wp.zeros(num_samples, dtype=wp.vec3, device=device)
+
+    wp.launch(
+        kernel=sampling_kernel,
+        dim=num_samples,
+        inputs=[
+            seed,
+            triangle_samples,
+            square_samples,
+            ring_samples,
+            disk_samples,
+            sphere_surface_samples,
+            sphere_samples,
+            hemisphere_surface_samples,
+            hemisphere_samples,
+            cube_samples,
+        ],
+        device=device,
+    )
+
+    # bounds check
+    test.assertTrue((triangle_samples.numpy()[:, 0] <= 1.0).all())
+    test.assertTrue((triangle_samples.numpy()[:, 0] >= 0.0).all())
+    test.assertTrue((triangle_samples.numpy()[:, 1] >= 0.0).all())
+    test.assertTrue((triangle_samples.numpy()[:, 1] >= 0.0).all())
+    test.assertTrue((square_samples.numpy()[:, 0] >= -0.5).all())
+    test.assertTrue((square_samples.numpy()[:, 0] <= 1.5).all())
+    test.assertTrue((square_samples.numpy()[:, 1] >= -0.5).all())
+    test.assertTrue((square_samples.numpy()[:, 1] <= 0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 0] >= -0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 0] <= 0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 1] >= -0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 1] <= 0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 2] >= -0.5).all())
+    test.assertTrue((cube_samples.numpy()[:, 2] <= 0.5).all())
+    test.assertTrue((hemisphere_surface_samples.numpy()[:, 2] >= 0.0).all())
+    test.assertTrue((hemisphere_samples.numpy()[:, 2] >= 0.0).all())
+    test.assertTrue((np.linalg.norm(ring_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+    test.assertTrue((np.linalg.norm(disk_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+    test.assertTrue((np.linalg.norm(sphere_surface_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+    test.assertTrue((np.linalg.norm(sphere_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+    test.assertTrue((np.linalg.norm(hemisphere_surface_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+    test.assertTrue((np.linalg.norm(hemisphere_samples.numpy(), axis=1) <= 1.0 + 1e6).all())
+
+
+@wp.kernel
+def sample_poisson_kernel(
+    kernel_seed: int, poisson_samples_low: wp.array(dtype=wp.uint32), poisson_samples_high: wp.array(dtype=wp.uint32)
+):
+    tid = wp.tid()
+    state = wp.rand_init(kernel_seed, tid)
+
+    x = wp.poisson(state, 3.0)
+    y = wp.poisson(state, 42.0)
+
+    poisson_samples_low[tid] = x
+    poisson_samples_high[tid] = y
+
+
+def test_poisson(test, device):
+    seed = 13
+    N = 20000
+    poisson_low = wp.zeros(N, dtype=wp.uint32, device=device)
+    poisson_high = wp.zeros(N, dtype=wp.uint32, device=device)
+
+    wp.launch(kernel=sample_poisson_kernel, dim=N, inputs=[seed, poisson_low, poisson_high], device=device)
+
+    # bins = np.arange(100)
+    # _ = plt.hist(poisson_high.numpy(), bins)
+    # plt.show()
+
+    rng = np.random.default_rng(seed)
+
+    np_poisson_low = rng.poisson(lam=3.0, size=N)
+    np_poisson_high = rng.poisson(lam=42.0, size=N)
+
+    poisson_low_mean = np.mean(poisson_low.numpy())
+    np_poisson_low_mean = np.mean(np_poisson_low)
+
+    poisson_high_mean = np.mean(poisson_high.numpy())
+    np_poisson_high_mean = np.mean(np_poisson_high)
+
+    poisson_low_std = np.std(poisson_low.numpy())
+    np_poisson_low_std = np.std(np_poisson_low)
+
+    poisson_high_std = np.std(poisson_high.numpy())
+    np_poisson_high_std = np.std(np_poisson_high)
+
+    # compare basic distribution characteristics
+    test.assertTrue(np.abs(poisson_low_mean - np_poisson_low_mean) <= 5e-1)
+    test.assertTrue(np.abs(poisson_high_mean - np_poisson_high_mean) <= 5e-1)
+    test.assertTrue(np.abs(poisson_low_std - np_poisson_low_std) <= 2e-1)
+    test.assertTrue(np.abs(poisson_high_std - np_poisson_high_std) <= 2e-1)
+
+
+devices = get_test_devices()
+
+
+class TestRand(unittest.TestCase):
+    pass
+
+
+add_function_test(TestRand, "test_rand", test_rand, devices=devices)
+add_function_test(TestRand, "test_randn", test_randn, devices=devices)
+add_function_test(TestRand, "test_sample_cdf", test_sample_cdf, devices=devices)
+add_function_test(TestRand, "test_sampling_methods", test_sampling_methods, devices=devices)
+add_function_test(TestRand, "test_poisson", test_poisson, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_reload.py

@@ -0,0 +1,302 @@
+# 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 importlib
+import os
+import unittest
+
+import numpy as np
+
+import warp as wp
+
+# dummy modules used for testing reload with dependencies
+import warp.tests.aux_test_dependent as test_dependent
+import warp.tests.aux_test_reference as test_reference
+import warp.tests.aux_test_reference_reference as test_reference_reference
+
+# dummy module used for testing reload
+import warp.tests.aux_test_square as test_square
+from warp.tests.unittest_utils import *
+
+
+def reload_module(module):
+    # Clearing the .pyc file associated with a module is a necessary workaround
+    # for `importlib.reload` to work as expected when run from within Kit.
+    cache_file = importlib.util.cache_from_source(module.__file__)
+    os.remove(cache_file)
+    importlib.reload(module)
+
+
+def test_redefine(test, device):
+    # --------------------------------------------
+    # first pass
+
+    @wp.kernel
+    def basic(x: wp.array(dtype=float)):
+        tid = wp.tid()
+
+        x[tid] = float(tid) * 1.0
+
+    n = 32
+
+    x = wp.zeros(n, dtype=float, device=device)
+
+    wp.launch(kernel=basic, dim=n, inputs=[x], device=device)
+
+    # --------------------------------------------
+    # redefine kernel, should trigger a recompile
+
+    @wp.kernel
+    def basic(x: wp.array(dtype=float)):
+        tid = wp.tid()
+
+        x[tid] = float(tid) * 2.0
+
+    y = wp.zeros(n, dtype=float, device=device)
+
+    wp.launch(kernel=basic, dim=n, inputs=[y], device=device)
+
+    assert_np_equal(np.arange(0, n, 1), x.numpy())
+    assert_np_equal(np.arange(0, n, 1) * 2.0, y.numpy())
+
+
+def test_redefine_command(test, device):
+    # Test whether executable modules are retained correctly.
+    # A module can have multiple executable versions in use if something
+    # still holds a reference to them.
+
+    @wp.kernel
+    def k(value: int):
+        wp.expect_eq(value, 17)
+
+    # record a command
+    cmd1 = wp.launch(k, dim=1, inputs=[17], device=device, record_cmd=True)
+    cmd1.launch()
+
+    # redefine the kernel, triggering module reload
+    @wp.kernel
+    def k(value: int):
+        wp.expect_eq(value, 42)
+
+    # re-record the command
+    cmd2 = wp.launch(k, dim=1, inputs=[42], device=device, record_cmd=True)
+    cmd2.launch()
+
+    # previous command should still work
+    cmd1.launch()
+
+
+square_two = """import warp as wp
+
+
+@wp.func
+def sqr(x: float):
+    return x * x
+
+
+@wp.kernel
+def kern(expect: float):
+    wp.expect_eq(sqr(2.0), expect)
+
+
+def run(expect, device):
+    wp.launch(kern, dim=1, inputs=[expect], device=device)
+    wp.synchronize_device(device)
+"""
+
+square_four = """import warp as wp
+
+
+@wp.func
+def multiply(x: float):
+    return x * x
+
+
+@wp.kernel
+def kern(expect: float):
+    wp.expect_eq(multiply(4.0), expect)
+
+
+def run(expect, device):
+    wp.launch(kern, dim=1, inputs=[expect], device=device)
+    wp.synchronize_device(device)
+"""
+
+
+def test_reload(test, device):
+    # write out the module python and import it
+    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_square.py")), "w")
+    f.writelines(square_two)
+    f.flush()
+    f.close()
+
+    reload_module(test_square)
+    test_square.run(expect=4.0, device=device)  # 2*2=4
+
+    f = open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_square.py")), "w")
+    f.writelines(square_four)
+    f.flush()
+    f.close()
+
+    # reload module, this should trigger all of the funcs / kernels to be updated
+    reload_module(test_square)
+    test_square.run(expect=16.0, device=device)  # 4*4 = 16
+
+
+def test_reload_class(test, device):
+    def test_func():
+        import importlib as imp
+
+        import warp.tests.aux_test_class_kernel
+        from warp.tests.aux_test_class_kernel import ClassKernelTest
+
+        imp.reload(warp.tests.aux_test_class_kernel)
+
+        ctest = ClassKernelTest(device)
+        expected = np.zeros((10, 3, 3), dtype=np.float32)
+        expected[:] = np.eye(3)
+        assert_np_equal(expected, ctest.identities.numpy())
+
+    test_func()
+    test_func()
+
+
+template_ref = """# This file is used to test reloading module references.
+
+import warp as wp
+import warp.tests.aux_test_reference_reference as refref
+
+
+@wp.func
+def magic():
+    return {} * refref.more_magic()
+"""
+
+template_refref = """# This file is used to test reloading module references.
+
+import warp as wp
+
+
+@wp.func
+def more_magic():
+    return {}
+"""
+
+
+def test_reload_references(test, device):
+    path_ref = os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_reference.py"))
+    path_refref = os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_reference_reference.py"))
+
+    # rewrite both dependency modules and reload them
+    with open(path_ref, "w") as f:
+        f.writelines(template_ref.format(1.0))
+    importlib.reload(test_reference)
+
+    with open(path_refref, "w") as f:
+        f.writelines(template_refref.format(1.0))
+    importlib.reload(test_reference_reference)
+
+    test_dependent.run(expect=1.0, device=device)  # 1 * 1 = 1
+
+    # rewrite and reload the first dependency module
+    with open(path_ref, "w") as f:
+        f.writelines(template_ref.format(2.0))
+    importlib.reload(test_reference)
+
+    test_dependent.run(expect=2.0, device=device)  # 2 * 1 = 1
+
+    # rewrite and reload the second dependency module
+    with open(path_refref, "w") as f:
+        f.writelines(template_refref.format(2.0))
+    importlib.reload(test_reference_reference)
+
+    test_dependent.run(expect=4.0, device=device)  # 2 * 2 = 4
+
+
+def test_graph_launch_after_module_reload(test, device):
+    @wp.kernel
+    def foo(a: wp.array(dtype=int)):
+        a[0] = 42
+
+    with wp.ScopedDevice(device):
+        a = wp.zeros(1, dtype=int)
+
+        # preload module before graph capture
+        wp.load_module(device=device)
+
+        # capture a launch
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            wp.launch(foo, dim=1, inputs=[a])
+
+        # unload the module
+        foo.module.unload()
+
+        # launch previously captured graph
+        wp.capture_launch(capture.graph)
+
+        test.assertEqual(a.numpy()[0], 42)
+
+
+def test_module_unload_during_graph_capture(test, device):
+    @wp.kernel
+    def foo(a: wp.array(dtype=int)):
+        a[0] = 42
+
+    # preload module before graph capture
+    wp.load_module(device=device)
+
+    # load another module to test unloading during graph capture
+    other_module = wp.get_module("warp.tests.aux_test_module_unload")
+    other_module.load(device)
+
+    with wp.ScopedDevice(device):
+        a = wp.zeros(1, dtype=int)
+
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            wp.launch(foo, dim=1, inputs=[a])
+
+            # unloading a module during graph capture should be fine (deferred until capture completes)
+            other_module.unload()
+
+        wp.capture_launch(capture.graph)
+
+        test.assertEqual(a.numpy()[0], 42)
+
+
+devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
+
+
+class TestReload(unittest.TestCase):
+    pass
+
+
+add_function_test(TestReload, "test_redefine", test_redefine, devices=devices)
+add_function_test(TestReload, "test_redefine_command", test_redefine_command, devices=devices)
+add_function_test(TestReload, "test_reload", test_reload, devices=devices)
+add_function_test(TestReload, "test_reload_class", test_reload_class, devices=devices)
+# TODO: test_reload_references sometimes has issues running on cuda:1
+add_function_test(TestReload, "test_reload_references", test_reload_references, devices=get_test_devices("basic"))
+add_function_test(
+    TestReload, "test_graph_launch_after_module_reload", test_graph_launch_after_module_reload, devices=cuda_devices
+)
+add_function_test(
+    TestReload, "test_module_unload_during_graph_capture", test_module_unload_during_graph_capture, devices=cuda_devices
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)