warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/benchmarks/benchmark_interop_paddle.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 paddle
+
+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_paddle(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    paddle_device = wp.device_to_paddle(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        paddle_shape = (array_size, *warp_dtype._shape_)
+        paddle_dtype = wp.dtype_to_paddle(warp_dtype._wp_scalar_type_)
+    else:
+        paddle_shape = (array_size,)
+        paddle_dtype = paddle.float32 if warp_dtype is None else wp.dtype_to_paddle(warp_dtype)
+
+    _a = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _b = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _c = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _d = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _e = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_paddle(_a, dtype=warp_dtype)
+        b = wp.from_paddle(_b, dtype=warp_dtype)
+        c = wp.from_paddle(_c, dtype=warp_dtype)
+        d = wp.from_paddle(_d, dtype=warp_dtype)
+        e = wp.from_paddle(_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_paddle(...)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_array_ctype_from_paddle(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    paddle_device = wp.device_to_paddle(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        paddle_shape = (array_size, *warp_dtype._shape_)
+        paddle_dtype = wp.dtype_to_paddle(warp_dtype._wp_scalar_type_)
+    else:
+        paddle_shape = (array_size,)
+        paddle_dtype = paddle.float32 if warp_dtype is None else wp.dtype_to_paddle(warp_dtype)
+
+    _a = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _b = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _c = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _d = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _e = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+
+    wp.synchronize()
+
+    # profiler = Profiler(interval=0.000001)
+    # profiler.start()
+
+    t1 = time.time_ns()
+
+    for _ in range(num_iters):
+        a = wp.from_paddle(_a, dtype=warp_dtype, return_ctype=True)
+        b = wp.from_paddle(_b, dtype=warp_dtype, return_ctype=True)
+        c = wp.from_paddle(_c, dtype=warp_dtype, return_ctype=True)
+        d = wp.from_paddle(_d, dtype=warp_dtype, return_ctype=True)
+        e = wp.from_paddle(_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_paddle(..., return_ctype=True)")
+
+    # profiler.stop()
+    # profiler.print()
+
+
+def test_direct_from_paddle(kernel, num_iters, array_size, device, warp_dtype=None):
+    warp_device = wp.get_device(device)
+    paddle_device = wp.device_to_paddle(warp_device)
+
+    if hasattr(warp_dtype, "_shape_"):
+        paddle_shape = (array_size, *warp_dtype._shape_)
+        paddle_dtype = wp.dtype_to_paddle(warp_dtype._wp_scalar_type_)
+    else:
+        paddle_shape = (array_size,)
+        paddle_dtype = paddle.float32 if warp_dtype is None else wp.dtype_to_paddle(warp_dtype)
+
+    _a = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _b = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _c = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _d = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_device)
+    _e = paddle.zeros(paddle_shape, dtype=paddle_dtype).to(device=paddle_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 paddle")
+
+    # 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_paddle(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_array_ctype_from_paddle(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)
+    test_direct_from_paddle(kernel, num_iters, 10, "cuda:0", warp_dtype=warp_dtype)

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("--------------------------------")