warp-lang 1.5.1__py3-none-manylinux2014_x86_64.whl → 1.6.0__py3-none-manylinux2014_x86_64.whl

warp/examples/sim/example_granular_collision_sdf.py

@@ -15,7 +15,6 @@
 # Note: requires a CUDA-capable device
 ###########################################################################
 
-import math
 import os
 
 import numpy as np
@@ -66,8 +65,8 @@ class Example:
             sdf=rock_sdf,
             body=-1,
             pos=wp.vec3(0.0, 0.0, 0.0),
-            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -0.5 * math.pi),
-            scale=wp.vec3(0.01, 0.01, 0.01),
+            rot=wp.quat(0.0, 0.0, 0.0, 1.0),
+            scale=wp.vec3(1.0, 1.0, 1.0),
         )
 
         mins = np.array([-3.0, -3.0, -3.0])
@@ -149,13 +148,13 @@ class Example:
         with wp.ScopedTimer("render"):
             self.renderer.begin_frame(self.sim_time)
 
-            # Note the extra wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi) is because .usd is oriented differently from .nvdb
             self.renderer.render_ref(
                 name="collision",
                 path=os.path.join(warp.examples.get_asset_directory(), "rocks.usd"),
                 pos=wp.vec3(0.0, 0.0, 0.0),
                 rot=wp.quat(0.0, 0.0, 0.0, 1.0),
-                scale=wp.vec3(0.01, 0.01, 0.01),
+                scale=wp.vec3(1.0, 1.0, 1.0),
+                color=(0.35, 0.55, 0.9),
             )
 
             self.renderer.render_sphere(

warp/examples/sim/example_jacobian_ik.py

@@ -43,8 +43,6 @@ class Example:
     def __init__(self, stage_path="example_jacobian_ik.usd", num_envs=10):
         rng = np.random.default_rng(42)
 
-        builder = wp.sim.ModelBuilder()
-
         self.num_envs = num_envs
 
         fps = 60

warp/examples/sim/example_quadruped.py

@@ -93,7 +93,7 @@ class Example:
         fps = 100
         self.frame_dt = 1.0 / fps
 
-        self.sim_substeps = 5
+        self.sim_substeps = 10
         self.sim_dt = self.frame_dt / self.sim_substeps
 
         self.num_envs = num_envs
@@ -116,10 +116,13 @@ class Example:
         self.model.joint_attach_ke = 16000.0
         self.model.joint_attach_kd = 200.0
         self.use_tile_gemm = False
+        self.fuse_cholesky = False
 
         # self.integrator = wp.sim.XPBDIntegrator()
         # self.integrator = wp.sim.SemiImplicitIntegrator()
-        self.integrator = wp.sim.FeatherstoneIntegrator(self.model, use_tile_gemm=self.use_tile_gemm)
+        self.integrator = wp.sim.FeatherstoneIntegrator(
+            self.model, use_tile_gemm=self.use_tile_gemm, fuse_cholesky=self.fuse_cholesky
+        )
 
         if stage_path:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage_path)

warp/examples/tile/example_tile_cholesky.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2025 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 Tile Cholesky
+#
+# Shows how to write a simple kernel computing a Cholesky factorize and
+# triangular solve using Warp Cholesky Tile APIs.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.init()
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 128
+TILE = 32
+
+# Both should work
+np_type, wp_type = np.float64, wp.float64
+# np_type, wp_type = np.float32, wp.float32
+
+
+@wp.kernel
+def cholesky(
+    A: wp.array2d(dtype=wp_type),
+    L: wp.array2d(dtype=wp_type),
+    X: wp.array1d(dtype=wp_type),
+    Y: wp.array1d(dtype=wp_type),
+):
+    i, j, _ = wp.tid()
+
+    a = wp.tile_load(A, shape=(TILE, TILE))
+    l = wp.tile_cholesky(a)
+    wp.tile_store(L, l)
+
+    x = wp.tile_load(X, shape=TILE)
+    y = wp.tile_cholesky_solve(l, x)
+    wp.tile_store(Y, y)
+
+
+if __name__ == "__main__":
+    wp.set_device("cuda:0")
+
+    A_h = np.ones((TILE, TILE), dtype=np_type) + 5 * np.diag(np.ones(TILE), 0)
+    L_h = np.zeros_like(A_h)
+
+    A_wp = wp.array2d(A_h, dtype=wp_type)
+    L_wp = wp.array2d(L_h, dtype=wp_type)
+
+    X_h = np.arange(TILE, dtype=np_type)
+    Y_h = np.zeros_like(X_h)
+
+    X_wp = wp.array1d(X_h, dtype=wp_type)
+    Y_wp = wp.array1d(Y_h, dtype=wp_type)
+
+    wp.launch_tiled(cholesky, dim=[1, 1], inputs=[A_wp, L_wp, X_wp, Y_wp], block_dim=BLOCK_DIM)
+
+    L_np = np.linalg.cholesky(A_h)
+    Y_np = np.linalg.solve(A_h, X_h)
+
+    print("A:\n", A_h)
+    print("L (Warp):\n", L_wp)
+    print("L (Numpy):\n", L_np)
+
+    print("x:\n", X_h)
+    print("A\\n (Warp):\n", Y_wp.numpy())
+    print("A\\x (Numpy):\n", Y_np)
+
+    assert np.allclose(Y_wp.numpy(), Y_np) and np.allclose(L_wp.numpy(), L_np)
+
+    print("Example Tile Cholesky passed")

