warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/tile/example_tile_fft.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Tile FFT
+#
+# Shows how to write a simple FFT kernel using Warp tile primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 8
+TILE_M = 1
+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, shape=(TILE_M, TILE_N))
+    wp.tile_fft(a)
+    wp.tile_ifft(a)
+    wp.tile_store(y, a)
+
+
+if __name__ == "__main__":
+    wp.set_device("cuda:0")
+
+    x_h = np.ones((TILE_M, TILE_N, 2), dtype=np.float64)
+    x_h[:, :, 1] = 0
+    y_h = 3 * np.ones((TILE_M, TILE_N, 2), dtype=np.float64)
+    x_wp = wp.array2d(x_h, dtype=wp.vec2d)
+    y_wp = wp.array2d(y_h, dtype=wp.vec2d)
+
+    wp.launch_tiled(fft_tiled, dim=[1, 1], inputs=[x_wp], outputs=[y_wp], block_dim=BLOCK_DIM)
+
+    print("Inputs:\n", x_wp)  # [1+0i, 1+0i, 1+0i, ...]
+    print("Output:\n", y_wp)  # [32+0i, 0, 0, ...]

warp/examples/tile/example_tile_filtering.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Tile Filtering
+#
+# Shows how to write a simple filtering kernel using Warp FFT tile
+# primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+wp.set_module_options({"enable_backward": False})
+
+BLOCK_DIM = 128
+TILE_M = 1
+TILE_N = 512
+
+scale = wp.vec2d(wp.float64(1 / TILE_N), wp.float64(1 / TILE_N))
+
+
+def cplx(array):
+    return array[..., 0] + 1j * array[..., 1]
+
+
+@wp.func
+def cplx_prod(x: wp.vec2d, y: wp.vec2d):
+    return wp.cw_mul(wp.vec2d(x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0]), scale)
+
+
+@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, 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, c)
+
+
+if __name__ == "__main__":
+    rng = np.random.default_rng(42)
+
+    # Create noisy input signal
+    t = np.linspace(0, 2 * np.pi, TILE_N, dtype=np.float64)
+    x = np.sin(t) + 0.5 * rng.random(TILE_N, dtype=np.float64)
+
+    # Create filter. This filter keeps only ~10% of the frequencies at the center
+    # of the spectrum.
+    f = np.ones_like(x)
+    freq = np.fft.fftfreq(TILE_N)
+    f[np.abs(freq) > 0.05] = 0.0
+    f[np.abs(freq) <= 0.05] = 1.0
+
+    # Create Warp input data
+    # We use vec2d to hold complex numbers
+    x_h = np.zeros((TILE_M, TILE_N, 2), dtype=np.float64)
+    f_h = np.zeros_like(x_h)
+    y_h = np.zeros_like(f_h)
+
+    x_h[:, :, 0] = x
+    f_h[:, :, 0] = f
+
+    x_wp = wp.array2d(x_h, dtype=wp.vec2d)
+    f_wp = wp.array2d(f_h, dtype=wp.vec2d)
+    y_wp = wp.array2d(y_h, dtype=wp.vec2d)
+
+    wp.launch_tiled(conv_tiled, dim=[1, 1], inputs=[x_wp, f_wp], outputs=[y_wp], block_dim=BLOCK_DIM)
+
+    # Extract output and compare with numpy
+    x_np = cplx(x_h)
+    f_np = cplx(f_h)
+    y_test = cplx(y_wp.numpy())
+    y_ref = np.fft.ifft(f_np * np.fft.fft(x_np))
+    np.testing.assert_allclose(y_ref, y_test)
+
+try:
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(10, 5))
+
+    ax.plot(
+        x,
+        color="#DDDDDD",
+        linewidth=2,
+        label="Original",
+    )
+    ax.plot(y_test[0, :].real, color="#76B900", linewidth=3, label="Smoothed")
+
+    ax.legend()
+    ax.grid(True)
+
+    plt.tight_layout()
+    plt.show()
+
+except ModuleNotFoundError:
+    print("Matplotlib not available; skipping figure")

