warp-lang 1.5.0__py3-none-macosx_10_13_universal2.whl → 1.6.0__py3-none-macosx_10_13_universal2.whl

warp/examples/benchmarks/benchmark_gemm.py

@@ -13,23 +13,28 @@ wp.clear_kernel_cache()
 wp.set_module_options({"fast_math": True, "enable_backward": False})
 
 
-def create_mlp_kernel(m, n, k):
+# returns a kernel to compute a GEMM given m,n,k tile sizes
+def create_gemm_kernel(m, n, k):
     TILE_M = m
     TILE_N = n
     TILE_K = k
 
     @wp.kernel
-    def mlp(x: wp.array2d(dtype=float), weights_wp: wp.array2d(dtype=float), n_k: int, output: wp.array2d(dtype=float)):
-        i_m, i_n = wp.tid()
-        sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=wp.float32)
-        for count in range(n_k):
-            feat = wp.tile_load(x, i_m, count, TILE_M, TILE_K)
-            weight = wp.tile_load(weights_wp, count, i_n, TILE_K, TILE_N)
-            wp.tile_matmul(feat, weight, sum)
+    def gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+        i, j = wp.tid()
+        sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)
 
-        wp.tile_store(output, i_m, i_n, sum)
+        count = A.shape[1] // TILE_K
 
-    return mlp
+        for k in range(count):
+            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))
+
+            wp.tile_matmul(a, b, sum)
+
+        wp.tile_store(output, sum, offset=(i * TILE_M, j * TILE_N))
+
+    return gemm
 
 
 def benchmark_torch(A, B, warm_up, iterations):
@@ -55,19 +60,25 @@ def benchmark_warp(A, B, config, warm_up, iterations):
     TILE_K = config[2]
     BLOCK_DIM = config[3]
 
-    mlp = create_mlp_kernel(TILE_M, TILE_N, TILE_K)
+    mlp = create_gemm_kernel(TILE_M, TILE_N, TILE_K)
 
     M = A.shape[0]
     N = B.shape[1]
-    K = A.shape[1]
 
     output = wp.zeros((M, N), dtype=float)
 
