warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/examples/benchmarks/benchmark_cloth_taichi.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 taichi as ti
+
+
+@ti.func
+def step(x):
+    ret = 0.0
+    if x < 0:
+        ret = 1
+    return ret
+
+
+@ti.data_oriented
+class TiIntegrator:
+    @ti.kernel
+    def eval_springs(self):
+        for tid in range(self.cloth.num_springs):
+            i = self.spring_indices[2 * tid]
+            j = self.spring_indices[2 * tid + 1]
+
+            ke = self.spring_stiffness[tid]
+            kd = self.spring_damping[tid]
+            rest = self.spring_lengths[tid]
+
+            xi = self.positions[i]
+            xj = self.positions[j]
+
+            vi = self.velocities[i]
+            vj = self.velocities[j]
+
+            xij = xi - xj
+            vij = vi - vj
+
+            l = xij.norm()
+            dir = xij.normalized()
+
+            c = l - rest
+            dcdt = dir.dot(vij)
+
+            fs = dir * (ke * c + kd * dcdt)
+
+            self.forces[i] -= fs
+            self.forces[j] += fs
+
+    @ti.kernel
+    def integrate_particles(self, dt: ti.f32):
+        for tid in range(self.cloth.num_particles):
+            x0 = self.positions[tid]
+            v0 = self.velocities[tid]
+            f0 = self.forces[tid]
+            w = self.inv_mass[tid]
+
+            g = ti.Vector([0.0, 0.0, 0.0])
+
+            if w > 0.0:
+                g = ti.Vector([0.0, -9.81, 0.0])
+
+            v1 = v0 + (f0 * w + g) * dt
+            x1 = x0 + v1 * dt
+
+            self.positions[tid] = x1
+            self.velocities[tid] = v1
+            self.forces[tid] = ti.Vector([0.0, 0.0, 0.0])
+
+    def __init__(self, cloth, device):
+        if device == "cpu":
+            ti.init(arch=ti.cpu)
+        elif device == "cuda":
+            ti.init(arch=ti.gpu)
+        else:
+            raise RuntimeError("Unsupported Taichi device")
+
+        self.cloth = cloth
+
+        self.positions = ti.Vector.field(3, dtype=ti.f32, shape=self.cloth.num_particles)
+        self.velocities = ti.Vector.field(3, dtype=ti.f32, shape=self.cloth.num_particles)
+        self.inv_mass = ti.field(ti.f32, shape=self.cloth.num_particles)
+
+        self.spring_indices = ti.field(ti.i32, shape=self.cloth.num_springs * 2)
+        self.spring_lengths = ti.field(ti.f32, shape=self.cloth.num_springs)
+        self.spring_stiffness = ti.field(ti.f32, shape=self.cloth.num_springs)
+        self.spring_damping = ti.field(ti.f32, shape=self.cloth.num_springs)
+
+        self.forces = ti.Vector.field(3, dtype=ti.f32, shape=self.cloth.num_particles)
+
+        # upload data
+        self.positions.from_numpy(cloth.positions)
+        self.velocities.from_numpy(cloth.velocities)
+        self.inv_mass.from_numpy(cloth.inv_masses)
+        self.forces.from_numpy(np.zeros_like(self.cloth.velocities))
+
+        self.spring_indices.from_numpy(cloth.spring_indices)
+        self.spring_lengths.from_numpy(cloth.spring_lengths)
+        self.spring_stiffness.from_numpy(cloth.spring_stiffness)
+        self.spring_damping.from_numpy(cloth.spring_damping)
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for _s in range(substeps):
+            self.eval_springs()
+
+            self.integrate_particles(sim_dt)
+
+        return self.positions.to_numpy()

warp/examples/benchmarks/benchmark_cloth_warp.py

