warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/examples/benchmarks/benchmark_cloth_warp.py

@@ -0,0 +1,146 @@
+# 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 warp as wp
+
+wp.init()
+wp.build.clear_kernel_cache()
+
+
+@wp.kernel
+def eval_springs(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    spring_indices: wp.array(dtype=int),
+    spring_rest_lengths: wp.array(dtype=float),
+    spring_stiffness: wp.array(dtype=float),
+    spring_damping: wp.array(dtype=float),
+    f: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+
+    i = spring_indices[tid * 2 + 0]
+    j = spring_indices[tid * 2 + 1]
+
+    ke = spring_stiffness[tid]
+    kd = spring_damping[tid]
+    rest = spring_rest_lengths[tid]
+
+    xi = x[i]
+    xj = x[j]
+
+    vi = v[i]
+    vj = v[j]
+
+    xij = xi - xj
+    vij = vi - vj
+
+    l = wp.length(xij)
+    l_inv = 1.0 / l
+
+    # normalized spring direction
+    dir = xij * l_inv
+
+    c = l - rest
+    dcdt = wp.dot(dir, vij)
+
+    # damping based on relative velocity.
+    fs = dir * (ke * c + kd * dcdt)
+
+    wp.atomic_sub(f, i, fs)
+    wp.atomic_add(f, j, fs)
+
+
+@wp.kernel
+def integrate_particles(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    f: wp.array(dtype=wp.vec3),
+    w: wp.array(dtype=float),
+    dt: float,
+):
+    tid = wp.tid()
+
+    x0 = x[tid]
+    v0 = v[tid]
+    f0 = f[tid]
+    inv_mass = w[tid]
+
+    g = wp.vec3()
+
+    # treat particles with inv_mass == 0 as kinematic
+    if inv_mass > 0.0:
+        g = wp.vec3(0.0, 0.0 - 9.81, 0.0)
+
+    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
+    v1 = v0 + (f0 * inv_mass + g) * dt
+    x1 = x0 + v1 * dt
+
+    x[tid] = x1
+    v[tid] = v1
+
+    # clear forces
+    f[tid] = wp.vec3()
+
+
+class WpIntegrator:
+    def __init__(self, cloth, device):
+        self.device = wp.get_device(device)
+
+        with wp.ScopedDevice(self.device):
+            self.positions = wp.from_numpy(cloth.positions, dtype=wp.vec3)
+            self.positions_host = wp.from_numpy(cloth.positions, dtype=wp.vec3, device="cpu")
+            self.invmass = wp.from_numpy(cloth.inv_masses, dtype=float)
+
+            self.velocities = wp.zeros(cloth.num_particles, dtype=wp.vec3)
+            self.forces = wp.zeros(cloth.num_particles, dtype=wp.vec3)
+
+            self.spring_indices = wp.from_numpy(cloth.spring_indices, dtype=int)
+            self.spring_lengths = wp.from_numpy(cloth.spring_lengths, dtype=float)
+            self.spring_stiffness = wp.from_numpy(cloth.spring_stiffness, dtype=float)
+            self.spring_damping = wp.from_numpy(cloth.spring_damping, dtype=float)
+
+        self.cloth = cloth
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for s in range(substeps):
+            wp.launch(
+                kernel=eval_springs,
+                dim=self.cloth.num_springs,
+                inputs=[
+                    self.positions,
+                    self.velocities,
+                    self.spring_indices,
+                    self.spring_lengths,
+                    self.spring_stiffness,
+                    self.spring_damping,
+                    self.forces,
+                ],
+                outputs=[],
+                device=self.device,
+            )
+
+            # integrate
+            wp.launch(
+                kernel=integrate_particles,
+                dim=self.cloth.num_particles,
+                inputs=[self.positions, self.velocities, self.forces, self.invmass, sim_dt],
+                outputs=[],
+                device=self.device,
+            )
+
+        # copy data back to host
+        if self.device.is_cuda:
+            wp.copy(self.positions_host, self.positions)
+            wp.synchronize()
+
+            return self.positions_host.numpy()
+
+        else:
+            return self.positions.numpy()

warp/examples/benchmarks/benchmark_launches.py

@@ -0,0 +1,295 @@
+# 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.
+
+###########################################################################
+# Benchmarks for kernel launches with different types of args
+###########################################################################
+
+import warp as wp
+
+
+@wp.struct
+class S0:
+    pass
+
+
+@wp.struct
+class Sf:
+    x: float
+    y: float
+    z: float
+
+
+@wp.struct
+class Sv:
+    u: wp.vec3
+    v: wp.vec3
+    w: wp.vec3
+
+
+@wp.struct
+class Sm:
+    M: wp.mat33
+    N: wp.mat33
+    O: wp.mat33
+
+
+@wp.struct
+class Sa:
+    a: wp.array(dtype=float)
+    b: wp.array(dtype=float)
+    c: wp.array(dtype=float)
+
+
+@wp.struct
+class Sz:
+    a: wp.array(dtype=float)
+    b: wp.array(dtype=float)
+    c: wp.array(dtype=float)
+    x: float
+    y: float
+    z: float
+    u: wp.vec3
+    v: wp.vec3
+    w: wp.vec3
+
+
+@wp.kernel
+def k0():
+    tid = wp.tid()
+
+
+@wp.kernel
+def kf(x: float, y: float, z: float):
+    tid = wp.tid()
+
+
+@wp.kernel
+def kv(u: wp.vec3, v: wp.vec3, w: wp.vec3):
+    tid = wp.tid()
+
+
+@wp.kernel
+def km(M: wp.mat33, N: wp.mat33, O: wp.mat33):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ka(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float)):
+    tid = wp.tid()
+
+
+@wp.kernel
+def kz(
+    a: wp.array(dtype=float),
+    b: wp.array(dtype=float),
+    c: wp.array(dtype=float),
+    x: float,
+    y: float,
+    z: float,
+    u: wp.vec3,
+    v: wp.vec3,
+    w: wp.vec3,
+):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ks0(s: S0):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ksf(s: Sf):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ksv(s: Sv):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ksm(s: Sm):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ksa(s: Sa):
+    tid = wp.tid()
+
+
+@wp.kernel
+def ksz(s: Sz):
+    tid = wp.tid()
+
+
+wp.init()
+
+wp.build.clear_kernel_cache()
+
+devices = wp.get_devices()
+num_launches = 100000
+
+for device in devices:
+    with wp.ScopedDevice(device):
+        print(f"\n=================== Device '{device}' ===================")
+
+        wp.force_load(device)
+
+        n = 1
+        a = wp.zeros(n, dtype=float)
+        b = wp.zeros(n, dtype=float)
+        c = wp.zeros(n, dtype=float)
+        x = 17.0
+        y = 42.0
+        z = 99.0
+        u = wp.vec3(1, 2, 3)
+        v = wp.vec3(10, 20, 30)
+        w = wp.vec3(100, 200, 300)
+        M = wp.mat33()
+        N = wp.mat33()
+        O = wp.mat33()
+
+        s0 = S0()
+
+        sf = Sf()
+        sf.x = x
+        sf.y = y
+        sf.z = z
+
+        sv = Sv()
+        sv.u = u
+        sv.v = v
+        sv.w = w
+
+        sm = Sm()
+        sm.M = M
+        sm.N = N
+        sm.O = O
+
+        sa = Sa()
+        sa.a = a
+        sa.b = b
+        sa.c = c
+
+        sz = Sz()
+        sz.a = a
+        sz.b = b
+        sz.c = c
+        sz.x = x
+        sz.y = y
+        sz.z = z
+        sz.u = u
+        sz.v = v
+        sz.w = w
+
+        tk0 = wp.ScopedTimer("k0")
+        tkf = wp.ScopedTimer("kf")
+        tkv = wp.ScopedTimer("kv")
+        tkm = wp.ScopedTimer("km")
+        tka = wp.ScopedTimer("ka")
+        tkz = wp.ScopedTimer("kz")
+
+        ts0 = wp.ScopedTimer("s0")
+        tsf = wp.ScopedTimer("sf")
+        tsv = wp.ScopedTimer("sv")
+        tsm = wp.ScopedTimer("sm")
+        tsa = wp.ScopedTimer("sa")
+        tsz = wp.ScopedTimer("sz")
+
+        wp.synchronize_device()
+
+        with tk0:
+            for _ in range(num_launches):
+                wp.launch(k0, dim=1, inputs=[])
+
+        wp.synchronize_device()
+
+        with tkf:
+            for _ in range(num_launches):
+                wp.launch(kf, dim=1, inputs=[x, y, z])
+
+        wp.synchronize_device()
+
+        with tkv:
+            for _ in range(num_launches):
+                wp.launch(kv, dim=1, inputs=[u, v, w])
+
+        wp.synchronize_device()
+
+        with tkm:
+            for _ in range(num_launches):
+                wp.launch(km, dim=1, inputs=[M, N, O])
+
+        wp.synchronize_device()
+
+        with tka:
+            for _ in range(num_launches):
+                wp.launch(ka, dim=1, inputs=[a, b, c])
+
+        wp.synchronize_device()
+
+        with tkz:
+            for _ in range(num_launches):
+                wp.launch(kz, dim=1, inputs=[a, b, c, x, y, z, u, v, w])
+
+        # structs
+
+        wp.synchronize_device()
+
+        with ts0:
+            for _ in range(num_launches):
+                wp.launch(ks0, dim=1, inputs=[s0])
+
+        wp.synchronize_device()
+
+        with tsf:
+            for _ in range(num_launches):
+                wp.launch(ksf, dim=1, inputs=[sf])
+
+        wp.synchronize_device()
+
+        with tsv:
+            for _ in range(num_launches):
+                wp.launch(ksv, dim=1, inputs=[sv])
+
+        wp.synchronize_device()
+
+        with tsm:
+            for _ in range(num_launches):
+                wp.launch(ksm, dim=1, inputs=[sm])
+
+        wp.synchronize_device()
+
+        with tsa:
+            for _ in range(num_launches):
+                wp.launch(ksa, dim=1, inputs=[sa])
+
+        wp.synchronize_device()
+
+        with tsz:
+            for _ in range(num_launches):
+                wp.launch(ksz, dim=1, inputs=[sz])
+
+        wp.synchronize_device()
+
+        timers = [
+            [tk0, ts0],
+            [tkf, tsf],
+            [tkv, tsv],
+            [tkm, tsm],
+            [tka, tsa],
+            [tkz, tsz],
+        ]
+
+        print("--------------------------------")
+        print("| args |    direct |    struct |")
+        print("--------------------------------")
+        for tk, ts in timers:
+            print(f"|  {tk.name}  |{tk.elapsed:10.0f} |{ts.elapsed:10.0f} |")
+        print("--------------------------------")

