warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/benchmarks/benchmark_tile_load_store.py

@@ -0,0 +1,103 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 numpy as np
+
+import warp as wp
+
+BLOCK_DIM = 128
+
+TILE = 32
+
+
+def create_test_kernel(storage_type: str):
+    @wp.kernel
+    def load_store(a: wp.array2d(dtype=wp.float32), b: wp.array2d(dtype=wp.float32)):
+        i, j = wp.tid()
+
+        if wp.static(storage_type == "shared"):
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="shared")
+        else:
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="register")
+
+        wp.tile_store(b, a_tile, offset=(i * TILE, j * TILE))
+
+    return load_store
+
+
+if __name__ == "__main__":
+    wp.config.quiet = True
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    iterations = 100
+    rng = np.random.default_rng(42)
+
+    shared_benchmark_data = {}
+    register_benchmark_data = {}
+    memcpy_benchmark_data = {}
+
+    sizes = list(range(128, 4097, 128))
+
+    print(f"{'Transfer Size (Bytes)':<23s} {'Shared (GiB/s)':<16s} {'Register (GiB/s)':<18s} {'memcpy (GiB/s)':<16s}")
+    print("-" * 79)
+
+    for size in sizes:
+        a = wp.array(rng.random((size, size), dtype=np.float32), dtype=wp.float32)
+        b = wp.empty_like(a)
+
+        for storage_type in ("shared", "register"):
+            load_store = create_test_kernel(storage_type)
+
+            cmd = wp.launch_tiled(
+                load_store,
+                dim=(a.shape[0] // TILE, a.shape[1] // TILE),
+                inputs=[a],
+                outputs=[b],
+                block_dim=BLOCK_DIM,
+                record_cmd=True,
+            )
+            # Warmup
+            for _ in range(5):
+                cmd.launch()
+
+            with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_KERNEL, print=False, synchronize=True) as timer:
+                for _ in range(iterations):
+                    cmd.launch()
+
+            np.testing.assert_equal(a.numpy(), b.numpy())
+
+            timing_results = [result.elapsed for result in timer.timing_results]
+            avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+
+            if storage_type == "shared":
+                shared_benchmark_data[a.capacity] = avg_bw
+            else:
+                register_benchmark_data[a.capacity] = avg_bw
+
+        # Compare with memcpy
+        with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_MEMCPY, print=False, synchronize=True) as timer:
+            for _ in range(iterations):
+                wp.copy(b, a)
+
+        timing_results = [result.elapsed for result in timer.timing_results]
+        avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+        memcpy_benchmark_data[a.capacity] = avg_bw
+
+        # Print results
+        print(
+            f"{a.capacity:<23d} {shared_benchmark_data[a.capacity]:<#16.4g} {register_benchmark_data[a.capacity]:<#18.4g} {memcpy_benchmark_data[a.capacity]:<#16.4g}"
+        )

warp/examples/benchmarks/benchmark_tile_sort.py