warp/examples/tile/example_tile_convolution.py

@@ -34,11 +34,11 @@ def filter(x: wp.vec2d):
 @wp.kernel
 def conv_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d)):
     i, j, _ = wp.tid()
-    a = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
     wp.tile_fft(a)
     b = wp.tile_map(filter, a)
     wp.tile_ifft(b)
-    wp.tile_store(y, i, j, b)
+    wp.tile_store(y, b)
 
 
 if __name__ == "__main__":

warp/examples/tile/example_tile_fft.py

@@ -26,10 +26,10 @@ TILE_N = 32
 @wp.kernel
 def fft_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d)):
     i, j, _ = wp.tid()
-    a = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
     wp.tile_fft(a)
     wp.tile_ifft(a)
-    wp.tile_store(y, i, j, a)
+    wp.tile_store(y, a)
 
 
 if __name__ == "__main__":

warp/examples/tile/example_tile_filtering.py

@@ -38,12 +38,12 @@ def cplx_prod(x: wp.vec2d, y: wp.vec2d):
 @wp.kernel
 def conv_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d), z: wp.array2d(dtype=wp.vec2d)):
     i, j, _ = wp.tid()
-    a = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)
-    b = wp.tile_load(y, i, j, m=TILE_M, n=TILE_N)
+    a = wp.tile_load(x, shape=(TILE_M, TILE_N))
+    b = wp.tile_load(y, shape=(TILE_M, TILE_N))
     wp.tile_fft(a)
     c = wp.tile_map(cplx_prod, a, b)
     wp.tile_ifft(c)
-    wp.tile_store(z, i, j, c)
+    wp.tile_store(z, c)
 
 
 if __name__ == "__main__":

warp/examples/tile/example_tile_matmul.py

@@ -30,7 +30,7 @@ def tile_gemm(A: wp.array2d(dtype=wp.float32), B: wp.array2d(dtype=wp.float16),
     # output tile index
     i, j = wp.tid()
 
-    sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=wp.float64)
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float64)
 
     _M = A.shape[0]
     _N = B.shape[1]
@@ -39,13 +39,13 @@ def tile_gemm(A: wp.array2d(dtype=wp.float32), B: wp.array2d(dtype=wp.float16),
     count = int(K / TILE_K)
 
     for k in range(0, count):
-        a = wp.tile_load(A, i, k, m=TILE_M, n=TILE_K)
-        b = wp.tile_load(B, k, j, m=TILE_K, n=TILE_N)
+        a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
+        b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))
 
         # sum += a*b
         wp.tile_matmul(a, b, sum)
 
-    wp.tile_store(C, i, j, sum)
+    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
 
 
 if __name__ == "__main__":

warp/examples/tile/example_tile_mlp.py