warp/examples/tile/example_tile_matmul.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 Tile MatMul
+#
+# Shows how to write a simple GEMM kernel using Warp tile primitives.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+
+# tile size
+TILE_M = wp.constant(8)
+TILE_N = wp.constant(4)
+TILE_K = wp.constant(8)
+
+# num threads per-tile
+TILE_THREADS = 64
+
+
+@wp.kernel
+def tile_gemm(A: wp.array2d(dtype=wp.float32), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float64)):
+    # output tile index
+    i, j = wp.tid()
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float64)
+
+    _M = A.shape[0]
+    _N = B.shape[1]
+    K = A.shape[1]
+
+    count = int(K / TILE_K)
+
+    for k in range(0, 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))
+
+        # sum += a*b
+        wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+
+
+if __name__ == "__main__":
+    # generate some tile aligned matrix dimensions
+    M = TILE_M * 7
+    K = TILE_K * 6
+    N = TILE_N * 5
+
+    rng = np.random.default_rng(42)
+    A = rng.random((M, K), dtype=np.float32)
+    B = rng.random((K, N), dtype=np.float32).astype(np.float16)
+    C = np.zeros((M, N), dtype=np.float64)
+
+    A_wp = wp.array(A, requires_grad=True)
+    B_wp = wp.array(B, requires_grad=True)
+    C_wp = wp.array(C, requires_grad=True)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_gemm,
+            dim=(M // TILE_M, N // TILE_N),
+            inputs=[A_wp, B_wp],
+            outputs=[C_wp],
+            block_dim=TILE_THREADS,
+        )
+
+    np.testing.assert_allclose(C_wp.numpy(), A @ B, atol=1.0e-4)
+
+    print("Example matrix multiplication passed")

warp/examples/tile/example_tile_mcgp.py

