warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

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.SetFramesPerSecond(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 = f"Dim ({dim}^2)"
+
+    # 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(f"{k:16} min: {np.min(v):8.2f} max: {np.max(v):8.2f} avg: {np.mean(v):8.2f}")
+
+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()

warp/examples/benchmarks/benchmark_cloth_cupy.py

@@ -0,0 +1,96 @@
+# 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 cupy as cp
+import cupyx as cpx
+
+
+def eval_springs(x, v, indices, rest, ke, kd, f):
+    i = indices[:, 0]
+    j = indices[:, 1]
+
+    xi = x[i]
+    xj = x[j]
+
+    vi = v[i]
+    vj = v[j]
+
+    xij = xi - xj
+    vij = vi - vj
+
+    l = cp.linalg.norm(xij, axis=1)
+    l_inv = 1.0 / l
+
+    # normalized spring direction
+    dir = (xij.T * l_inv).T
+
+    c = l - rest
+    dcdt = cp.sum(dir * vij, axis=1)
+
+    # damping based on relative velocity.
+    fs = dir.T * (ke * c + kd * dcdt)
+
+    cpx.scatter_add(f, i, -fs.T)
+    cpx.scatter_add(f, j, fs.T)
+
+
+def integrate_particles(x, v, f, w, dt):
+    g = cp.array((0.0, 0.0 - 9.8, 0.0))
+    s = w > 0.0
+
+    a_ext = g * s[:, None]
+
+    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
+    v += ((f.T * w).T + a_ext) * dt
+    x += v * dt
+
+    # clear forces
+    f *= 0.0
+
+
+class CpIntegrator:
+    def __init__(self, cloth):
+        self.cloth = cloth
+
+        self.positions = cp.array(self.cloth.positions)
+        self.velocities = cp.array(self.cloth.velocities)
+        self.inv_mass = cp.array(self.cloth.inv_masses)
+
+        self.spring_indices = cp.array(self.cloth.spring_indices)
+        self.spring_lengths = cp.array(self.cloth.spring_lengths)
+        self.spring_stiffness = cp.array(self.cloth.spring_stiffness)
+        self.spring_damping = cp.array(self.cloth.spring_damping)
+
+        self.forces = cp.zeros((self.cloth.num_particles, 3), dtype=cp.float32)
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for _s in range(substeps):
+            eval_springs(
+                self.positions,
+                self.velocities,
+                self.spring_indices.reshape((self.cloth.num_springs, 2)),
+                self.spring_lengths,
+                self.spring_stiffness,
+                self.spring_damping,
+                self.forces,
+            )
+
+            # integrate
+            integrate_particles(self.positions, self.velocities, self.forces, self.inv_mass, sim_dt)
+
+        # return np.array(self.positions)
+        return self.positions.get()

warp/examples/benchmarks/benchmark_cloth_jax.py

@@ -0,0 +1,105 @@
+# 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 jax.lax
+import jax.numpy as jnp
+import numpy as np
+
+
+@jax.jit
+def eval_springs(x, v, indices, rest, ke, kd):
+    i = indices[:, 0]
+    j = indices[:, 1]
+
+    xi = x[i]
+    xj = x[j]
+
+    vi = v[i]
+    vj = v[j]
+
+    xij = xi - xj
+    vij = vi - vj
+
+    l = jnp.linalg.norm(xij, axis=1)
+    l_inv = 1.0 / l
+
+    # normalized spring direction
+    dir = (xij.T * l_inv).T
+
+    c = l - rest
+    dcdt = jnp.sum(dir * vij, axis=1)
+
+    # damping based on relative velocity.
+    fs = dir.T * (ke * c + kd * dcdt)
+
+    f = jnp.zeros_like(v)
+
+    # f = jax.ops.index_add(f, i, -fs.T, indices_are_sorted=False, unique_indices=False)
+    # f = jax.ops.index_add(f, j, fs.T, indices_are_sorted=False, unique_indices=False)
+
+    f.at[i].add(-fs.T)
+    f.at[j].add(fs.T)
+
+    return f
+
+
+@jax.jit
+def integrate_particles(x, v, f, w, dt):
+    g = jnp.array((0.0, 0.0 - 9.8, 0.0))
+    s = w > 0.0
+
+    a_ext = g * s[:, None]
+
+    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
+    v += ((f.T * w).T + a_ext) * dt
+    x += v * dt
+
+    return (x, v)
+
+
+class JxIntegrator:
+    def __init__(self, cloth):
+        self.cloth = cloth
+
+        self.positions = jnp.array(self.cloth.positions)
+        self.velocities = jnp.array(self.cloth.velocities)
+        self.inv_mass = jnp.array(self.cloth.inv_masses)
+
+        print(self.positions.device_buffer.device())
+
+        self.spring_indices = jnp.array(self.cloth.spring_indices)
+        self.spring_lengths = jnp.array(self.cloth.spring_lengths)
+        self.spring_stiffness = jnp.array(self.cloth.spring_stiffness)
+        self.spring_damping = jnp.array(self.cloth.spring_damping)
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for _s in range(substeps):
+            f = eval_springs(
+                self.positions,
+                self.velocities,
+                self.spring_indices.reshape((self.cloth.num_springs, 2)),
+                self.spring_lengths,
+                self.spring_stiffness,
+                self.spring_damping,
+            )
+
+            # integrate
+            (self.positions, self.velocities) = integrate_particles(
+                self.positions, self.velocities, f, self.inv_mass, sim_dt
+            )
+
+        return np.array(self.positions)

warp/examples/benchmarks/benchmark_cloth_numba.py

@@ -0,0 +1,161 @@
+# 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 math
+
+import cupy as cp
+import numpy as np
+from numba import cuda, float32
+
+# Notes:
+#
+# Current implementation requires some familarity of writing custom cuda kernels
+# May be improved with cuda ufuncs and/or writing custom numba type extensions.
+
+
+@cuda.jit(device=True)
+def norm(x):
+    s = float32(0.0)
+    for i in range(3):
+        s += x[i] * x[i]
+    return math.sqrt(s)
+
+
+@cuda.jit(device=True)
+def dot(x, y):
+    s = float32(0.0)
+    for i in range(3):
+        s += x[i] * y[i]
+    return s
+
+
+@cuda.jit
+def eval_springs_cuda(
+    num_springs,  # (1,)
+    xs,  # position              (N, 3)
+    vs,  # velocities            (N, 3)
+    indices,  # spring indices        (S, 2)
+    rests,  # spring rest length    (S,)
+    kes,  # stiffness             (S,)
+    kds,  # damping               (S,)
+    fs,
+):  # forces                (N, 3)
+    tidx = cuda.grid(1)
+
+    if tidx < num_springs:
+        i, j = indices[tidx][0], indices[tidx][1]
+        xi, xj = xs[i], xs[j]
+        vi, vj = vs[i], vs[j]
+        rest, ke, kd = rests[tidx], kes[tidx], kds[tidx]
+
+        xij = cuda.local.array(3, dtype=cp.float32)
+        vij = cuda.local.array(3, dtype=cp.float32)
+        for k in range(3):
+            xij[k] = xi[k] - xj[k]
+        for k in range(3):
+            vij[k] = vi[k] - vj[k]
+
+        l = norm(xij)
+
+        l_inv = float32(1.0) / l
+
+        # normalized spring direction
+        xij_unit = cuda.local.array(3, dtype=cp.float32)
+        for k in range(3):
+            xij_unit[k] = xij[k] * l_inv
+        c = l - rest
+        dcdt = dot(xij_unit, vij)
+
+        # mass-spring-damper model
+        fac = ke * c + kd * dcdt
+        df = cuda.local.array(3, dtype=cp.float32)
+        for k in range(3):
+            df[k] = xij_unit[k] * fac
+
+        for k in range(3):
+            cuda.atomic.add(fs[i], k, -df[k])
+            cuda.atomic.add(fs[j], k, df[k])
+
+
+# Support const array with cp array?
+g = np.array([0.0, 0.0 - 9.8, 0.0], dtype=np.float32)
+z = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+
+
+@cuda.jit
+def integrate_particles_cuda(
+    xs,  # position  (N, 3)
+    vs,  # velocity  (N, 3)
+    fs,  # force     (N, 3)
+    ws,  # inverse of mass (N,)
+    dt,
+):  # dt        (1,)
+    i = cuda.grid(1)
+
+    if i < xs.shape[0]:
+        w = ws[i]
+        a = cuda.const.array_like(g) if w > 0.0 else cuda.const.array_like(z)
+
+        for j in range(3):
+            # vs[i] += ((f * w) + a) * dt (ideally)
+            vs[i][j] = vs[i][j] + ((fs[i][j] * w) + a[j]) * dt
+            xs[i][j] = xs[i][j] + vs[i][j] * dt
+
+        fs[i] = 0.0
+
+
+class NbIntegrator:
+    def __init__(self, cloth):
+        self.cloth = cloth
+
+        self.positions = cp.array(self.cloth.positions)
+        self.velocities = cp.array(self.cloth.velocities)
+        self.inv_mass = cp.array(self.cloth.inv_masses)
+
+        self.spring_indices = cp.array(self.cloth.spring_indices)
+        self.spring_lengths = cp.array(self.cloth.spring_lengths)
+        self.spring_stiffness = cp.array(self.cloth.spring_stiffness)
+        self.spring_damping = cp.array(self.cloth.spring_damping)
+
+        self.forces = cp.zeros((self.cloth.num_particles, 3), dtype=cp.float32)
+
+        self.num_particles = self.positions.shape[0]
+        self.integrate_tpb = 4
+        self.integrate_nb = self.num_particles // self.integrate_tpb + 1
+
+        self.spring_tpb = 4
+        self.spring_nb = self.cloth.num_springs // self.spring_tpb + 1
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for _s in range(substeps):
+            eval_springs_cuda[self.spring_nb, self.spring_tpb](
+                self.cloth.num_springs,
+                self.positions,
+                self.velocities,
+                self.spring_indices.reshape((self.cloth.num_springs, 2)),
+                self.spring_lengths,
+                self.spring_stiffness,
+                self.spring_damping,
+                self.forces,
+            )
+
+            # integrate
+            integrate_particles_cuda[self.integrate_nb, self.integrate_tpb](
+                self.positions, self.velocities, self.forces, self.inv_mass, sim_dt
+            )
+
+        return self.positions.get()