@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 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 warp as wp
+
+wp.clear_kernel_cache()
+
+
+@wp.kernel
+def eval_springs(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    spring_indices: wp.array(dtype=int),
+    spring_rest_lengths: wp.array(dtype=float),
+    spring_stiffness: wp.array(dtype=float),
+    spring_damping: wp.array(dtype=float),
+    f: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+
+    i = spring_indices[tid * 2 + 0]
+    j = spring_indices[tid * 2 + 1]
+
+    ke = spring_stiffness[tid]
+    kd = spring_damping[tid]
+    rest = spring_rest_lengths[tid]
+
+    xi = x[i]
+    xj = x[j]
+
+    vi = v[i]
+    vj = v[j]
+
+    xij = xi - xj
+    vij = vi - vj
+
+    l = wp.length(xij)
+    l_inv = 1.0 / l
+
+    # normalized spring direction
+    dir = xij * l_inv
+
+    c = l - rest
+    dcdt = wp.dot(dir, vij)
+
+    # damping based on relative velocity.
+    fs = dir * (ke * c + kd * dcdt)
+
+    wp.atomic_sub(f, i, fs)
+    wp.atomic_add(f, j, fs)
+
+
+@wp.kernel
+def integrate_particles(
+    x: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    f: wp.array(dtype=wp.vec3),
+    w: wp.array(dtype=float),
+    dt: float,
+):
+    tid = wp.tid()
+
+    x0 = x[tid]
+    v0 = v[tid]
+    f0 = f[tid]
+    inv_mass = w[tid]
+
+    g = wp.vec3()
+
+    # treat particles with inv_mass == 0 as kinematic
+    if inv_mass > 0.0:
+        g = wp.vec3(0.0, 0.0 - 9.81, 0.0)
+
+    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
+    v1 = v0 + (f0 * inv_mass + g) * dt
+    x1 = x0 + v1 * dt
+
+    x[tid] = x1
+    v[tid] = v1
+
+    # clear forces
+    f[tid] = wp.vec3()
+
+
+class WpIntegrator:
+    def __init__(self, cloth, device):
+        self.device = wp.get_device(device)
+
+        with wp.ScopedDevice(self.device):
+            self.positions = wp.from_numpy(cloth.positions, dtype=wp.vec3)
+            self.positions_host = wp.from_numpy(cloth.positions, dtype=wp.vec3, device="cpu")
+            self.invmass = wp.from_numpy(cloth.inv_masses, dtype=float)
+
+            self.velocities = wp.zeros(cloth.num_particles, dtype=wp.vec3)
+            self.forces = wp.zeros(cloth.num_particles, dtype=wp.vec3)
+
+            self.spring_indices = wp.from_numpy(cloth.spring_indices, dtype=int)
+            self.spring_lengths = wp.from_numpy(cloth.spring_lengths, dtype=float)
+            self.spring_stiffness = wp.from_numpy(cloth.spring_stiffness, dtype=float)
+            self.spring_damping = wp.from_numpy(cloth.spring_damping, dtype=float)
+
+        self.cloth = cloth
+
+    def simulate(self, dt, substeps):
+        sim_dt = dt / substeps
+
+        for _s in range(substeps):
+            wp.launch(
+                kernel=eval_springs,
+                dim=self.cloth.num_springs,
+                inputs=[
+                    self.positions,
+                    self.velocities,
+                    self.spring_indices,
+                    self.spring_lengths,
+                    self.spring_stiffness,
+                    self.spring_damping,
+                    self.forces,
+                ],
+                outputs=[],
+                device=self.device,
+            )
+
+            # integrate
+            wp.launch(
+                kernel=integrate_particles,
+                dim=self.cloth.num_particles,
+                inputs=[self.positions, self.velocities, self.forces, self.invmass, sim_dt],
+                outputs=[],
+                device=self.device,
+            )
+
+        # copy data back to host
+        if self.device.is_cuda:
+            wp.copy(self.positions_host, self.positions)
+            wp.synchronize()
+
+            return self.positions_host.numpy()
+
+        else:
+            return self.positions.numpy()

warp/examples/benchmarks/benchmark_gemm.py

@@ -0,0 +1,164 @@
+# 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.
+
+"""Compare GEMM performance between Torch and Warp (Tiled).
+
+This script can be used to identify optimal tile parameters for a fixed-size
+matrix multiplication.
+"""
+
+from itertools import product
+from statistics import mean, stdev
+from typing import List
+
+import numpy as np
+import torch
+
+import warp as wp
+
+
+# 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 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)
+
+        count = A.shape[1] // TILE_K
+
+        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: torch.Tensor, B: torch.Tensor, warm_up: int, iterations: int):
+    # warm-up
+    for _ in range(warm_up):
+        torch.matmul(A, B)
+
+    torch.cuda.synchronize()
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    timing_results = []
+
+    for _i in range(iterations):
+        start_event.record()
+        torch.matmul(A, B)
+        end_event.record()
+
+        torch.cuda.synchronize()
+        timing_results.append(start_event.elapsed_time(end_event))
+
+    return mean(timing_results), stdev(timing_results)
+
+
+def benchmark_warp(A: wp.array, B: wp.array, config: List[int], warm_up: int, iterations: int):
+    TILE_M = config[0]
+    TILE_N = config[1]
+    TILE_K = config[2]
+    BLOCK_DIM = config[3]
+
+    mlp = create_gemm_kernel(TILE_M, TILE_N, TILE_K)
+
+    M = A.shape[0]
+    N = B.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):
+        cmd.launch()
+
+    # check output
+    if warm_up > 0:
+        try:
+            np.testing.assert_allclose(output.numpy(), A.numpy() @ B.numpy(), atol=1e-3, rtol=1e-3)
+        except AssertionError as e:
+            print(f"Failed with {TILE_M=}, {TILE_N=}, {TILE_K=}, {BLOCK_DIM=}")
+            raise e
+
+    # benchmark
+    with wp.ScopedTimer("warp", print=False, synchronize=True, cuda_filter=wp.TIMING_KERNEL) as timer:
+        for _ in range(iterations):
+            cmd.launch()
+
+    timing_results = [result.elapsed for result in timer.timing_results]
+
+    return mean(timing_results), stdev(timing_results)
+
+
+if __name__ == "__main__":
+    torch.backends.cuda.matmul.allow_tf32 = False  # Disable TF32 for matrix multiplications
+    torch.backends.cudnn.allow_tf32 = False  # Disable TF32 for cuDNN operations
+
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    tile_m = [8, 16, 32, 64]
+    tile_n = [8, 16, 32, 64]
+    tile_k = [8, 16, 64]
+    block = [32, 64, 128]
+
+    M = 1024
+    N = 1024
+    K = 1024
+    print(f"{M=}, {N=}, {K=}")
+
+    A = torch.randn(M, K).cuda()
+    B = torch.randn(K, N).cuda()
+
+    iterations = 100
+    warm_up = 5
+
+    time_torch_mean, time_torch_std = benchmark_torch(A, B, warm_up, iterations)
+    print(f"Torch: {time_torch_mean:.6g}±{time_torch_std:.2g} ms")
+
+    configs = list(product(tile_m, tile_n, tile_k, block))
+
+    wp.config.quiet = True
+
+    # header
+    print(
+        f"{'TILE_M':<8s} {'TILE_N':<8s} {'TILE_K':<8s} {'BLOCK':<8s} {'Time (ms)':<10s} {'Std dev (ms)':<14s} {'Warp/Torch':<12s}"
+    )
+    print("-" * 79)
+
+    for c in configs:
+        time_mean, time_std = benchmark_warp(wp.from_torch(A), wp.from_torch(B), c, warm_up, iterations)
+        print(
+            f"{c[0]:<8d} {c[1]:<8d} {c[2]:<8d} {c[3]:<8d} {time_mean:<10.6g} {time_std:<#14.2g} {time_mean / time_torch_mean:<12.6g}"
+        )

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)