@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 time
+
+import numpy as np
+
+import warp as wp
+
+BLOCK_DIM = 128
+
+
+def create_test_kernel(KEY_TYPE, MAX_SORT_LENGTH):
+    @wp.kernel
+    def tile_sort_kernel(
+        input_keys: wp.array(dtype=KEY_TYPE, ndim=2),
+        input_values: wp.array(dtype=wp.int32, ndim=2),
+        output_keys: wp.array(dtype=KEY_TYPE, ndim=2),
+        output_values: wp.array(dtype=wp.int32, ndim=2),
+    ):
+        batch_id, i = wp.tid()
+
+        # Load input into shared memory
+        keys = wp.tile_load(input_keys[batch_id], shape=MAX_SORT_LENGTH, storage="shared")
+        values = wp.tile_load(input_values[batch_id], shape=MAX_SORT_LENGTH, storage="shared")
+
+        # Perform in-place sorting
+        wp.tile_sort(keys, values)
+
+        # Store sorted shared memory into output arrays
+        wp.tile_store(output_keys[batch_id], keys)
+        wp.tile_store(output_values[batch_id], values)
+
+    return tile_sort_kernel
+
+
+if __name__ == "__main__":
+    wp.config.quiet = True
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    iterations = 100
+    rng = np.random.default_rng(42)
+
+    shared_benchmark_data = {}
+    cub_segmented_sort_benchmark_data = {}
+
+    array_length = list(range(16, 257, 16))
+
+    print(
+        f"{'Type':<12s} {'Batch Size':<12s} {'Length':<12s} {'Tile Sort (ms)':<16s} {'Cub Segmented Sort (ms)':<24s} {'CubTime/TileTime':<16s}"
+    )
+    print("-" * 100)
+
+    for dtype in [int, float]:
+        for batch_size_exponent in range(5, 11):
+            batch_size = 2**batch_size_exponent
+            for length in array_length:
+                if dtype == int:
+                    np_keys = rng.choice(1000000000, size=(batch_size, length), replace=False)
+                else:  # dtype == float
+                    np_keys = rng.choice(1000000, size=(batch_size, length), replace=False).astype(np.float32)
+
+                np_values = np.tile(np.arange(length), (batch_size, 1))
+
+                # Sort using NumPy for validation
+                np_sorted_keys = np.zeros_like(np_keys)
+                np_sorted_values = np.zeros_like(np_values)
+                for b in range(batch_size):
+                    sorted_indices = np.argsort(np_keys[b])
+                    np_sorted_keys[b] = np_keys[b][sorted_indices]
+                    np_sorted_values[b] = np_values[b][sorted_indices]
+
+                # Generate random keys and iota indexer
+                input_keys = wp.array(np_keys, dtype=dtype, ndim=2, device="cuda")
+                input_values = wp.array(np_values, dtype=int, ndim=2, device="cuda")
+                output_keys = wp.zeros_like(input_keys, device="cuda")
+                output_values = wp.zeros_like(input_values, device="cuda")
+
+                kernel = create_test_kernel(dtype, length)
+
+                cmd = wp.launch_tiled(
+                    kernel,
+                    dim=batch_size,
+                    inputs=[input_keys, input_values, output_keys, output_values],
+                    block_dim=BLOCK_DIM,
+                    record_cmd=True,
+                )
+                # Warmup
+                for _ in range(5):
+                    cmd.launch()
+
+                with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_KERNEL, print=False, synchronize=True) as timer:
+                    for _ in range(iterations):
+                        cmd.launch()
+                    wp.synchronize()
+
+                if dtype == int:
+                    keys_match = np.array_equal(output_keys.numpy(), np_sorted_keys)
+                else:  # dtype == float
+                    keys_match = np.allclose(output_keys.numpy(), np_sorted_keys, atol=1e-6)  # Use tolerance for floats
+
+                values_match = np.array_equal(output_values.numpy(), np_sorted_values)
+
+                # Validate results
+                assert keys_match, f"Key sorting mismatch for dtype={dtype}!"
+                assert values_match, f"Value sorting mismatch for dtype={dtype}!"
+
+                timing_results = [result.elapsed for result in timer.timing_results]
+                mean_timing = np.mean(timing_results)
+
+                shared_benchmark_data[length] = mean_timing
+
+                # Allocate memory
+                input_keys = wp.zeros(shape=(batch_size * 2, length), dtype=dtype, device="cuda")
+                input_values = wp.zeros(shape=(batch_size * 2, length), dtype=int, device="cuda")
+
+                # Copy data
+                input_keys.assign(np_keys)
+                input_values.assign(np_values)
+
+                input_keys = input_keys.reshape(-1)
+                input_values = input_values.reshape(-1)
+
+                segments = wp.array(np.arange(0, batch_size + 1) * length, dtype=int, device="cuda")
+
+                # Compare with cub segmented radix sort
+                # Warmup
+                for _ in range(5):
+                    wp.utils.segmented_sort_pairs(input_keys, input_values, batch_size * length, segments)
+
+                t1 = time.time_ns()
+                for _ in range(iterations):
+                    wp.utils.segmented_sort_pairs(input_keys, input_values, batch_size * length, segments)
+                wp.synchronize()
+                t2 = time.time_ns()
+                cub_segmented_sort_benchmark_data[length] = (t2 - t1) / (1_000_000 * iterations)
+
+                # Print results
+                print(
+                    f"{dtype!s:<12s} {batch_size:<12d} {length:<12d} {shared_benchmark_data[length]:<16.4g} {cub_segmented_sort_benchmark_data[length]:<24.4g} {cub_segmented_sort_benchmark_data[length] / shared_benchmark_data[length]:<16.4g}"
+                )