@@ -117,23 +117,23 @@ def compute(
     f = wp.tile(local)
 
     # input layer
-    w0 = wp.tile_load(weights_0, 0, 0, m=DIM_HID, n=DIM_IN)
-    b0 = wp.tile_load(bias_0, 0, 0, m=DIM_HID, n=1)
-    z = wp.tile_map(relu, wp.tile_matmul(w0, f) + wp.tile_broadcast(b0, m=DIM_HID, n=NUM_THREADS))
+    w0 = wp.tile_load(weights_0, shape=(DIM_HID, DIM_IN))
+    b0 = wp.tile_load(bias_0, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w0, f) + wp.tile_broadcast(b0, shape=(DIM_HID, NUM_THREADS)))
 
     # hidden layer
-    w1 = wp.tile_load(weights_1, 0, 0, m=DIM_HID, n=DIM_HID)
-    b1 = wp.tile_load(bias_1, 0, 0, m=DIM_HID, n=1)
-    z = wp.tile_map(relu, wp.tile_matmul(w1, z) + wp.tile_broadcast(b1, m=DIM_HID, n=NUM_THREADS))
+    w1 = wp.tile_load(weights_1, shape=(DIM_HID, DIM_HID))
+    b1 = wp.tile_load(bias_1, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w1, z) + wp.tile_broadcast(b1, shape=(DIM_HID, NUM_THREADS)))
 
-    w2 = wp.tile_load(weights_2, 0, 0, m=DIM_HID, n=DIM_HID)
-    b2 = wp.tile_load(bias_2, 0, 0, m=DIM_HID, n=1)
-    z = wp.tile_map(relu, wp.tile_matmul(w2, z) + wp.tile_broadcast(b2, m=DIM_HID, n=NUM_THREADS))
+    w2 = wp.tile_load(weights_2, shape=(DIM_HID, DIM_HID))
+    b2 = wp.tile_load(bias_2, shape=(DIM_HID, 1))
+    z = wp.tile_map(relu, wp.tile_matmul(w2, z) + wp.tile_broadcast(b2, shape=(DIM_HID, NUM_THREADS)))
 
     # output layer
-    w3 = wp.tile_load(weights_3, 0, 0, m=DIM_OUT, n=DIM_HID)
-    b3 = wp.tile_load(bias_3, 0, 0, m=DIM_OUT, n=1)
-    o = wp.tile_map(relu, wp.tile_matmul(w3, z) + wp.tile_broadcast(b3, m=DIM_OUT, n=NUM_THREADS))
+    w3 = wp.tile_load(weights_3, shape=(DIM_OUT, DIM_HID))
+    b3 = wp.tile_load(bias_3, shape=(DIM_OUT, 1))
+    o = wp.tile_map(relu, wp.tile_matmul(w3, z) + wp.tile_broadcast(b3, shape=(DIM_OUT, NUM_THREADS)))
 
     # untile back to SIMT
     output = wp.untile(o)

