warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/examples/interop/example_jax_callable.py

@@ -0,0 +1,116 @@
+# 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.
+
+###########################################################################
+# Example jax_callable()
+#
+# Examples of calling annotated Python functions from JAX.
+###########################################################################
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import warp as wp
+from warp.jax_experimental.ffi import jax_callable
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+# The Python function to call.
+# Note the argument annotations, just like Warp kernels.
+def example_func(
+    # inputs
+    a: wp.array(dtype=float),
+    b: wp.array(dtype=wp.vec2),
+    s: float,
+    # outputs
+    c: wp.array(dtype=float),
+    d: wp.array(dtype=wp.vec2),
+):
+    wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+    wp.launch(scale_vec_kernel, dim=b.shape, inputs=[b, s], outputs=[d])
+
+
+def example1():
+    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+
+    @jax.jit
+    def f():
+        # inputs
+        a = jnp.arange(10, dtype=jnp.float32)
+        b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+        s = 2.0
+
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    r1, r2 = f()
+    print(r1)
+    print(r2)
+
+
+def example2():
+    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, b, s):
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    # inputs
+    a = jnp.arange(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+    s = 3.0
+
+    r1, r2 = f(a, b, s)
+    print(r1)
+    print(r2)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()

warp/examples/interop/example_jax_ffi_callback.py

@@ -0,0 +1,132 @@
+# 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.
+
+###########################################################################
+# Example register_ffi_callback()
+#
+# Examples of calling Python functions from JAX.
+# Target functions must have the form func(inputs, outputs, attrs, ctx).
+###########################################################################
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import warp as wp
+from warp.jax import get_jax_device
+from warp.jax_experimental.ffi import register_ffi_callback
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+def example1():
+    # the Python function to call
+    def print_args(inputs, outputs, attrs, ctx):
+        def buffer_to_string(b):
+            return str(b.dtype) + str(list(b.shape)) + " @%x" % b.data
+
+        print("Inputs:     ", ", ".join([buffer_to_string(b) for b in inputs]))
+        print("Outputs:    ", ", ".join([buffer_to_string(b) for b in outputs]))
+        print("Attributes: ", "".join(["\n  %s: %s" % (k, str(v)) for k, v in attrs.items()]))
+
+    # register callback
+    register_ffi_callback("print_args", print_args)
+
+    # set up call
+    call = jax.ffi.ffi_call("print_args", jax.ShapeDtypeStruct((1, 2, 3), jnp.int8))
+
+    # call it
+    call(
+        jnp.arange(16),
+        jnp.arange(32.0).reshape((4, 8)),
+        str_attr="hi",
+        f32_attr=np.float32(4.2),
+        dict_attr={"a": 1, "b": 6.4},
+    )
+
+
+def example2():
+    # the Python function to call
+    def warp_func(inputs, outputs, attrs, ctx):
+        # input arrays
+        a = inputs[0]
+        b = inputs[1]
+
+        # scalar attributes
+        s = attrs["scale"]
+
+        # output arrays
+        c = outputs[0]
+        d = outputs[1]
+
+        device = wp.device_from_jax(get_jax_device())
+        stream = wp.Stream(device, cuda_stream=ctx.stream)
+
+        with wp.ScopedStream(stream):
+            # launch with arrays of scalars
+            wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+
+            # launch with arrays of vec2
+            # NOTE: the input shapes are from JAX arrays, we need to strip the inner dimension for vec2 arrays
+            wp.launch(scale_vec_kernel, dim=b.shape[0], inputs=[b, s], outputs=[d])
+
+    # register callback
+    register_ffi_callback("warp_func", warp_func)
+
+    n = 10
+
+    # inputs
+    a = jnp.arange(n, dtype=jnp.float32)
+    b = jnp.arange(n, dtype=jnp.float32).reshape((n // 2, 2))  # array of wp.vec2
+    s = 2.0
+
+    # set up call
+    out_types = [
+        jax.ShapeDtypeStruct(a.shape, jnp.float32),
+        jax.ShapeDtypeStruct(b.shape, jnp.float32),  # array of wp.vec2
+    ]
+    call = jax.ffi.ffi_call("warp_func", out_types)
+
+    # call it
+    c, d = call(a, b, scale=s)
+
+    print(c)
+    print(d)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()

warp/examples/interop/example_jax_kernel.py

@@ -0,0 +1,205 @@
+# 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.
+
+###########################################################################
+# Example jax_kernel()
+#
+# Examples of calling a Warp kernel from JAX.
+###########################################################################
+
+import math
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import warp as wp
+from warp.jax_experimental.ffi import jax_kernel
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=int), b: wp.array(dtype=int), output: wp.array(dtype=int)):
+    tid = wp.tid()
+    output[tid] = a[tid] + b[tid]
+
+
+@wp.kernel
+def sincos_kernel(angle: wp.array(dtype=float), sin_out: wp.array(dtype=float), cos_out: wp.array(dtype=float)):
+    tid = wp.tid()
+    sin_out[tid] = wp.sin(angle[tid])
+    cos_out[tid] = wp.cos(angle[tid])
+
+
+@wp.kernel
+def diagonal_kernel(output: wp.array(dtype=wp.mat33)):
+    tid = wp.tid()
+    output[tid] = wp.mat33(1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0)
+
+
+@wp.kernel
+def matmul_kernel(
+    a: wp.array2d(dtype=float),  # NxK
+    b: wp.array2d(dtype=float),  # KxM
+    c: wp.array2d(dtype=float),  # NxM
+):
+    # launch dims should be (N, M)
+    i, j = wp.tid()
+    N = a.shape[0]
+    K = a.shape[1]
+    M = b.shape[1]
+    if i < N and j < M:
+        s = wp.float32(0)
+        for k in range(K):
+            s += a[i, k] * b[k, j]
+        c[i, j] = s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+def example1():
+    # two inputs and one output
+    jax_add = jax_kernel(add_kernel)
+
+    @jax.jit
+    def f():
+        n = 10
+        a = jnp.arange(n, dtype=jnp.int32)
+        b = jnp.ones(n, dtype=jnp.int32)
+        return jax_add(a, b)
+
+    print(f())
+
+
+def example2():
+    # one input and two outputs
+    jax_sincos = jax_kernel(sincos_kernel, num_outputs=2)
+
+    @jax.jit
+    def f():
+        n = 32
+        a = jnp.linspace(0, 2 * math.pi, n)
+        return jax_sincos(a)
+
+    s, c = f()
+    print(s)
+    print()
+    print(c)
+
+
+def example3():
+    # multiply vectors by scalar
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    @jax.jit
+    def f():
+        a = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # array of vec2
+        s = 2.0
+        return jax_scale_vec(a, s)
+
+    b = f()
+    print(b)
+
+
+def example4():
+    # multiply vectors by scalar (static arg)
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, s):
+        return jax_scale_vec(a, s)
+
+    a = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # array of vec2
+    s = 3.0
+
+    b = f(a, s)
+    print(b)
+
+
+def example5():
+    N, M, K = 3, 4, 2
+
+    # specify default launch dims
+    jax_matmul = jax_kernel(matmul_kernel, launch_dims=(N, M))
+
+    @jax.jit
+    def f():
+        a = jnp.full((N, K), 2, dtype=jnp.float32)
+        b = jnp.full((K, M), 3, dtype=jnp.float32)
+
+        # use default launch dims
+        return jax_matmul(a, b)
+
+    print(f())
+
+
+def example6():
+    # don't specify default launch dims
+    jax_matmul = jax_kernel(matmul_kernel)
+
+    @jax.jit
+    def f():
+        N1, M1, K1 = 3, 4, 2
+        a1 = jnp.full((N1, K1), 2, dtype=jnp.float32)
+        b1 = jnp.full((K1, M1), 3, dtype=jnp.float32)
+
+        # use custom launch dims
+        result1 = jax_matmul(a1, b1, launch_dims=(N1, M1))
+
+        N2, M2, K2 = 4, 3, 2
+        a2 = jnp.full((N2, K2), 2, dtype=jnp.float32)
+        b2 = jnp.full((K2, M2), 3, dtype=jnp.float32)
+
+        # use custom launch dims
+        result2 = jax_matmul(a2, b2, launch_dims=(N2, M2))
+
+        return result1, result2
+
+    r1, r2 = f()
+    print(r1)
+    print()
+    print(r2)
+
+
+def example7():
+    # no inputs and one output
+    jax_diagonal = jax_kernel(diagonal_kernel)
+
+    @jax.jit
+    def f():
+        # launch dimensions determine output size
+        return jax_diagonal(launch_dims=4)
+
+    print(f())
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2, example3, example4, example5, example6, example7]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()

