warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_dlpack.py

@@ -5,13 +5,14 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
+import ctypes
+import os
 import unittest
-import sys
 
+import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -64,57 +65,69 @@ def test_dlpack_dtypes_and_shapes(test, device):
         test.assertEqual(a1.shape, a2.shape)
         test.assertEqual(a1.strides, a2.strides)
 
-    # convert vec arrays to float arrays
-    def wrap_vec_to_float_tensor(vec_dtype):
+    # convert vector arrays to scalar arrays
+    def wrap_vector_to_scalar_tensor(vec_dtype):
+        scalar_type = vec_dtype._wp_scalar_type_
+        scalar_size = ctypes.sizeof(vec_dtype._type_)
+
         a1 = wp.zeros(10, dtype=vec_dtype, device=device)
-        a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=wp.float32)
+        a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=scalar_type)
 
         test.assertEqual(a1.ptr, a2.ptr)
         test.assertEqual(a1.device, a2.device)
         test.assertEqual(a2.ndim, a1.ndim + 1)
         test.assertEqual(a1.dtype, vec_dtype)
-        test.assertEqual(a2.dtype, wp.float32)
+        test.assertEqual(a2.dtype, scalar_type)
         test.assertEqual(a2.shape, (*a1.shape, vec_dtype._length_))
-        test.assertEqual(a2.strides, (*a1.strides, 4))
+        test.assertEqual(a2.strides, (*a1.strides, scalar_size))
 
-    # convert float arrays to vec arrays
-    def wrap_float_to_vec_tensor(vec_dtype):
-        a1 = wp.zeros((10, vec_dtype._length_), dtype=wp.float32, device=device)
+    # convert scalar arrays to vector arrays
+    def wrap_scalar_to_vector_tensor(vec_dtype):
+        scalar_type = vec_dtype._wp_scalar_type_
+        scalar_size = ctypes.sizeof(vec_dtype._type_)
+
+        a1 = wp.zeros((10, vec_dtype._length_), dtype=scalar_type, device=device)
         a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=vec_dtype)
 
         test.assertEqual(a1.ptr, a2.ptr)
         test.assertEqual(a1.device, a2.device)
         test.assertEqual(a2.ndim, a1.ndim - 1)
-        test.assertEqual(a1.dtype, wp.float32)
+        test.assertEqual(a1.dtype, scalar_type)
         test.assertEqual(a2.dtype, vec_dtype)
         test.assertEqual(a1.shape, (*a2.shape, vec_dtype._length_))
-        test.assertEqual(a1.strides, (*a2.strides, 4))
+        test.assertEqual(a1.strides, (*a2.strides, scalar_size))
+
+    # convert matrix arrays to scalar arrays
+    def wrap_matrix_to_scalar_tensor(mat_dtype):
+        scalar_type = mat_dtype._wp_scalar_type_
+        scalar_size = ctypes.sizeof(mat_dtype._type_)
 
-    # convert mat arrays to float arrays
-    def wrap_mat_to_float_tensor(mat_dtype):
         a1 = wp.zeros(10, dtype=mat_dtype, device=device)
-        a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=wp.float32)
+        a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=scalar_type)
 
         test.assertEqual(a1.ptr, a2.ptr)
         test.assertEqual(a1.device, a2.device)
         test.assertEqual(a2.ndim, a1.ndim + 2)
         test.assertEqual(a1.dtype, mat_dtype)
-        test.assertEqual(a2.dtype, wp.float32)
+        test.assertEqual(a2.dtype, scalar_type)
         test.assertEqual(a2.shape, (*a1.shape, *mat_dtype._shape_))
-        test.assertEqual(a2.strides, (*a1.strides, *wp.types.strides_from_shape(mat_dtype._shape_, wp.float32)))
+        test.assertEqual(a2.strides, (*a1.strides, scalar_size * mat_dtype._shape_[1], scalar_size))
 