warp/examples/core/example_dem.py

@@ -0,0 +1,221 @@
+# 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.
+
+###########################################################################
+# 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 os
+
+import numpy as np
+
+import warp as wp
+import warp.render
+
+wp.init()
+
+
+@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):
+        self.frame_dt = 1.0 / 60
+        self.frame_count = 400
+
+        self.sim_substeps = 64
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_steps = self.frame_count * 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, 128, 32, (0.0, 0.3, 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, 10.0]), dtype=wp.vec3)
+        self.f = wp.zeros_like(self.v)
+
+        self.renderer = None
+        if stage is not None:
+            self.renderer = wp.render.UsdRenderer(stage)
+            self.renderer.render_ground()
+
+        self.use_graph = wp.get_device().is_cuda
+        if self.use_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", active=True):
+            with wp.ScopedTimer("grid build", active=False):
+                self.grid.build(self.x, self.grid_cell_size)
+
+            if self.use_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", active=True):
+            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),) * len(self.x))
+            self.renderer.end_frame()
+
+    # creates a grid of particles
+    def particle_grid(self, dim_x, dim_y, dim_z, lower, radius, jitter):
+        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 + np.random.rand(*points_t.shape) * radius * jitter
+
+        return points_t.reshape((-1, 3))
+
+
+if __name__ == "__main__":
+    stage_path = os.path.join(os.path.dirname(__file__), "example_dem.usd")
+
+    example = Example(stage_path)
+
+    for i in range(example.frame_count):
+        example.step()
+        example.render()
+
+    if example.renderer:
+        example.renderer.save()