@@ -0,0 +1,191 @@
+# 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 Tile Monte Carlo Geometry Processing
+#
+# Shows how to use the built-in wp.Mesh data structure and
+# wp.mesh_eval_position() function to implement a tile-based, grid-free
+# Laplace solver using a Monte Carlo approach.
+#
+# References:
+#   Rohan Sawhney and Keenan Crane. Monte Carlo Geometry Processing:
+#   A Grid-Free Approach to PDE-Based Methods on Volumetric Domains.
+#   ACM Trans. Graph., Vol. 38, No. 4, Article 1. Published July 2020
+#
+##############################################################################
+
+import os
+
+import numpy as np
+from pxr import Usd, UsdGeom
+
+import warp as wp
+import warp.examples
+
+TILE_SIZE = wp.constant(64)  # number of samples for random walk
+STEPS = wp.constant(8)
+SEED = wp.constant(42)
+TOL = wp.constant(0.001)
+
+wp.config.enable_backward = False
+
+
+@wp.func
+def update_radius(p: wp.vec3):
+    """new radius is the distance to the closest point on the mesh"""
+    radius = 0.0
+    query = wp.mesh_query_point(MESH_ID, p, 1e6)
+
+    if query.result:
+        q = wp.mesh_eval_position(MESH_ID, query.face, query.u, query.v)
+        radius = -wp.length(q - p) * query.sign
+
+    return radius
+
+
+@wp.func
+def get_boundary_value(p: wp.vec3):
+    """analytic boundary condition on the surface of the mesh"""
+    return wp.abs(wp.sin(5.0 * wp.length(p)))
+
+
+@wp.func
+def walk(p: wp.vec3, rand_offset: int):
+    """random walk on spheres"""
+    rng = wp.rand_init(SEED, rand_offset)
+    for _ in range(0, STEPS):
+        r = update_radius(p)
+        if r < 0.0:  # outside the mesh
+            return 0.0
+        elif r < TOL:  # within the epsilon boundary
+            return get_boundary_value(p)
+        pr = wp.sample_unit_sphere_surface(rng) * r
+        p += pr
+    return get_boundary_value(p)  # closest choice if epsilon boundary not reached
+
+
+@wp.kernel
+def sphere_walk(delta_z: float, samples: wp.array2d(dtype=wp.vec3), solutions: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    sample_origin = samples[i, j] + wp.vec3(0.0, 0.0, delta_z)
+
+    rand_samples = wp.tile_full(TILE_SIZE, value=sample_origin, dtype=wp.vec3)
+
+    # every random sample gets a unique offset for rng
+    rand_offsets = wp.tile_arange(TILE_SIZE, dtype=int)
+    rand_offsets += wp.tile_full(TILE_SIZE, value=(i * TILE_SIZE + j) * TILE_SIZE, dtype=int)
+
+    # mcgp
+    walk_results = wp.tile_map(walk, rand_samples, rand_offsets)
+
+    # solution is an average of all walks originating from this position
+    walk_sum = wp.tile_sum(walk_results)
+    result = walk_sum * (1.0 / wp.float32(TILE_SIZE))
+
+    wp.tile_store(solutions[i], result, offset=(j,))
+
+
+class Example:
+    def __init__(self, height=256, slices=60):
+        usd_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+        self.mesh = wp.Mesh(
+            points=wp.array(usd_geom.GetPointsAttr().Get(), dtype=wp.vec3),
+            indices=wp.array(usd_geom.GetFaceVertexIndicesAttr().Get(), dtype=int),
+            bvh_leaf_size=1,
+        )
+
+        # z-slice scanning grid
+        self.height = height
+        self.width = self.height
+
+        x = np.linspace(-1.15, 0.85, self.width)
+        y = np.linspace(-0.2, 1.8, self.height)
+        xv, yv = np.meshgrid(x, y, indexing="ij")
+        zv = np.zeros_like(xv)
+        grid = np.stack((xv, yv, zv), axis=-1)
+
+        self.grid = wp.array(grid, dtype=wp.vec3)
+        self.pixels = wp.zeros((self.height, self.width), dtype=float)
+
+        self.slices = slices
+        self.z = np.linspace(-0.6, 0.8, self.slices)
+
+        # storage for animation
+        self.images = np.zeros((self.slices, self.height, self.width))
+
+    def render(self, slice):
+        wp.launch_tiled(
+            sphere_walk,
+            dim=[self.height, self.width],
+            inputs=[self.z[slice], self.grid],
+            outputs=[self.pixels],
+            block_dim=TILE_SIZE,
+        )
+
+        self.images[slice] = self.pixels.numpy()
+        print(f"slice: {slice}")
+
+    def get_animation(self):
+        fig, ax = plt.subplots()
+        plt.axis("off")
+        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+        plt.margins(0, 0)
+
+        slices = []
+        for i in range(self.slices):
+            slice = ax.imshow(np.flip(self.images[i, :, :].T), animated=True)
+            slices.append([slice])
+
+        ani = animation.ArtistAnimation(fig, slices, interval=60, blit=True, repeat_delay=1000)
+        return ani
+
+
+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("--height", type=int, default=256, help="Height of rendered image in pixels.")
+    parser.add_argument("--slices", type=int, default=60, help="Number of planar z-slices to scan.")
+    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]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(height=args.height, slices=args.slices)
+
+        # todo: permit runtime constants to be passed to tile map functions
+        MESH_ID = wp.constant(wp.uint64(example.mesh.id))
+
+        for i in range(args.slices):
+            example.render(i)
+
+        if not args.headless:
+            import matplotlib.animation as animation
+            import matplotlib.pyplot as plt
+
+            print("Creating the animation")
+            anim = example.get_animation()
+            anim_filename = "example_tile_mcgp_animation.gif"
+            anim.save(anim_filename, dpi=300, writer=animation.PillowWriter(fps=5))
+            print(f"Saved the animation at `{anim_filename}`")