warp/examples/optim/example_bounce.py

@@ -0,0 +1,266 @@
+# 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 Sim Grad Bounce
+#
+# Shows how to use Warp to optimize the initial velocity of a particle
+# such that it bounces off the wall and floor in order to hit a target.
+#
+# This example uses the built-in wp.Tape() object to compute gradients of
+# the distance to target (loss) w.r.t the initial velocity, followed by
+# a simple gradient-descent optimization step.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+
+
+@wp.kernel
+def loss_kernel(pos: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):
+    # distance to target
+    delta = pos[0] - target
+    loss[0] = wp.dot(delta, delta)
+
+
+@wp.kernel
+def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
+    tid = wp.tid()
+
+    # gradient descent step
+    x[tid] = x[tid] - grad[tid] * alpha
+
+
+class Example:
+    def __init__(self, stage_path="example_bounce.usd", verbose=False):
+        self.verbose = verbose
+
+        # seconds
+        sim_duration = 0.6
+
+        # control frequency
+        fps = 60
+        self.frame_dt = 1.0 / fps
+        frame_steps = int(sim_duration / self.frame_dt)
+
+        # sim frequency
+        self.sim_substeps = 8
+        self.sim_steps = frame_steps * self.sim_substeps
+        self.sim_dt = self.frame_dt / self.sim_substeps
+
+        self.iter = 0
+        self.render_time = 0.0
+
+        self.train_rate = 0.02
+
+        ke = 1.0e4
+        kf = 0.0
+        kd = 1.0e1
+        mu = 0.2
+
+        builder = wp.sim.ModelBuilder()
+        builder.add_particle(pos=wp.vec3(-0.5, 1.0, 0.0), vel=wp.vec3(5.0, -5.0, 0.0), mass=1.0)
+        builder.add_shape_box(body=-1, pos=wp.vec3(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=ke, kf=kf, kd=kd, mu=mu)
+
+        # use `requires_grad=True` to create a model for differentiable simulation
+        self.model = builder.finalize(requires_grad=True)
+        self.model.ground = True
+
+        self.model.soft_contact_ke = ke
+        self.model.soft_contact_kf = kf
+        self.model.soft_contact_kd = kd
+        self.model.soft_contact_mu = mu
+        self.model.soft_contact_margin = 10.0
+        self.model.soft_contact_restitution = 1.0
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+
+        self.target = (-2.0, 1.5, 0.0)
+        self.loss = wp.zeros(1, dtype=wp.float32, requires_grad=True)
+
+        # allocate sim states for trajectory
+        self.states = []
+        for _i in range(self.sim_steps + 1):
+            self.states.append(self.model.state())
+
+        # one-shot contact creation (valid if we're doing simple collision against a constant normal plane)
+        wp.sim.collide(self.model, self.states[0])
+
+        if stage_path:
+            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, scaling=1.0)
+        else:
+            self.renderer = None
+
+        # capture forward/backward passes
+        self.use_cuda_graph = wp.get_device().is_cuda
+        if self.use_cuda_graph:
+            with wp.ScopedCapture() as capture:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(self.loss)
+            self.graph = capture.graph
+
+    def forward(self):
+        # run control loop
+        for i in range(self.sim_steps):
+            self.states[i].clear_forces()
+            self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)
+
+        # compute loss on final state
+        wp.launch(loss_kernel, dim=1, inputs=[self.states[-1].particle_q, self.target, self.loss])
+
+        return self.loss
+
+    def step(self):
+        with wp.ScopedTimer("step"):
+            if self.use_cuda_graph:
+                wp.capture_launch(self.graph)
+            else:
+                self.tape = wp.Tape()
+                with self.tape:
+                    self.forward()
+                self.tape.backward(self.loss)
+
+            # gradient descent step
+            x = self.states[0].particle_qd
+            wp.launch(step_kernel, dim=len(x), inputs=[x, x.grad, self.train_rate])
+
+            x_grad = self.tape.gradients[self.states[0].particle_qd]
+
+            if self.verbose:
+                print(f"Iter: {self.iter} Loss: {self.loss}")
+                print(f"    x: {x} g: {x_grad}")
+
+            # clear grads for next iteration
+            self.tape.zero()
+
+            self.iter = self.iter + 1
+
+    def render(self):
+        if self.renderer is None:
+            return
+
+        with wp.ScopedTimer("render"):
+            # draw trajectory
+            traj_verts = [self.states[0].particle_q.numpy()[0].tolist()]
+
+            for i in range(0, self.sim_steps, self.sim_substeps):
+                traj_verts.append(self.states[i].particle_q.numpy()[0].tolist())
+
+                self.renderer.begin_frame(self.render_time)
+                self.renderer.render(self.states[i])
+                self.renderer.render_box(
+                    pos=self.target,
+                    rot=wp.quat_identity(),
+                    extents=(0.1, 0.1, 0.1),
+                    name="target",
+                    color=(0.0, 0.0, 0.0),
+                )
+                self.renderer.render_line_strip(
+                    vertices=traj_verts,
+                    color=wp.render.bourke_color_map(0.0, 7.0, self.loss.numpy()[0]),
+                    radius=0.02,
+                    name=f"traj_{self.iter - 1}",
+                )
+                self.renderer.end_frame()
+
+                from pxr import Gf, UsdGeom
+
+                particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
+                particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
+                particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)
+
+                self.render_time += self.frame_dt
+
+    def check_grad(self):
+        param = self.states[0].particle_qd
+
+        # initial value
+        x_c = param.numpy().flatten()
+
+        # compute numeric gradient
+        x_grad_numeric = np.zeros_like(x_c)
+
+        for i in range(len(x_c)):
+            eps = 1.0e-3
+
+            step = np.zeros_like(x_c)
+            step[i] = eps
+
+            x_1 = x_c + step
+            x_0 = x_c - step
+
+            param.assign(x_1)
+            l_1 = self.forward().numpy()[0]
+
+            param.assign(x_0)
+            l_0 = self.forward().numpy()[0]
+
+            dldx = (l_1 - l_0) / (eps * 2.0)
+
+            x_grad_numeric[i] = dldx
+
+        # reset initial state
+        param.assign(x_c)
+
+        # compute analytic gradient
+        tape = wp.Tape()
+        with tape:
+            l = self.forward()
+
+        tape.backward(l)
+
+        x_grad_analytic = tape.gradients[param]
+
+        print(f"numeric grad: {x_grad_numeric}")
+        print(f"analytic grad: {x_grad_analytic}")
+
+        tape.zero()
+
+
+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_bounce.usd",
+        help="Path to the output USD file.",
+    )
+    parser.add_argument("--train_iters", type=int, default=250, help="Total number of training iterations.")
+    parser.add_argument("--verbose", action="store_true", help="Print out additional status messages during execution.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(stage_path=args.stage_path, verbose=args.verbose)
+
+        example.check_grad()
+
+        # replay and optimize
+        for i in range(args.train_iters):
+            example.step()
+            if i % 16 == 0:
+                example.render()
+
+        if example.renderer:
+            example.renderer.save()