-    # convert float arrays to mat arrays
-    def wrap_float_to_mat_tensor(mat_dtype):
-        a1 = wp.zeros((10, *mat_dtype._shape_), dtype=wp.float32, device=device)
+    # convert scalar arrays to matrix arrays
+    def wrap_scalar_to_matrix_tensor(mat_dtype):
+        scalar_type = mat_dtype._wp_scalar_type_
+        scalar_size = ctypes.sizeof(mat_dtype._type_)
+
+        a1 = wp.zeros((10, *mat_dtype._shape_), dtype=scalar_type, device=device)
         a2 = wp.from_dlpack(wp.to_dlpack(a1), dtype=mat_dtype)
 
         test.assertEqual(a1.ptr, a2.ptr)
         test.assertEqual(a1.device, a2.device)
         test.assertEqual(a2.ndim, a1.ndim - 2)
-        test.assertEqual(a1.dtype, wp.float32)
+        test.assertEqual(a1.dtype, scalar_type)
         test.assertEqual(a2.dtype, mat_dtype)
         test.assertEqual(a1.shape, (*a2.shape, *mat_dtype._shape_))
-        test.assertEqual(a1.strides, (*a2.strides, *wp.types.strides_from_shape(mat_dtype._shape_, wp.float32)))
+        test.assertEqual(a1.strides, (*a2.strides, scalar_size * mat_dtype._shape_[1], scalar_size))
 
     for t in wp.types.scalar_types:
         wrap_scalar_tensor_implicit(t)
@@ -132,25 +145,38 @@ def test_dlpack_dtypes_and_shapes(test, device):
     wrap_scalar_tensor_explicit(wp.int64, wp.uint64)
     wrap_scalar_tensor_explicit(wp.uint64, wp.int64)
 
-    wrap_vec_to_float_tensor(wp.vec2)
-    wrap_vec_to_float_tensor(wp.vec3)
-    wrap_vec_to_float_tensor(wp.vec4)
-    wrap_vec_to_float_tensor(wp.spatial_vector)
-    wrap_vec_to_float_tensor(wp.transform)
-
-    wrap_float_to_vec_tensor(wp.vec2)
-    wrap_float_to_vec_tensor(wp.vec3)
-    wrap_float_to_vec_tensor(wp.vec4)
-    wrap_float_to_vec_tensor(wp.spatial_vector)
-    wrap_float_to_vec_tensor(wp.transform)
+    vec_types = []
+    for t in wp.types.scalar_types:
+        for vec_len in [2, 3, 4, 5]:
+            vec_types.append(wp.types.vector(vec_len, t))
+
+    vec_types.append(wp.quath)
+    vec_types.append(wp.quatf)
+    vec_types.append(wp.quatd)
+    vec_types.append(wp.transformh)
+    vec_types.append(wp.transformf)
+    vec_types.append(wp.transformd)
+    vec_types.append(wp.spatial_vectorh)
+    vec_types.append(wp.spatial_vectorf)
+    vec_types.append(wp.spatial_vectord)
+
+    for vec_type in vec_types:
+        wrap_vector_to_scalar_tensor(vec_type)
+        wrap_scalar_to_vector_tensor(vec_type)
+
+    mat_shapes = [(2, 2), (3, 3), (4, 4), (5, 5), (2, 3), (3, 2), (3, 4), (4, 3)]
+    mat_types = []
+    for t in wp.types.scalar_types:
+        for mat_shape in mat_shapes:
+            mat_types.append(wp.types.matrix(mat_shape, t))
 
-    wrap_mat_to_float_tensor(wp.mat22)
-    wrap_mat_to_float_tensor(wp.mat33)
-    wrap_mat_to_float_tensor(wp.mat44)
+    mat_types.append(wp.spatial_matrixh)
+    mat_types.append(wp.spatial_matrixf)
+    mat_types.append(wp.spatial_matrixd)
 