+    # create launch command
+    cmd = wp.launch_tiled(
+        kernel=mlp,
+        dim=[M // TILE_M, N // TILE_N],
+        inputs=[A, B, output],
+        block_dim=BLOCK_DIM,
+        record_cmd=True,
+    )
+
     # warm-up
     for _ in range(warm_up):
-        wp.launch_tiled(
-            kernel=mlp, dim=[M // TILE_M, N // TILE_N], inputs=[A, B, K // TILE_K, output], block_dim=BLOCK_DIM
-        )
+        cmd.launch()
 
     # check output
     if warm_up > 0:
@@ -77,9 +88,7 @@ def benchmark_warp(A, B, config, warm_up, iterations):
     timers = {}
     with wp.ScopedTimer("warp", print=False, dict=timers, synchronize=True):
         for _ in range(iterations):
-            wp.launch_tiled(
-                kernel=mlp, dim=[M // TILE_M, N // TILE_N], inputs=[A, B, K // TILE_K, output], block_dim=BLOCK_DIM
-            )
+            cmd.launch()
 
     return timers["warp"][0]
 

warp/examples/benchmarks/benchmark_interop_paddle.py

@@ -58,7 +58,7 @@ def test_from_paddle(kernel, num_iters, array_size, device, warp_dtype=None):
         wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_paddle(...)")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  from_paddle(...)")
 
     # profiler.stop()
     # profiler.print()
@@ -97,7 +97,7 @@ def test_array_ctype_from_paddle(kernel, num_iters, array_size, device, warp_dty
         wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_paddle(..., return_ctype=True)")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  from_paddle(..., return_ctype=True)")
 
     # profiler.stop()
     # profiler.print()
@@ -131,7 +131,7 @@ def test_direct_from_paddle(kernel, num_iters, array_size, device, warp_dtype=No
         wp.launch(kernel, dim=array_size, inputs=[_a, _b, _c, _d, _e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  direct from paddle")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  direct from paddle")
 
     # profiler.stop()
     # profiler.print()

warp/examples/benchmarks/benchmark_interop_torch.py

@@ -58,7 +58,7 @@ def test_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
         wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_torch(...)")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  from_torch(...)")
 
     # profiler.stop()
     # profiler.print()
@@ -97,7 +97,7 @@ def test_array_ctype_from_torch(kernel, num_iters, array_size, device, warp_dtyp
         wp.launch(kernel, dim=array_size, inputs=[a, b, c, d, e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  from_torch(..., return_ctype=True)")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  from_torch(..., return_ctype=True)")
 
     # profiler.stop()
     # profiler.print()
@@ -131,7 +131,7 @@ def test_direct_from_torch(kernel, num_iters, array_size, device, warp_dtype=Non
         wp.launch(kernel, dim=array_size, inputs=[_a, _b, _c, _d, _e])
 
     t2 = time.time_ns()
-    print(f"{(t2 - t1) / 1_000_000 :8.0f} ms  direct from torch")
+    print(f"{(t2 - t1) / 1_000_000:8.0f} ms  direct from torch")
 
     # profiler.stop()
     # profiler.print()

warp/examples/core/example_torch.py

@@ -23,8 +23,6 @@ import torch
 
 import warp as wp
 
-pvec2 = wp.types.vector(length=2, dtype=wp.float32)
-
 
 # Define the Rosenbrock function
 @wp.func
@@ -33,11 +31,7 @@ def rosenbrock(x: float, y: float):
 
 
 @wp.kernel
-def eval_rosenbrock(
-    xs: wp.array(dtype=pvec2),
-    # outputs
-    z: wp.array(dtype=float),
-):
+def eval_rosenbrock(xs: wp.array(dtype=wp.vec2), z: wp.array(dtype=float)):
     i = wp.tid()
     x = xs[i]
     z[i] = rosenbrock(x[0], x[1])
@@ -46,7 +40,7 @@ def eval_rosenbrock(
 class Rosenbrock(torch.autograd.Function):
     @staticmethod
     def forward(ctx, xy, num_particles):
-        ctx.xy = wp.from_torch(xy, dtype=pvec2, requires_grad=True)
+        ctx.xy = wp.from_torch(xy, dtype=wp.vec2, requires_grad=True)
         ctx.num_particles = num_particles
 
         # allocate output
@@ -76,10 +70,10 @@ class Rosenbrock(torch.autograd.Function):
 
 
 class Example:
-    def __init__(self, headless=False, train_iters=10):
-        self.num_particles = 1500
+    def __init__(self, headless=False, train_iters=10, num_particles=1500):
+        self.num_particles = num_particles
         self.train_iters = train_iters
-        self.frame = 0
+        self.train_iter = 0
 
         self.learning_rate = 5e-2
 
@@ -113,7 +107,7 @@ class Example:
         xy = np.column_stack((X.flatten(), Y.flatten()))
         N = len(xy)
 
-        xy = wp.array(xy, dtype=pvec2)
+        xy = wp.array(xy, dtype=wp.vec2)
         Z = wp.empty(N, dtype=wp.float32)
 
         wp.launch(eval_rosenbrock, dim=N, inputs=[xy], outputs=[Z])
@@ -122,23 +116,8 @@ class Example:
         # Plot the function as a heatmap
         self.fig = plt.figure(figsize=(6, 6))
         ax = plt.gca()
-        plt.imshow(
-            Z,
-            extent=[min_x, max_x, min_y, max_y],
-            origin="lower",
-            interpolation="bicubic",
-            cmap="coolwarm",
-        )
-        plt.contour(
-            X,
-            Y,
-            Z,
-            extent=[min_x, max_x, min_y, max_y],
-            levels=150,
-            colors="k",
-            alpha=0.5,
-            linewidths=0.5,
-        )
+        plt.imshow(Z, extent=[min_x, max_x, min_y, max_y], origin="lower", interpolation="bicubic", cmap="coolwarm")
+        plt.contour(X, Y, Z, extent=[min_x, max_x, min_y, max_y], levels=150, colors="k", alpha=0.5, linewidths=0.5)
 
         # Plot optimum
         plt.plot(1, 1, "*", color="r", markersize=10)
@@ -167,16 +146,16 @@ class Example:
 
         # Compute mean
         self.mean_pos = np.mean(self.xy_np, axis=0)
-        print(f"\rFrame {self.frame:5d} particle mean: {self.mean_pos[0]:.8f}, {self.mean_pos[1]:.8f}    ", end="")
+        print(f"\rIter {self.train_iter:5d} particle mean: {self.mean_pos[0]:.8f}, {self.mean_pos[1]:.8f}    ", end="")
 
-        self.frame += 1
+        self.train_iter += 1
 
     def render(self):
         if self.scatter_plot is None:
             return
 
         self.scatter_plot.set_offsets(np.c_[self.xy_np[:, 0], self.xy_np[:, 1]])
-        self.mean_marker.set_data(self.mean_pos[0], self.mean_pos[1])
+        self.mean_marker.set_data([self.mean_pos[0]], [self.mean_pos[1]])
 
     # Function to update the scatter plot
     def step_and_render(self, frame):
@@ -193,6 +172,9 @@ if __name__ == "__main__":
     parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
     parser.add_argument("--num_frames", type=int, default=10000, help="Total number of frames.")
     parser.add_argument("--train_iters", type=int, default=10, help="Total number of training iterations per frame.")
+    parser.add_argument(
+        "--num_particles", type=int, default=1500, help="Total number of particles to use in optimization."
+    )
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -202,14 +184,16 @@ if __name__ == "__main__":
     args = parser.parse_known_args()[0]
 
     with wp.ScopedDevice(args.device):
-        example = Example(headless=args.headless, train_iters=args.train_iters)
+        example = Example(headless=args.headless, train_iters=args.train_iters, num_particles=args.num_particles)
 
         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=100)
+            ani = FuncAnimation(
+                example.fig, example.step_and_render, frames=args.num_frames, interval=100, repeat=False
+            )
 
             # Display the animation
             plt.show()

warp/examples/fem/example_apic_fluid.py

@@ -220,6 +220,7 @@ class Example:
             jitter=self.radius * 1.0,
         )
         self.model: Model = builder.finalize()
+        self.model.ground = False
 
         # Storage for temporary variables
         self.temporary_store = fem.TemporaryStore()

warp/examples/fem/example_mixed_elasticity.py

@@ -236,7 +236,7 @@ class Example:
         final_area = fem.integrate(
             area_form, quadrature=fem.RegularQuadrature(domain, order=4), fields={"u_cur": self._u_field}
         )
-        print(f"Area gain: {final_area}  (using Poisson ratio={self._lame[0] / (self._lame[0] + 2.0*self._lame[1])})")
+        print(f"Area gain: {final_area}  (using Poisson ratio={self._lame[0] / (self._lame[0] + 2.0 * self._lame[1])})")
 
     def render(self):
         self.renderer.add_field("solution", self._u_field)

warp/examples/optim/example_bounce.py

@@ -169,7 +169,7 @@ class Example:
                     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}",
+                    name=f"traj_{self.iter - 1}",
                 )
                 self.renderer.end_frame()
 

warp/examples/optim/example_cloth_throw.py

@@ -184,7 +184,7 @@ class Example:
                     vertices=traj_verts,
                     color=wp.render.bourke_color_map(0.0, 269.0, self.loss.numpy()[0]),
                     radius=0.02,
-                    name=f"traj_{self.iter-1}",
+                    name=f"traj_{self.iter - 1}",
                 )
                 self.renderer.end_frame()
 

warp/examples/optim/example_diffray.py

@@ -166,24 +166,12 @@ def draw_kernel(
         j = mesh.indices[query.face * 3 + 1]
         k = mesh.indices[query.face * 3 + 2]
 
-        a = mesh.vertices[i]
-        b = mesh.vertices[j]
-        c = mesh.vertices[k]
-
-        p = wp.mesh_eval_position(mesh.id, query.face, query.u, query.v)
-
-        # barycentric coordinates
-        tri_area = wp.length(wp.cross(b - a, c - a))
-        w = wp.length(wp.cross(b - a, p - a)) / tri_area
-        v = wp.length(wp.cross(p - a, c - a)) / tri_area
-        u = 1.0 - w - v
-
         a_n = mesh.vertex_normals[i]
         b_n = mesh.vertex_normals[j]
         c_n = mesh.vertex_normals[k]
 
         # vertex normal interpolation
-        normal = u * a_n + v * b_n + w * c_n
+        normal = query.u * a_n + query.v * b_n + (1.0 - query.u - query.v) * c_n
 
         if mode == 0 or mode == 1:
             if mode == 0:  # grayscale
@@ -194,7 +182,7 @@ def draw_kernel(
                 tex_b = mesh.tex_coords[mesh.tex_indices[query.face * 3 + 1]]
                 tex_c = mesh.tex_coords[mesh.tex_indices[query.face * 3 + 2]]
 
-                tex = u * tex_a + v * tex_b + w * tex_c
+                tex = query.u * tex_a + query.v * tex_b + (1.0 - query.u - query.v) * tex_c
 
                 color = texture_interpolation(tex, texture)
 
@@ -352,7 +340,8 @@ class Example:
         # construct RenderMesh
         self.render_mesh = RenderMesh()
         self.mesh = wp.Mesh(
-            points=wp.array(points, dtype=wp.vec3, requires_grad=True), indices=wp.array(indices, dtype=int)
+            points=wp.array(points, dtype=wp.vec3, requires_grad=True),
+            indices=wp.array(indices, dtype=int),
         )
         self.render_mesh.id = self.mesh.id
         self.render_mesh.vertices = self.mesh.points

warp/examples/optim/example_drone.py

@@ -702,7 +702,7 @@ class Example:
     def step(self):
         if self.frame % int((self.num_frames / len(self.targets))) == 0:
             if self.verbose:
-                print(f"Choosing new flight target: {self.target_idx+1}")
+                print(f"Choosing new flight target: {self.target_idx + 1}")
 
             self.target_idx += 1
             self.target_idx %= len(self.targets)