warp/examples/tile/example_tile_nbody.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2025 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 N-Body
+#
+# Shows how to simulate an N-Body gravitational problem using an all-pairs
+# approach with Warp tile primitives.
+#
+# References:
+#   L. Nyland, M. Harris, and J. Prins. "Fast N-Body Simulation with
+#   CUDA" in GPU Gems 3. H. Nguyen, Addison-Wesley Professional, 2007.
+#   https://developer.nvidia.com/gpugems/gpugems3/part-v-physics-simulation/chapter-31-fast-n-body-simulation-cuda
+#
+###########################################################################
+
+import argparse
+
+import numpy as np
+
+import warp as wp
+
+wp.init()
+
+DT = wp.constant(0.01)
+SOFTENING_SQ = wp.constant(0.1**2)  # Softening factor for numerical stability
+TILE_SIZE = wp.constant(64)
+PARTICLE_MASS = wp.constant(1.0)
+
+
+@wp.func
+def body_body_interaction(p0: wp.vec3, pi: wp.vec3):
+    """Return the acceleration of the particle at position `p0` due to the
+    particle at position `pi`."""
+    r = pi - p0
+
+    dist_sq = wp.length_sq(r) + SOFTENING_SQ
+
+    inv_dist = 1.0 / wp.sqrt(dist_sq)
+    inv_dist_cubed = inv_dist * inv_dist * inv_dist
+
+    acc = PARTICLE_MASS * inv_dist_cubed * r
+
+    return acc
+
+
+@wp.kernel
+def integrate_bodies_tiled(
+    old_position: wp.array(dtype=wp.vec3),
+    velocity: wp.array(dtype=wp.vec3),
+    new_position: wp.array(dtype=wp.vec3),
+    num_bodies: int,
+):
+    i = wp.tid()
+
+    p0 = old_position[i]
+
+    accel = wp.vec3(0.0, 0.0, 0.0)
+
+    for k in range(num_bodies / TILE_SIZE):
+        k_tile = wp.tile_load(old_position, shape=TILE_SIZE, offset=k * TILE_SIZE)
+        for idx in range(TILE_SIZE):
+            pi = k_tile[idx]
+            accel += body_body_interaction(p0, pi)
+
+    # Advance the velocity one timestep (in-place)
+    velocity[i] = velocity[i] + accel * DT
+
+    # Advance the positions (using a second array)
+    new_position[i] = old_position[i] + DT * velocity[i]
+
+
+class Example:
+    def __init__(self, headless=False, num_bodies=1024):
+        self.num_bodies = num_bodies
+
+        rng = np.random.default_rng(42)
+
+        # Sample the surface of a sphere
+        r = 10.0 * (num_bodies / 1024) ** (1 / 2)  # Scale factor to maintain a constant density
+        phi = np.arccos(1.0 - 2.0 * rng.uniform(size=self.num_bodies))
+        theta = rng.uniform(low=0.0, high=2.0 * np.pi, size=self.num_bodies)
+        x = r * np.cos(theta) * np.sin(phi)
+        y = r * np.sin(theta) * np.sin(phi)
+        z = r * np.cos(phi)
+
+        self.scale = r
+        init_pos_np = np.stack((x, y, z), axis=1)
+
+        self.pos_array_0 = wp.array(init_pos_np, dtype=wp.vec3)
+        self.pos_array_1 = wp.empty_like(self.pos_array_0)
+        self.vel_array = wp.zeros(self.num_bodies, dtype=wp.vec3)
+
+        if headless:
+            self.scatter_plot = None
+        else:
+            self.scatter_plot = self.create_plot()
+
+    def create_plot(self):
+        import matplotlib.pyplot as plt
+
+        # Create a figure and a 3D axis for the plot
+        self.fig = plt.figure()
+        ax = self.fig.add_subplot(111, projection="3d")
+
+        # Scatter plot of initial positions
+        init_pos_np = self.pos_array_0.numpy()
+        scatter_plot = ax.scatter(init_pos_np[:, 0], init_pos_np[:, 1], init_pos_np[:, 2], c="#76b900", alpha=0.5)
+
+        # Set axis limits
+        ax.set_xlim(-self.scale, self.scale)
+        ax.set_ylim(-self.scale, self.scale)
+        ax.set_zlim(-self.scale, self.scale)
+
+        return scatter_plot
+
+    def step(self):
+        wp.launch(
+            integrate_bodies_tiled,
+            dim=self.num_bodies,
+            inputs=[self.pos_array_0, self.vel_array, self.pos_array_1, self.num_bodies],
+            block_dim=TILE_SIZE,
+        )
+
+        # Swap arrays
+        (self.pos_array_0, self.pos_array_1) = (self.pos_array_1, self.pos_array_0)
+
+    def render(self):
+        positions_cpu = self.pos_array_0.numpy()
+
+        # Update scatter plot positions
+        self.scatter_plot._offsets3d = (
+            positions_cpu[:, 0],
+            positions_cpu[:, 1],
+            positions_cpu[:, 2],
+        )
+
+    # Function to update the scatter plot
+    def step_and_render(self, frame):
+        self.step()
+        self.render()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--num_frames", type=int, default=1000, help="Total number of frames.")
+    parser.add_argument("-N", help="Number of bodies. Should be a multiple of 64.", type=int, default=1024)
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    if args.device == "cpu":
+        print("This example only runs on CUDA devices.")
+        exit()
+
+    with wp.ScopedDevice(args.device):
+        example = Example(headless=args.headless, num_bodies=args.N)
+
+        if not args.headless:
+            import matplotlib.pyplot as plt
+            from matplotlib.animation import FuncAnimation
+
+            # Create the animation
+            ani = FuncAnimation(example.fig, example.step_and_render, frames=args.num_frames, interval=50, repeat=False)
+
+            # Display the animation
+            plt.show()
+
+        else:
+            for _ in range(args.num_frames):
+                example.step()