warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/benchmarks/benchmark_interop_torch.py

@@ -0,0 +1,166 @@
+# 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.
+
+import time
+
+import torch
+
+import warp as wp
+
+
+def create_simple_kernel(dtype):
+    def simple_kernel(
+        a: wp.array(dtype=dtype),
+        b: wp.array(dtype=dtype),
+        c: wp.array(dtype=dtype),
+        d: wp.array(dtype=dtype),
+        e: wp.array(dtype=dtype),
+    ):
+        pass
+
+    return wp.Kernel(simple_kernel)
+
+
+def test_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_torch(_a, dtype=warp_dtype)
+        b = wp.from_torch(_b, dtype=warp_dtype)
+        c = wp.from_torch(_c, dtype=warp_dtype)
+        d = wp.from_torch(_d, dtype=warp_dtype)
+        e = wp.from_torch(_e, dtype=warp_dtype)
+        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(...)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_array_ctype_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_torch(_a, dtype=warp_dtype, return_ctype=True)
+        b = wp.from_torch(_b, dtype=warp_dtype, return_ctype=True)
+        c = wp.from_torch(_c, dtype=warp_dtype, return_ctype=True)
+        d = wp.from_torch(_d, dtype=warp_dtype, return_ctype=True)
+        e = wp.from_torch(_e, dtype=warp_dtype, return_ctype=True)
+        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)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_direct_from_torch(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    torch_device = wp.device_to_torch(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        torch_shape = (array_size, *warp_dtype._shape_)
+        torch_dtype = wp.dtype_to_torch(warp_dtype._wp_scalar_type_)
+    else:
+        torch_shape = (array_size,)
+        torch_dtype = torch.float32 if warp_dtype is None else wp.dtype_to_torch(warp_dtype)
+
+    _a = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _b = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _c = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _d = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+    _e = torch.zeros(torch_shape, dtype=torch_dtype, device=torch_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        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")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+wp.init()
+
+params = [
+    # (warp_dtype arg, kernel)
+    (None, create_simple_kernel(wp.float32)),
+    (wp.float32, create_simple_kernel(wp.float32)),
+    (wp.vec3f, create_simple_kernel(wp.vec3f)),
+    (wp.mat22f, create_simple_kernel(wp.mat22f)),
+]
+
+wp.load_module()
+
+num_iters = 100000
+
+for warp_dtype, kernel in params:
+    print(f"\ndtype={wp.context.type_str(warp_dtype)}")
+    test_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_array_ctype_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_direct_from_torch(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)

warp/examples/benchmarks/benchmark_launches.py

@@ -0,0 +1,301 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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.
+
+###########################################################################
+# 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()  # noqa: F841
+
+
+@wp.kernel
+def kf(x: float, y: float, z: float):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def kv(u: wp.vec3, v: wp.vec3, w: wp.vec3):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def km(M: wp.mat33, N: wp.mat33, O: wp.mat33):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ka(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float)):
+    tid = wp.tid()  # noqa: F841
+
+
+@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()  # noqa: F841
+
+
+@wp.kernel
+def ks0(s: S0):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ksf(s: Sf):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ksv(s: Sv):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ksm(s: Sm):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ksa(s: Sa):
+    tid = wp.tid()  # noqa: F841
+
+
+@wp.kernel
+def ksz(s: Sz):
+    tid = wp.tid()  # noqa: F841
+
+
+wp.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/benchmarks/benchmark_tile_load_store.py

@@ -0,0 +1,103 @@
+# 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.
+
+import numpy as np
+
+import warp as wp
+
+BLOCK_DIM = 128
+
+TILE = 32
+
+
+def create_test_kernel(storage_type: str):
+    @wp.kernel
+    def load_store(a: wp.array2d(dtype=wp.float32), b: wp.array2d(dtype=wp.float32)):
+        i, j = wp.tid()
+
+        if wp.static(storage_type == "shared"):
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="shared")
+        else:
+            a_tile = wp.tile_load(a, shape=(TILE, TILE), offset=(i * TILE, j * TILE), storage="register")
+
+        wp.tile_store(b, a_tile, offset=(i * TILE, j * TILE))
+
+    return load_store
+
+
+if __name__ == "__main__":
+    wp.config.quiet = True
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    iterations = 100
+    rng = np.random.default_rng(42)
+
+    shared_benchmark_data = {}
+    register_benchmark_data = {}
+    memcpy_benchmark_data = {}
+
+    sizes = list(range(128, 4097, 128))
+
+    print(f"{'Transfer Size (Bytes)':<23s} {'Shared (GiB/s)':<16s} {'Register (GiB/s)':<18s} {'memcpy (GiB/s)':<16s}")
+    print("-" * 79)
+
+    for size in sizes:
+        a = wp.array(rng.random((size, size), dtype=np.float32), dtype=wp.float32)
+        b = wp.empty_like(a)
+
+        for storage_type in ("shared", "register"):
+            load_store = create_test_kernel(storage_type)
+
+            cmd = wp.launch_tiled(
+                load_store,
+                dim=(a.shape[0] // TILE, a.shape[1] // TILE),
+                inputs=[a],
+                outputs=[b],
+                block_dim=BLOCK_DIM,
+                record_cmd=True,
+            )
+            # Warmup
+            for _ in range(5):
+                cmd.launch()
+
+            with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_KERNEL, print=False, synchronize=True) as timer:
+                for _ in range(iterations):
+                    cmd.launch()
+
+            np.testing.assert_equal(a.numpy(), b.numpy())
+
+            timing_results = [result.elapsed for result in timer.timing_results]
+            avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+
+            if storage_type == "shared":
+                shared_benchmark_data[a.capacity] = avg_bw
+            else:
+                register_benchmark_data[a.capacity] = avg_bw
+
+        # Compare with memcpy
+        with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_MEMCPY, print=False, synchronize=True) as timer:
+            for _ in range(iterations):
+                wp.copy(b, a)
+
+        timing_results = [result.elapsed for result in timer.timing_results]
+        avg_bw = 2.0 * (a.capacity / (1024 * 1024 * 1024)) / (1e-3 * np.mean(timing_results))
+        memcpy_benchmark_data[a.capacity] = avg_bw
+
+        # Print results
+        print(
+            f"{a.capacity:<23d} {shared_benchmark_data[a.capacity]:<#16.4g} {register_benchmark_data[a.capacity]:<#18.4g} {memcpy_benchmark_data[a.capacity]:<#16.4g}"
+        )