warp/examples/browse.py

@@ -0,0 +1,37 @@
+# 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 os
+import subprocess
+import sys
+
+
+def open_file(filename):
+    if sys.platform == "win32":
+        os.startfile(filename)
+    else:
+        subprocess.call(["xdg-open", filename])
+
+
+if __name__ == "__main__":
+    import warp.examples
+
+    source_dir = warp.examples.get_source_directory()
+    print(f"Example source directory: {source_dir}")
+
+    try:
+        open_file(source_dir)
+    except Exception:
+        pass

warp/examples/core/example_cupy.py

@@ -0,0 +1,86 @@
+# 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.
+
+###########################################################################
+# Example CuPy
+#
+# The example demonstrates interoperability with CuPy on CUDA devices
+# and NumPy on CPU devices.
+###########################################################################
+
+import warp as wp
+
+
+@wp.kernel
+def saxpy(x: wp.array(dtype=float), y: wp.array(dtype=float), a: float):
+    i = wp.tid()
+    y[i] = a * x[i] + y[i]
+
+
+class Example:
+    def __init__(self):
+        device = wp.get_device()
+
+        self.n = 10
+        self.a = 1.0
+
+        if device.is_cuda:
+            # use CuPy arrays on CUDA devices
+            import cupy as cp
+
+            print(f"Using CuPy on device {device}")
+
+            # tell CuPy to use the same device
+            with cp.cuda.Device(device.ordinal):
+                self.x = cp.arange(self.n, dtype=cp.float32)
+                self.y = cp.ones(self.n, dtype=cp.float32)
+        else:
+            # use NumPy arrays on CPU
+            import numpy as np
+
+            print("Using NumPy on CPU")
+
+            self.x = np.arange(self.n, dtype=np.float32)
+            self.y = np.ones(self.n, dtype=np.float32)
+
+    def step(self):
+        # Launch a Warp kernel on the pre-allocated arrays.
+        # When running on a CUDA device, these are CuPy arrays.
+        # When running on the CPU, these are NumPy arrays.
+        #
+        # Note that the arrays can be passed to Warp kernels directly.  Under the hood,
+        # Warp uses the __cuda_array_interface__ and __array_interface__ protocols to
+        # access the data.
+        wp.launch(saxpy, dim=self.n, inputs=[self.x, self.y, self.a])
+
+    def render(self):
+        print(self.y)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--num_frames", type=int, default=10, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example()
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()

warp/examples/core/example_dem.py