-    wrap_float_to_mat_tensor(wp.mat22)
-    wrap_float_to_mat_tensor(wp.mat33)
-    wrap_float_to_mat_tensor(wp.mat44)
+    for mat_type in mat_types:
+        wrap_matrix_to_scalar_tensor(mat_type)
+        wrap_scalar_to_matrix_tensor(mat_type)
 
 
 def test_dlpack_warp_to_torch(test, device):
@@ -274,75 +300,77 @@ def test_dlpack_jax_to_warp(test, device):
         assert_np_equal(a2.numpy(), np.asarray(j))
 
 
-def register(parent):
-    class TestDLPack(parent):
-        pass
-
-    devices = get_test_devices()
-
-    add_function_test(TestDLPack, "test_dlpack_warp_to_warp", test_dlpack_warp_to_warp, devices=devices)
-    add_function_test(TestDLPack, "test_dlpack_dtypes_and_shapes", test_dlpack_dtypes_and_shapes, devices=devices)
-
-    # torch interop via dlpack
-    try:
-        import torch
-        import torch.utils.dlpack
-
-        # check which Warp devices work with Torch
-        # CUDA devices may fail if Torch was not compiled with CUDA support
-        test_devices = get_test_devices()
-        torch_compatible_devices = []
-        for d in test_devices:
-            try:
-                t = torch.arange(10, device=wp.device_to_torch(d))
-                t += 1
-                torch_compatible_devices.append(d)
-            except Exception as e:
-                print(f"Skipping Torch DLPack tests on device '{d}' due to exception: {e}")
-
-        if torch_compatible_devices:
-            add_function_test(
-                TestDLPack, "test_dlpack_warp_to_torch", test_dlpack_warp_to_torch, devices=torch_compatible_devices
-            )
-            add_function_test(
-                TestDLPack, "test_dlpack_torch_to_warp", test_dlpack_torch_to_warp, devices=torch_compatible_devices
-            )
-
-    except Exception as e:
-        print(f"Skipping Torch DLPack tests due to exception: {e}")
-
-    # jax interop via dlpack
-    try:
-        import jax
-        import jax.dlpack
-
-        # check which Warp devices work with Jax
-        # CUDA devices may fail if Jax cannot find a CUDA Toolkit
-        test_devices = get_test_devices()
-        jax_compatible_devices = []
-        for d in test_devices:
-            try:
-                with jax.default_device(wp.device_to_jax(d)):
-                    j = jax.numpy.arange(10, dtype=jax.numpy.float32)
-                    j += 1
-                jax_compatible_devices.append(d)
-            except Exception as e:
-                print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
-
-        if jax_compatible_devices:
-            add_function_test(
-                TestDLPack, "test_dlpack_warp_to_jax", test_dlpack_warp_to_jax, devices=jax_compatible_devices
-            )
-            add_function_test(
-                TestDLPack, "test_dlpack_jax_to_warp", test_dlpack_jax_to_warp, devices=jax_compatible_devices
-            )
-
-    except Exception as e:
-        print(f"Skipping Jax DLPack tests due to exception: {e}")
-
-    return TestDLPack
+class TestDLPack(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestDLPack, "test_dlpack_warp_to_warp", test_dlpack_warp_to_warp, devices=devices)
+add_function_test(TestDLPack, "test_dlpack_dtypes_and_shapes", test_dlpack_dtypes_and_shapes, devices=devices)
+
+# torch interop via dlpack
+try:
+    import torch
+    import torch.utils.dlpack
+
+    # check which Warp devices work with Torch
+    # CUDA devices may fail if Torch was not compiled with CUDA support
+    test_devices = get_test_devices()
+    torch_compatible_devices = []
+    for d in test_devices:
+        try:
+            t = torch.arange(10, device=wp.device_to_torch(d))
+            t += 1
+            torch_compatible_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Torch DLPack tests on device '{d}' due to exception: {e}")
+
+    if torch_compatible_devices:
+        add_function_test(
+            TestDLPack, "test_dlpack_warp_to_torch", test_dlpack_warp_to_torch, devices=torch_compatible_devices
+        )
+        add_function_test(
+            TestDLPack, "test_dlpack_torch_to_warp", test_dlpack_torch_to_warp, devices=torch_compatible_devices
+        )
+
+except Exception as e:
+    print(f"Skipping Torch DLPack tests due to exception: {e}")
+
+# jax interop via dlpack
+try:
+    # prevent Jax from gobbling up GPU memory
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+    os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+
+    import jax
+    import jax.dlpack
+
+    # check which Warp devices work with Jax
+    # CUDA devices may fail if Jax cannot find a CUDA Toolkit
+    test_devices = get_test_devices()
+    jax_compatible_devices = []
+    for d in test_devices:
+        try:
+            with jax.default_device(wp.device_to_jax(d)):
+                j = jax.numpy.arange(10, dtype=jax.numpy.float32)
+                j += 1
+            jax_compatible_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
+
+    if jax_compatible_devices:
+        add_function_test(
+            TestDLPack, "test_dlpack_warp_to_jax", test_dlpack_warp_to_jax, devices=jax_compatible_devices
+        )
+        add_function_test(
+            TestDLPack, "test_dlpack_jax_to_warp", test_dlpack_jax_to_warp, devices=jax_compatible_devices
+        )
+
+except Exception as e:
+    print(f"Skipping Jax DLPack tests due to exception: {e}")
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_examples.py

