warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/tests/test_examples.py

@@ -1,119 +1,237 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
 import importlib
 import os
 import unittest
 
 import warp as wp
+from warp.tests.unittest_utils import get_unique_cuda_test_devices, sanitize_identifier
+
+wp.init()
 
 
 # registers an example to run as a TestCase
-def add_example_test(cls, name, options):
+def add_example_test(cls, name, devices=None, options={}):
     def run(test, device):
-        # disable scoped timer to avoid log spam
-        wp.ScopedTimer.enabled = False
+        # The copy() is needed because pop() is used to avoid passing extra args to Example()
+        # Can remove once all examples accept **kwargs and we no longer need to pop()
+        test_options = options.copy()
 
         try:
             module = importlib.import_module(f"examples.{name}")
+
+            torch_cuda_required = test_options.setdefault("torch_cuda_required", False)
+            test_options.pop("torch_cuda_required", None)
+            if torch_cuda_required and wp.get_device(device).is_cuda:
+                # Ensure torch has CUDA support
+                import torch
+
+                if not torch.cuda.is_available():
+                    test.skipTest("Torch not compiled with CUDA support")
+
         except Exception as e:
-            print(f"Skipping example: {name}\n    Reason: {e}")
-            return
+            test.skipTest(f"{e}")
 
         # create default USD stage output path which many examples expect
