warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.0__py3-none-win_amd64.whl

warp/examples/benchmarks/benchmark_cloth_cupy.py

@@ -1,88 +1,88 @@
-# 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 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()
+# 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 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

@@ -1,100 +1,97 @@
-# 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 numpy as np
-
-import jax.lax
-import jax.numpy as jnp
-
-import os
-
-
-@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)
+# 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 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

@@ -1,142 +1,146 @@
-from numba import cuda, float32
-import numpy as np
-import cupy as cp
-
-import math
-
-# 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, vs, fs, ws, dt  # position  (N, 3)  # velocity  (N, 3)  # force     (N, 3)  # inverse of mass (N,)
-):  # 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()
+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()