@@ -0,0 +1,241 @@
+# 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.
+
+###########################################################################
+# Example DEM
+#
+# Shows how to implement a DEM particle simulation with cohesion between
+# particles. Neighbors are found using the wp.HashGrid class, and
+# wp.hash_grid_query(), wp.hash_grid_query_next() kernel methods.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+
+@wp.func
+def contact_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float):
+    vn = wp.dot(n, v)
+    jn = c * k_n
+    jd = min(vn, 0.0) * k_d
+
+    # contact force
+    fn = jn + jd
+
+    # friction force
+    vt = v - n * vn
+    vs = wp.length(vt)
+
+    if vs > 0.0:
+        vt = vt / vs
+
+    # Coulomb condition
+    ft = wp.min(vs * k_f, k_mu * wp.abs(fn))
+
+    # total force
+    return -n * fn - vt * ft
+
+
+@wp.kernel
+def apply_forces(
+    grid: wp.uint64,
+    particle_x: wp.array(dtype=wp.vec3),
+    particle_v: wp.array(dtype=wp.vec3),
+    particle_f: wp.array(dtype=wp.vec3),
+    radius: float,
+    k_contact: float,
+    k_damp: float,
+    k_friction: float,
+    k_mu: float,
+):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    x = particle_x[i]
+    v = particle_v[i]
+
+    f = wp.vec3()
+
+    # ground contact
+    n = wp.vec3(0.0, 1.0, 0.0)
+    c = wp.dot(n, x)
+
+    cohesion_ground = 0.02
+    cohesion_particle = 0.0075
+
+    if c < cohesion_ground:
+        f = f + contact_force(n, v, c, k_contact, k_damp, 100.0, 0.5)
+
+    # particle contact
+    neighbors = wp.hash_grid_query(grid, x, radius * 5.0)
+
+    for index in neighbors:
+        if index != i:
+            # compute distance to point
+            n = x - particle_x[index]
+            d = wp.length(n)
+            err = d - radius * 2.0
+
+            if err <= cohesion_particle:
+                n = n / d
+                vrel = v - particle_v[index]
+
+                f = f + contact_force(n, vrel, err, k_contact, k_damp, k_friction, k_mu)
+
+    particle_f[i] = f
+
+
+@wp.kernel
+def integrate(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    f: wp.array(dtype=wp.vec3),
+    gravity: wp.vec3,
+    dt: float,
+    inv_mass: float,
+):
+    tid = wp.tid()
+
+    v_new = v[tid] + f[tid] * inv_mass * dt + gravity * dt
+    x_new = x[tid] + v_new * dt
+
+    v[tid] = v_new
+    x[tid] = x_new
+
+
+class Example:
+    def __init__(self, stage_path="example_dem.usd"):
+        fps = 60
+        self.frame_dt = 1.0 / fps
+
+        self.sim_substeps = 64
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_time = 0.0
+
+        self.point_radius = 0.1
+
+        self.k_contact = 8000.0
+        self.k_damp = 2.0
+        self.k_friction = 1.0
+        self.k_mu = 100000.0  # for cohesive materials
+
+        self.inv_mass = 64.0
+
+        self.grid = wp.HashGrid(128, 128, 128)
+        self.grid_cell_size = self.point_radius * 5.0
+
+        self.points = self.particle_grid(32, 64, 32, (0.0, 0.5, 0.0), self.point_radius, 0.1)
+
+        self.x = wp.array(self.points, dtype=wp.vec3)
+        self.v = wp.array(np.ones([len(self.x), 3]) * np.array([0.0, 0.0, 15.0]), dtype=wp.vec3)
+        self.f = wp.zeros_like(self.v)
+
+        if stage_path:
+            self.renderer = wp.render.UsdRenderer(stage_path)
+            self.renderer.render_ground()
+        else:
+            self.renderer = None
+
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.simulate()
+            self.graph = capture.graph
+
+    def simulate(self):
+        for _ in range(self.sim_substeps):
+            wp.launch(
+                kernel=apply_forces,
+                dim=len(self.x),
+                inputs=[
+                    self.grid.id,
+                    self.x,
+                    self.v,
+                    self.f,
+                    self.point_radius,
+                    self.k_contact,
+                    self.k_damp,
+                    self.k_friction,
+                    self.k_mu,
+                ],
+            )
+            wp.launch(
+                kernel=integrate,
+                dim=len(self.x),
+                inputs=[self.x, self.v, self.f, (0.0, -9.8, 0.0), self.sim_dt, self.inv_mass],
+            )
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            with wp.ScopedTimer("grid build", active=False):
+                self.grid.build(self.x, self.grid_cell_size)
+
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.simulate()
+
+            self.sim_time += self.frame_dt
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            self.renderer.begin_frame(self.sim_time)
+            self.renderer.render_points(
+                points=self.x.numpy(), radius=self.point_radius, name="points", colors=(0.8, 0.3, 0.2)
+            )
+            self.renderer.end_frame()
+
+    # creates a grid of particles
+    def particle_grid(self, dim_x, dim_y, dim_z, lower, radius, jitter):
+        rng = np.random.default_rng(42)
+        points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
+        points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
+        points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
+
+        return points_t.reshape((-1, 3))
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument(
+        "--stage_path",
+        type=lambda x: None if x == "None" else str(x),
+        default="example_dem.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--num_frames", type=int, default=200, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path)
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()
+
+        if example.renderer:
+            example.renderer.save()