warp/examples/browse.py

@@ -0,0 +1,37 @@
+# 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.
+
+import os
+import subprocess
+import sys
+
+
+def open_file(filename):
+    if sys.platform == "win32":
+        os.startfile(filename)
+    else:
+        subprocess.call(["xdg-open", filename])
+
+
+if __name__ == "__main__":
+    import warp.examples
+
+    source_dir = warp.examples.get_source_directory()
+    print(f"Example source directory: {source_dir}")
+
+    try:
+        open_file(source_dir)
+    except Exception:
+        pass

warp/examples/core/example_cupy.py

@@ -0,0 +1,86 @@
+# 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 CuPy
+#
+# The example demonstrates interoperability with CuPy on CUDA devices
+# and NumPy on CPU devices.
+###########################################################################
+
+import warp as wp
+
+
+@wp.kernel
+def saxpy(x: wp.array(dtype=float), y: wp.array(dtype=float), a: float):
+    i = wp.tid()
+    y[i] = a * x[i] + y[i]
+
+
+class Example:
+    def __init__(self):
+        device = wp.get_device()
+
+        self.n = 10
+        self.a = 1.0
+
+        if device.is_cuda:
+            # use CuPy arrays on CUDA devices
+            import cupy as cp
+
+            print(f"Using CuPy on device {device}")
+
+            # tell CuPy to use the same device
+            with cp.cuda.Device(device.ordinal):
+                self.x = cp.arange(self.n, dtype=cp.float32)
+                self.y = cp.ones(self.n, dtype=cp.float32)
+        else:
+            # use NumPy arrays on CPU
+            import numpy as np
+
+            print("Using NumPy on CPU")
+
+            self.x = np.arange(self.n, dtype=np.float32)
+            self.y = np.ones(self.n, dtype=np.float32)
+
+    def step(self):
+        # Launch a Warp kernel on the pre-allocated arrays.
+        # When running on a CUDA device, these are CuPy arrays.
+        # When running on the CPU, these are NumPy arrays.
+        #
+        # Note that the arrays can be passed to Warp kernels directly.  Under the hood,
+        # Warp uses the __cuda_array_interface__ and __array_interface__ protocols to
+        # access the data.
+        wp.launch(saxpy, dim=self.n, inputs=[self.x, self.y, self.a])
+
+    def render(self):
+        print(self.y)
+
+
+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("--num_frames", type=int, default=10, help="Total number of frames.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example()
+
+        for _ in range(args.num_frames):
+            example.step()
+            example.render()