warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/examples/benchmarks/benchmark_api.py

@@ -0,0 +1,389 @@
+# 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 gc
+import statistics as stats
+
+import warp as wp
+
+ENABLE_MEMPOOLS = False
+ENABLE_PEER_ACCESS = False
+ENABLE_MEMPOOL_ACCESS = False
+ENABLE_MEMPOOL_RELEASE_THRESHOLD = False
+
+MEMPOOL_RELEASE_THRESHOLD = 1024 * 1024 * 1024
+
+DO_SYNC = False
+VERBOSE = False
+USE_NVTX = False
+
+num_elems = 10000
+num_runs = 10000
+trim_runs = 2500
+
+
+@wp.kernel
+def inc_kernel(a: wp.array(dtype=float)):
+    tid = wp.tid()
+    a[tid] = a[tid] + 1.0
+
+
+# configure devices
+for target_device in wp.get_cuda_devices():
+    try:
+        wp.set_mempool_enabled(target_device, ENABLE_MEMPOOLS)
+        if ENABLE_MEMPOOL_RELEASE_THRESHOLD:
+            wp.set_mempool_release_threshold(target_device, MEMPOOL_RELEASE_THRESHOLD)
+    except Exception as e:
+        print(f"Error: {e}")
+
+    for peer_device in wp.get_cuda_devices():
+        try:
+            wp.set_peer_access_enabled(target_device, peer_device, ENABLE_PEER_ACCESS)
+        except Exception as e:
+            print(f"Error: {e}")
+
+        try:
+            wp.set_mempool_access_enabled(target_device, peer_device, ENABLE_MEMPOOL_ACCESS)
+        except Exception as e:
+            print(f"Error: {e}")
+
+cuda_device_count = wp.get_cuda_device_count()
+
+cuda0 = wp.get_device("cuda:0")
+
+# preallocate some arrays
+arr_host = wp.zeros(num_elems, dtype=float, device="cpu", pinned=False)
+arr_host_pinned = wp.zeros(num_elems, dtype=float, device="cpu", pinned=True)
+arr_cuda0 = wp.zeros(num_elems, dtype=float, device=cuda0)
+arr_cuda0_src = wp.zeros(num_elems, dtype=float, device=cuda0)
+arr_cuda0_dst = wp.zeros(num_elems, dtype=float, device=cuda0)
+
+# mgpu support
+if cuda_device_count > 1:
+    cuda1 = wp.get_device("cuda:1")
+    arr_cuda1 = wp.zeros(num_elems, dtype=float, device=cuda1)
+
+stream0 = wp.Stream(cuda0)
+
+# preload module
+wp.force_load(cuda0)
+if cuda_device_count > 1:
+    wp.force_load(cuda1)
+
+# capture graph
+with wp.ScopedDevice(cuda0):
+    wp.capture_begin()
+    wp.launch(inc_kernel, dim=arr_cuda0.size, inputs=[arr_cuda0])
+    graph0 = wp.capture_end()
+
+
+g_allocs = [None] * num_runs
+
+
+def test_alloc(num_elems, device, idx):
+    wp.synchronize()
+
+    with wp.ScopedTimer("alloc", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        g_allocs[idx] = wp.empty(num_elems, dtype=float, device=device)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_free(device, idx):
+    wp.synchronize()
+
+    with wp.ScopedTimer("free", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        g_allocs[idx] = None
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_zeros(num_elems, device, idx):
+    wp.synchronize()
+
+    with wp.ScopedTimer("zeros", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        g_allocs[idx] = wp.zeros(num_elems, dtype=float, device=device)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_h2d(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("h2d", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_cuda0, arr_host)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_d2h(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("d2h", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_host, arr_cuda0)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_h2d_pinned(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("h2d pinned", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_cuda0, arr_host_pinned)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_d2h_pinned(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("d2h pinned", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_host_pinned, arr_cuda0)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_d2d(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("d2d", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_cuda0_dst, arr_cuda0_src)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_p2p(num_elems, src_device, dst_device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("p2p", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_cuda0, arr_cuda1)
+
+        if DO_SYNC:
+            wp.synchronize_device(src_device)
+            wp.synchronize_device(dst_device)
+
+    return timer.elapsed
+
+
+def test_p2p_stream(num_elems, src_device, dst_device):
+    stream = stream0
+
+    wp.synchronize()
+
+    with wp.ScopedTimer("p2p stream", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.copy(arr_cuda0, arr_cuda1, stream=stream)
+
+        if DO_SYNC:
+            wp.synchronize_device(src_device)
+            wp.synchronize_device(dst_device)
+
+    return timer.elapsed
+
+
+def test_launch(num_elems, device):
+    a = arr_cuda0
+
+    wp.synchronize()
+
+    with wp.ScopedTimer("launch", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.launch(inc_kernel, dim=a.size, inputs=[a], device=device)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_launch_stream(num_elems, device):
+    a = arr_cuda0
+    stream = stream0
+
+    wp.synchronize()
+
+    with wp.ScopedTimer("launch stream", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.launch(inc_kernel, dim=a.size, inputs=[a], stream=stream)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_graph(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("graph", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.capture_launch(graph0)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+def test_graph_stream(num_elems, device):
+    wp.synchronize()
+
+    with wp.ScopedTimer("graph", print=VERBOSE, use_nvtx=USE_NVTX) as timer:
+        wp.capture_launch(graph0, stream=stream0)
+
+        if DO_SYNC:
+            wp.synchronize_device(device)
+
+    return timer.elapsed
+
+
+alloc_times = [0] * num_runs
+free_times = [0] * num_runs
+zeros_times = [0] * num_runs
+d2h_times = [0] * num_runs
+h2d_times = [0] * num_runs
+d2h_pinned_times = [0] * num_runs
+h2d_pinned_times = [0] * num_runs
+d2d_times = [0] * num_runs
+p2p_times = [0] * num_runs
+p2p_stream_times = [0] * num_runs
+launch_times = [0] * num_runs
+launch_stream_times = [0] * num_runs
+graph_times = [0] * num_runs
+graph_stream_times = [0] * num_runs
+
+wp.set_device(cuda0)
+
+# alloc
+for i in range(num_runs):
+    gc.disable()
+    alloc_times[i] = test_alloc(num_elems, cuda0, i)
+    gc.enable()
+
+# free
+for i in range(num_runs):
+    gc.disable()
+    free_times[i] = test_free(cuda0, i)
+    gc.enable()
+
+# zeros
+for i in range(num_runs):
+    gc.disable()
+    zeros_times[i] = test_zeros(num_elems, cuda0, i)
+    gc.enable()
+
+# free zeros
+for i in range(num_runs):
+    g_allocs[i] = None
+
+# h2d, d2h pageable copy
+for i in range(num_runs):
+    gc.disable()
+    h2d_times[i] = test_h2d(num_elems, cuda0)
+    d2h_times[i] = test_d2h(num_elems, cuda0)
+    gc.enable()
+
+# h2d, d2h pinned copy
+for i in range(num_runs):
+    gc.disable()
+    h2d_pinned_times[i] = test_h2d_pinned(num_elems, cuda0)
+    d2h_pinned_times[i] = test_d2h_pinned(num_elems, cuda0)
+    gc.enable()
+
+# d2d copy
+for i in range(num_runs):
+    gc.disable()
+    d2d_times[i] = test_d2d(num_elems, cuda0)
+    gc.enable()
+
+# p2p copy
+if cuda_device_count > 1:
+    for i in range(num_runs):
+        gc.disable()
+        p2p_times[i] = test_p2p(num_elems, cuda1, cuda0)
+        p2p_stream_times[i] = test_p2p_stream(num_elems, cuda1, cuda0)
+        gc.enable()
+
+# launch
+for i in range(num_runs):
+    gc.disable()
+    launch_times[i] = test_launch(num_elems, cuda0)
+    launch_stream_times[i] = test_launch_stream(num_elems, cuda0)
+    gc.enable()
+
+# graph
+for i in range(num_runs):
+    gc.disable()
+    graph_times[i] = test_graph(num_elems, cuda0)
+    graph_stream_times[i] = test_graph_stream(num_elems, cuda0)
+    gc.enable()
+
+
+def print_stat(name, data, trim=trim_runs):
+    assert len(data) - 2 * trim > 0
+    if trim > 0:
+        data = sorted(data)[trim:-trim]
+    print(f"{name:15s} {1000000 * stats.mean(data):.0f}")
+
+
+print("=========================")
+print_stat("Alloc", alloc_times)
+print_stat("Free", free_times)
+print_stat("Zeros", zeros_times)
+print_stat("H2D", h2d_times)
+print_stat("D2H", d2h_times)
+print_stat("H2D pinned", h2d_pinned_times)
+print_stat("D2H pinned", d2h_pinned_times)
+print_stat("D2D", d2d_times)
+print_stat("P2P", p2p_times)
+print_stat("P2P stream", p2p_stream_times)
+print_stat("Launch", launch_times)
+print_stat("Launch stream", launch_stream_times)
+print_stat("Graph", graph_times)
+print_stat("Graph stream", graph_stream_times)
+
+
+# ========= profiling ==========
+
+# from pyinstrument import Profiler
+# profiler = Profiler()
+# profiler.start()
+# for i in range(10):
+#     # test_alloc(num_elems, cuda0)
+#     # test_h2d(num_elems, cuda0)
+#     test_p2p(num_elems, cuda0, cuda1)
+# profiler.stop()
+# print(profiler.output_text(show_all=True))

warp/examples/benchmarks/benchmark_cloth.py

@@ -0,0 +1,296 @@
+# 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.
+
+# include parent path
+import csv
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".")))
+
+from pxr import Usd, UsdGeom
+
+import warp as wp
+
+
+class Cloth:
+    def __init__(
+        self, lower, dx, dy, radius, stretch_stiffness, bend_stiffness, shear_stiffness, mass, fix_corners=True
+    ):
+        self.triangles = []
+
+        self.positions = []
+        self.velocities = []
+        self.inv_masses = []
+
+        self.spring_indices = []
+        self.spring_lengths = []
+        self.spring_stiffness = []
+        self.spring_damping = []
+
+        def grid(x, y, stride):
+            return y * stride + x
+
+        def create_spring(i, j, stiffness, damp=10.0):
+            length = np.linalg.norm(np.array(self.positions[i]) - np.array(self.positions[j]))
+
+            self.spring_indices.append(i)
+            self.spring_indices.append(j)
+            self.spring_lengths.append(length)
+            self.spring_stiffness.append(stiffness)
+            self.spring_damping.append(damp)
+
+        for y in range(dy):
+            for x in range(dx):
+                p = np.array(lower) + radius * np.array((float(x), float(0.0), float(y)))
+
+                self.positions.append(p)
+                self.velocities.append(np.zeros(3))
+
+                if x > 0 and y > 0:
+                    self.triangles.append(grid(x - 1, y - 1, dx))
+                    self.triangles.append(grid(x, y - 1, dx))
+                    self.triangles.append(grid(x, y, dx))
+
+                    self.triangles.append(grid(x - 1, y - 1, dx))
+                    self.triangles.append(grid(x, y, dx))
+                    self.triangles.append(grid(x - 1, y, dx))
+
+                if fix_corners and y == 0 and (x == 0 or x == dx - 1):
+                    w = 0.0
+                else:
+                    w = 1.0 / mass
+
+                self.inv_masses.append(w)
+
+        # horizontal springs
+        for y in range(dy):
+            for x in range(dx):
+                index0 = y * dx + x
+
+                if x > 0:
+                    index1 = y * dx + x - 1
+                    create_spring(index0, index1, stretch_stiffness)
+
+                if x > 1 and bend_stiffness > 0.0:
+                    index2 = y * dx + x - 2
+                    create_spring(index0, index2, bend_stiffness)
+
+                if y > 0 and x < dx - 1 and shear_stiffness > 0.0:
+                    indexDiag = (y - 1) * dx + x + 1
+                    create_spring(index0, indexDiag, shear_stiffness)
+
+                if y > 0 and x > 0 and shear_stiffness > 0.0:
+                    indexDiag = (y - 1) * dx + x - 1
+                    create_spring(index0, indexDiag, shear_stiffness)
+
+        # vertical
+        for x in range(dx):
+            for y in range(dy):
+                index0 = y * dx + x
+
+                if y > 0:
+                    index1 = (y - 1) * dx + x
+                    create_spring(index0, index1, stretch_stiffness)
+
+                if y > 1 and bend_stiffness > 0.0:
+                    index2 = (y - 2) * dx + x
+                    create_spring(index0, index2, bend_stiffness)
+
+        # harden to np arrays
+        self.positions = np.array(self.positions, dtype=np.float32)
+        self.velocities = np.array(self.velocities, dtype=np.float32)
+        self.inv_masses = np.array(self.inv_masses, dtype=np.float32)
+        self.spring_lengths = np.array(self.spring_lengths, dtype=np.float32)
+        self.spring_indices = np.array(self.spring_indices, dtype=np.int32)
+        self.spring_stiffness = np.array(self.spring_stiffness, dtype=np.float32)
+        self.spring_damping = np.array(self.spring_damping, dtype=np.float32)
+
+        self.num_particles = len(self.positions)
+        self.num_springs = len(self.spring_lengths)
+        self.num_tris = int(len(self.triangles) / 3)
+
+
+def run_benchmark(mode, dim, timers, render=False):
+    # params
+    sim_width = dim
+    sim_height = dim
+
+    sim_fps = 60.0
+    sim_substeps = 16
+    sim_duration = 1.0
+    sim_frames = int(sim_duration * sim_fps)
+    sim_dt = 1.0 / sim_fps
+    sim_time = 0.0
+
+    # wave constants
+    k_stretch = 1000.0
+    k_shear = 1000.0
+    k_bend = 1000.0
+    # k_damp = 0.0
+
+    cloth = Cloth(
+        lower=(0.0, 0.0, 0.0),
+        dx=sim_width,
+        dy=sim_height,
+        radius=0.1,
+        stretch_stiffness=k_stretch,
+        bend_stiffness=k_bend,
+        shear_stiffness=k_shear,
+        mass=0.1,
+        fix_corners=True,
+    )
+
+    if render:
+        # set up grid for visualization
+        stage = Usd.Stage.CreateNew("benchmark.usd")
+        stage.SetStartTimeCode(0.0)
+        stage.SetEndTimeCode(sim_duration * sim_fps)
+        stage.SetTimeCodesPerSecond(sim_fps)
+
+        grid = UsdGeom.Mesh.Define(stage, "/root")
+        grid.GetPointsAttr().Set(cloth.positions, 0.0)
+        grid.GetFaceVertexIndicesAttr().Set(cloth.triangles, 0.0)
+        grid.GetFaceVertexCountsAttr().Set([3] * cloth.num_tris, 0.0)
+
+    with wp.ScopedTimer("Initialization", dict=timers):
+        if mode == "warp_cpu":
+            import benchmark_cloth_warp
+
+            integrator = benchmark_cloth_warp.WpIntegrator(cloth, "cpu")
+
+        elif mode == "warp_gpu":
+            import benchmark_cloth_warp
+
+            integrator = benchmark_cloth_warp.WpIntegrator(cloth, "cuda")
+
+        elif mode == "taichi_cpu":
+            import benchmark_cloth_taichi
+
+            integrator = benchmark_cloth_taichi.TiIntegrator(cloth, "cpu")
+
+        elif mode == "taichi_gpu":
+            import benchmark_cloth_taichi
+
+            integrator = benchmark_cloth_taichi.TiIntegrator(cloth, "cuda")
+
+        elif mode == "numpy":
+            import benchmark_cloth_numpy
+
+            integrator = benchmark_cloth_numpy.NpIntegrator(cloth)
+
+        elif mode == "cupy":
+            import benchmark_cloth_cupy
+
+            integrator = benchmark_cloth_cupy.CpIntegrator(cloth)
+
+        elif mode == "numba":
+            import benchmark_cloth_numba
+
+            integrator = benchmark_cloth_numba.NbIntegrator(cloth)
+
+        elif mode == "torch_cpu":
+            import benchmark_cloth_pytorch
+
+            integrator = benchmark_cloth_pytorch.TrIntegrator(cloth, "cpu")
+
+        elif mode == "torch_gpu":
+            import benchmark_cloth_pytorch
+
+            integrator = benchmark_cloth_pytorch.TrIntegrator(cloth, "cuda")
+
+        elif mode == "jax_cpu":
+            os.environ["JAX_PLATFORM_NAME"] = "cpu"
+
+            import benchmark_cloth_jax
+
+            integrator = benchmark_cloth_jax.JxIntegrator(cloth)
+
+        elif mode == "jax_gpu":
+            os.environ["JAX_PLATFORM_NAME"] = "gpu"
+
+            import benchmark_cloth_jax
+
+            integrator = benchmark_cloth_jax.JxIntegrator(cloth)
+
+        elif mode == "paddle_cpu":
+            import benchmark_cloth_paddle
+
+            integrator = benchmark_cloth_paddle.TrIntegrator(cloth, "cpu")
+
+        elif mode == "paddle_gpu":
+            import benchmark_cloth_paddle
+
+            integrator = benchmark_cloth_paddle.TrIntegrator(cloth, "gpu")
+
+        else:
+            raise RuntimeError("Unknown simulation backend")
+
+            # run one warm-up iteration to accurately measure initialization time (some engines do lazy init)
+            positions = integrator.simulate(sim_dt, sim_substeps)
+
+    label = "Dim ({}^2)".format(dim)
+
+    # run simulation
+    for _i in range(sim_frames):
+        # simulate
+        with wp.ScopedTimer(label, dict=timers):
+            positions = integrator.simulate(sim_dt, sim_substeps)
+
+        if render:
+            grid.GetPointsAttr().Set(positions, sim_time * sim_fps)
+
+        sim_time += sim_dt
+
+    if render:
+        stage.Save()
+
+
+# record profiling information
+timers = {}
+
+if len(sys.argv) > 1:
+    mode = sys.argv[1]
+else:
+    mode = "warp_gpu"
+
+run_benchmark(mode, 32, timers, render=False)
+run_benchmark(mode, 64, timers, render=False)
+run_benchmark(mode, 128, timers, render=False)
+
+# write results
+
+for k, v in timers.items():
+    print("{:16} min: {:8.2f} max: {:8.2f} avg: {:8.2f}".format(k, np.min(v), np.max(v), np.mean(v)))
+
+report = open(os.path.join("benchmark.csv"), "a")
+writer = csv.writer(report, delimiter=",")
+
+if report.tell() == 0:
+    writer.writerow(["Name", "Init", "Dim (32^2)", "Dim (64^2)", "Dim (128^2)"])
+
+writer.writerow(
+    [
+        mode,
+        np.max(timers["Initialization"]),
+        np.mean(timers["Dim (32^2)"]),
+        np.mean(timers["Dim (64^2)"]),
+        np.mean(timers["Dim (128^2)"]),
+    ]
+)
+
+report.close()