-        if "stage" not in options:
-            stage_path = os.path.join(os.path.dirname(__file__), f"outputs/{name}.usd")
-            options["stage"] = stage_path
-
-        if "num_frames" in options:
-            num_frames = options["num_frames"]
-            del options["num_frames"]
-        else:
-            num_frames = 10
-
-        e = module.Example(**options)
-
-        for _ in range(num_frames):
-            e.update()
-            e.render()
-
-        wp.ScopedTimer.enabled = True
-
-    from warp.tests.test_base import add_function_test
-
-    add_function_test(cls, f"test_{name}", run, check_output=False)
-
-
-def register(parent):
-    class TestExamples(parent):
-        pass
-
-    # Exclude unless we can run headless somehow
-    # add_example_test(TestExamples, name="example_render_opengl", options={})
-
-    # TODO: Test CPU and GPU versions
-    if wp.is_cuda_available():
-        add_example_test(TestExamples, name="example_dem", options={})
-        add_example_test(TestExamples, name="example_diffray", options={})
-        add_example_test(TestExamples, name="example_fluid", options={})
-        add_example_test(TestExamples, name="example_jacobian_ik", options={})
-        add_example_test(TestExamples, name="example_marching_cubes", options={})
-        add_example_test(TestExamples, name="example_mesh", options={})
-        add_example_test(TestExamples, name="example_mesh_intersect", options={"num_frames": 1})
-        add_example_test(TestExamples, name="example_nvdb", options={})
-        add_example_test(TestExamples, name="example_raycast", options={})
-        add_example_test(TestExamples, name="example_raymarch", options={})
-        add_example_test(TestExamples, name="example_sim_cartpole", options={})
-        add_example_test(TestExamples, name="example_sim_cloth", options={})
-        add_example_test(TestExamples, name="example_sim_fk_grad", options={})
-        # add_example_test(TestExamples, name="example_sim_fk_grad_torch", options={})   # disabling due to failure on TC machines that have torch but not CUDA torch
-        add_example_test(TestExamples, name="example_sim_grad_bounce", options={})
-        add_example_test(TestExamples, name="example_sim_grad_cloth", options={})
-        add_example_test(TestExamples, name="example_sim_granular", options={})
-        add_example_test(TestExamples, name="example_sim_granular_collision_sdf", options={})
-        add_example_test(TestExamples, name="example_sim_neo_hookean", options={})
-        add_example_test(TestExamples, name="example_sim_particle_chain", options={})
-        add_example_test(TestExamples, name="example_sim_quadruped", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_chain", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_contact", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_fem", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_force", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_gyroscopic", options={})
-        add_example_test(TestExamples, name="example_sim_rigid_kinematics", options={})
-        add_example_test(TestExamples, name="example_sim_trajopt", options={})
-        add_example_test(TestExamples, name="example_sph", options={})
-        add_example_test(TestExamples, name="example_wave", options={"resx": 256, "resy": 256})
-        add_example_test(TestExamples, name="fem.example_diffusion_mgpu", options={"quiet": True, "num_frames": 1})
-
-    # The following examples do not need cuda
-    add_example_test(TestExamples, name="fem.example_apic_fluid", options={"quiet": True, "res": [16, 16, 16]})
-    add_example_test(
-        TestExamples,
-        name="fem.example_diffusion",
-        options={"quiet": True, "resolution": 10, "mesh": "tri", "num_frames": 1},
-    )
-    add_example_test(
-        TestExamples, name="fem.example_diffusion_3d", options={"quiet": True, "resolution": 10, "num_frames": 1}
-    )
-    add_example_test(
-        TestExamples,
-        name="fem.example_deformed_geometry",
-        options={"quiet": True, "resolution": 10, "num_frames": 1, "mesh": "tri"},
-    )
-    add_example_test(TestExamples, name="fem.example_convection_diffusion", options={"quiet": True, "resolution": 20})
-    add_example_test(
-        TestExamples,
-        name="fem.example_mixed_elasticity",
-        options={"quiet": True, "nonconforming_stresses": True, "mesh": "quad", "num_frames": 1},
-    )
-    add_example_test(TestExamples, name="fem.example_stokes_transfer", options={"quiet": True, "num_frames": 1})
-
-    return TestExamples
+        test_options.setdefault(
+            "stage", os.path.join(os.path.dirname(__file__), f"outputs/{name}_{sanitize_identifier(device)}.usd")
+        )
 
+        try:
+            os.remove(test_options["stage"])
+        except OSError:
+            pass
 
-if __name__ == "__main__":
-    wp.init()
+        num_frames = test_options.get("num_frames", 10)
+        test_options.pop("num_frames", None)
+
+        # Don't want to force load all modules by default for serial test runner
+        wp.config.graph_capture_module_load_default = False
+
+        try:
+            enable_backward = test_options.get("enable_backward", True)
+            wp.set_module_options({"enable_backward": enable_backward}, module)
+            test_options.pop("enable_backward", None)
+
+            with wp.ScopedDevice(device):
+                wp.load_module(module, device=wp.get_device())
+                extra_load_modules = test_options.get("load_modules", [])
+                for module_name in extra_load_modules:
+                    wp.load_module(module_name, device=wp.get_device())
+                test_options.pop("load_modules", None)
+
+                e = module.Example(**test_options)
 
+                # disable scoped timer to avoid log spam from time steps
+                wp.ScopedTimer.enabled = False
+
+                for _ in range(num_frames):
+                    e.update()
+                    e.render()
+        except Exception as e:
+            test.fail(f"{e}")
+        finally:
+            wp.ScopedTimer.enabled = True
+            wp.config.graph_capture_module_load_default = True
+
+    from warp.tests.unittest_utils import add_function_test
+
+    add_function_test(cls, f"test_{name}", run, devices=devices, check_output=False)
+
+
+# TODO: Make the example classes use the passed in device
+cuda_test_devices = get_unique_cuda_test_devices()
+
+# NOTE: To give the parallel test runner more opportunities to parallelize test cases,
+# we break up the tests into multiple TestCase classes that should be non-conflicting
+# w.r.t. kernel compilation
+
+
+class TestExamples(unittest.TestCase):
+    pass
+
+
+# Exclude unless we can run headless somehow
+# add_example_test(TestExamples, name="example_render_opengl", options={})
+
+add_example_test(TestExamples, name="example_dem", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_diffray", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_fluid", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_jacobian_ik", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_marching_cubes", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_mesh", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_mesh_intersect", devices=cuda_test_devices, options={"num_frames": 1})
+add_example_test(TestExamples, name="example_nvdb", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_raycast", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_raymarch", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_sph", devices=cuda_test_devices)
+add_example_test(TestExamples, name="example_wave", devices=cuda_test_devices, options={"resx": 256, "resy": 256})
+
+
+class TestSimExamples(unittest.TestCase):
+    pass
+
+
+add_example_test(
+    TestSimExamples,
+    name="example_sim_cartpole",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.collide", "warp.sim.integrator_euler", "warp.sim.articulation"]},
+)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_cloth",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.collide", "warp.sim.integrator_euler", "warp.sim.particles"]},
+)
+add_example_test(TestSimExamples, name="example_sim_fk_grad", devices=cuda_test_devices)
+add_example_test(
+    TestSimExamples, name="example_sim_fk_grad_torch", devices=cuda_test_devices, options={"torch_cuda_required": True}
+)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_grad_bounce",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.integrator_euler", "warp.sim.particles"]},
+)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_grad_cloth",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.integrator_euler", "warp.sim.particles"]},
+)
+add_example_test(TestSimExamples, name="example_sim_granular", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_granular_collision_sdf", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_neo_hookean", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_particle_chain", devices=cuda_test_devices)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_quadruped",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.integrator_xpbd", "warp.sim.integrator_euler"]},
+)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_rigid_chain",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.integrator_xpbd", "warp.sim.integrator_euler"]},
+)
+add_example_test(
+    TestSimExamples,
+    name="example_sim_rigid_contact",
+    devices=cuda_test_devices,
+    options={"load_modules": ["warp.sim.integrator_euler"]},
+)
+add_example_test(TestSimExamples, name="example_sim_rigid_fem", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_rigid_force", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_rigid_gyroscopic", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_rigid_kinematics", devices=cuda_test_devices)
+add_example_test(TestSimExamples, name="example_sim_trajopt", devices=cuda_test_devices)
+
+
+class TestFemExamples(unittest.TestCase):
+    pass
+
+
+add_example_test(
+    TestFemExamples,
+    name="fem.example_diffusion_mgpu",
+    devices=cuda_test_devices,
+    options={"quiet": True, "num_frames": 1, "enable_backward": False},
+)
+
+# The following examples do not need CUDA, but they need USD
+add_example_test(
+    TestFemExamples,
+    name="fem.example_apic_fluid",
+    devices=cuda_test_devices,
+    options={"quiet": True, "res": [16, 16, 16], "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_diffusion",
+    devices=cuda_test_devices,
+    options={"quiet": True, "resolution": 10, "mesh": "tri", "num_frames": 1, "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_diffusion_3d",
+    devices=cuda_test_devices,
+    options={"quiet": True, "resolution": 10, "num_frames": 1, "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_deformed_geometry",
+    devices=cuda_test_devices,
+    options={"quiet": True, "resolution": 10, "num_frames": 1, "mesh": "tri", "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_convection_diffusion",
+    devices=cuda_test_devices,
+    options={"quiet": True, "resolution": 20, "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_mixed_elasticity",
+    devices=cuda_test_devices,
+    options={"quiet": True, "nonconforming_stresses": True, "mesh": "quad", "num_frames": 1, "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_stokes_transfer",
+    devices=cuda_test_devices,
+    options={"quiet": True, "num_frames": 1, "enable_backward": False},
+)
+
+
+if __name__ == "__main__":
     # force rebuild of all kernels
     wp.build.clear_kernel_cache()
 
-    _ = register(unittest.TestCase)
-
     unittest.main(verbosity=2, failfast=True)

warp/tests/test_fabricarray.py

@@ -12,7 +12,7 @@ from typing import Any
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -930,32 +930,26 @@ for T in _fabric_types:
     wp.overload(fa_generic_sums_kernel_indexed, [wp.indexedfabricarrayarray(dtype=T), wp.array(dtype=T)])
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestFabricArray(parent):
-        pass
 
-    # fabric arrays
-    add_function_test(TestFabricArray, "test_fabricarray_kernel", test_fabricarray_kernel, devices=devices)
-    add_function_test(TestFabricArray, "test_fabricarray_empty", test_fabricarray_empty, devices=devices)
-    add_function_test(
-        TestFabricArray, "test_fabricarray_generic_dtype", test_fabricarray_generic_dtype, devices=devices
-    )
-    add_function_test(
-        TestFabricArray, "test_fabricarray_generic_array", test_fabricarray_generic_array, devices=devices
-    )
-    add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
-    add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
-    add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)
+class TestFabricArray(unittest.TestCase):
+    pass
 
-    # fabric arrays of arrays
-    add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)
 
-    return TestFabricArray
+# fabric arrays
+add_function_test(TestFabricArray, "test_fabricarray_kernel", test_fabricarray_kernel, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_empty", test_fabricarray_empty, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_generic_dtype", test_fabricarray_generic_dtype, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_generic_array", test_fabricarray_generic_array, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)
+
+# fabric arrays of arrays
+add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_fast_math.py

@@ -1,11 +1,14 @@
-import warp as wp
-import numpy as np
+# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import unittest
 
 import warp as wp
-from warp.tests.test_base import *
-
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -33,19 +36,19 @@ def test_fast_math(test, device):
             with CheckOutput():
                 wp.launch(test_pow, dim=1, inputs=[-2.0, 2.0, 2.0], device=device)
 
+        # Turn fast math back off
+        wp.set_module_options({"fast_math": False})
+
 
-def register(parent):
-    class TestFastMath(parent):
-        pass
+class TestFastMath(unittest.TestCase):
+    pass
 
-    devices = get_test_devices()
 
-    add_function_test(TestFastMath, "test_fast_math", test_fast_math, devices=devices)
+devices = get_test_devices()
 
-    return TestFastMath
+add_function_test(TestFastMath, "test_fast_math", test_fast_math, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)