@@ -0,0 +1,277 @@
+# 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 (
+    USD_AVAILABLE,
+    get_test_devices,
+    get_unique_cuda_test_devices,
+    sanitize_identifier,
+)
+
+wp.init()
+
+
+# registers an example to run as a TestCase
+def add_example_test(cls, name, devices=None, options={}):
+    def run(test, device):
+        # 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:
+            test.skipTest(f"{e}")
+
+        # create default USD stage output path which many examples expect
+        test_options.setdefault(
+            "stage",
+            os.path.join(os.path.dirname(__file__), f"outputs/{name}_{sanitize_identifier(device)}.usd")
+            if USD_AVAILABLE
+            else None,
+        )
+
+        if test_options["stage"]:
+            try:
+                os.remove(test_options["stage"])
+            except OSError:
+                pass
+
+        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()
+test_devices = get_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
+add_example_test(
+    TestFemExamples,
+    name="fem.example_apic_fluid",
+    devices=test_devices,
+    options={"quiet": True, "res": [16, 16, 16], "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_diffusion",
+    devices=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=test_devices,
+    options={"quiet": True, "resolution": 10, "num_frames": 1, "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_deformed_geometry",
+    devices=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=test_devices,
+    options={"quiet": True, "resolution": 20, "enable_backward": False},
+)
+add_example_test(
+    TestExamples,
+    name="fem.example_convection_diffusion_dg0",
+    devices=test_devices,
+    options={"quiet": True, "resolution": 20, "num_frames": 25, "mesh": "quad", "enable_backward": False},
+)
+add_example_test(
+    TestExamples,
+    name="fem.example_convection_diffusion_dg",
+    devices=test_devices,
+    options={"quiet": True, "resolution": 20, "num_frames": 25, "mesh": "tri", "enable_backward": False},
+)
+add_example_test(
+    TestFemExamples,
+    name="fem.example_mixed_elasticity",
+    devices=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=test_devices,
+    options={"quiet": True, "num_frames": 1, "enable_backward": False},
+)
+add_example_test(
+    TestExamples,
+    name="fem.example_stokes",
+    devices=test_devices,
+    options={
+        "quiet": True,
+        "num_frames": 1,
+        "resolution": 10,
+        "nonconforming_pressures": True,
+        "enable_backward": False,
+    },
+)
+add_example_test(
+    TestExamples,
+    name="fem.example_navier_stokes",
+    devices=test_devices,
+    options={"quiet": True, "num_frames": 100, "resolution": 10, "tri_mesh": True, "enable_backward": False},
+)
+
+
+if __name__ == "__main__":
+    # force rebuild of all kernels
+    wp.build.clear_kernel_cache()
+
+    unittest.main(verbosity=2, failfast=True)