【CUDA编程】CUDA Warp 与 Warp-Python 学习文档

__global__ void add_vectors(float *a, float *b, float *c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) c[idx] = a[idx] + b[idx];  // 保护尾部线程，避免越界
}

__device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

python -m pip install warp-lang

import warp as wp
print(wp.__version__)  # 输出版本号（如 1.7.0）

import warp as wp
import numpy as np

# 初始化上下文并明确默认设备；验证后可改为 cuda:0。
wp.init()
wp.set_device("cpu")
rng = np.random.default_rng(42)

# 在选定设备上定义数组
n = 1024
a = wp.array(rng.random(n), dtype=wp.float32)
b = wp.array(rng.random(n), dtype=wp.float32)
c = wp.zeros(n, dtype=wp.float32)

# 定义 GPU 内核（@wp.kernel 装饰器）
@wp.kernel
def add_vectors(a: wp.array(dtype=wp.float32),
                b: wp.array(dtype=wp.float32),
                c: wp.array(dtype=wp.float32)):
    i = wp.tid()  # 获取线程索引
    c[i] = a[i] + b[i]

# 启动内核（指定线程数）
wp.launch(kernel=add_vectors, dim=n, inputs=[a, b, c])

# 将结果拷贝回 CPU
result = c.numpy()
np.testing.assert_allclose(result, a.numpy() + b.numpy(), rtol=1e-6)
print("vector addition passed on", c.device)

@wp.kernel
def matrix_mult(a: wp.array2d(dtype=wp.float32),
                b: wp.array2d(dtype=wp.float32),
                c: wp.array2d(dtype=wp.float32)):
    i, j = wp.tid()
    c[i, j] = 0.0
    for k in range(a.shape[1]):
        c[i, j] += a[i, k] * b[k, j]

# 定义矩阵
a = wp.array(rng.random((64, 64)), dtype=wp.float32)
b = wp.array(rng.random((64, 64)), dtype=wp.float32)
c = wp.zeros((64, 64), dtype=wp.float32)

# 执行矩阵乘法
wp.launch(matrix_mult, dim=(64, 64), inputs=[a, b, c])
np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-5)

@wp.kernel
def atomic_add_example(counter: wp.array(dtype=wp.int32)):
    wp.atomic_add(counter, 0, 1)  # 原子加操作

counter = wp.zeros(1, dtype=wp.int32)
wp.launch(atomic_add_example, dim=100, inputs=[counter])
print(counter.numpy())  # 输出 [100]

@wp.kernel
def update_particles(positions: wp.array(dtype=wp.vec3),
                     velocities: wp.array(dtype=wp.vec3),
                     dt: float):
    tid = wp.tid()
    velocities[tid] += wp.vec3(0.0, -9.8, 0.0) * dt  # 重力加速度
    positions[tid] += velocities[tid] * dt

# 初始化粒子
num_particles = 1000
positions = wp.array(rng.random((num_particles, 3)), dtype=wp.vec3)
velocities = wp.zeros(num_particles, dtype=wp.vec3)

# 模拟多步
for _ in range(100):
    wp.launch(update_particles, dim=num_particles, inputs=[positions, velocities, 0.01])

@wp.kernel
def loss_function(x: wp.array(dtype=wp.float32),
                  y: wp.array(dtype=wp.float32),
                  loss: wp.array(dtype=wp.float32)):
    tid = wp.tid()
    difference = x[tid] - y[tid]
    wp.atomic_add(loss, 0, difference * difference)

# kernel 通过输出数组写结果；wp.launch 不返回一个标量损失。
x = wp.array(rng.random(10), dtype=wp.float32, requires_grad=True)
y = wp.zeros(10, dtype=wp.float32, device=x.device)
loss = wp.zeros(1, dtype=wp.float32, device=x.device, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(loss_function, dim=10, inputs=[x, y], outputs=[loss],
              device=x.device)

tape.backward(loss=loss)
np.testing.assert_allclose(x.grad.numpy(), 2.0 * x.numpy(), rtol=1e-5)
print(x.grad.numpy())

import warp as wp
from warp import float32 as f32
import numpy as np
import cv2

# 设置画布大小
n = 800
pixel = wp.zeros((n, n), dtype=f32, device='cuda:0')

@wp.func
def mandelbrot_func(z: wp.vec2, c: wp.vec2) -> wp.vec2:
    return wp.vec2(z[0] * z[0] - z[1] * z[1] + c[0],
                   2.0 * z[0] * z[1] + c[1])

@wp.kernel
def paint(p: wp.array2d(dtype=f32), t: f32):
    i, j = wp.tid()

    # 动态缩放和平移
    zoom = 2.8 + wp.sin(t * 0.2) * 0.5
    center_x = -0.5 + wp.cos(t * 0.1) * 0.1
    center_y = wp.sin(t * 0.15) * 0.1

    x = (f32(j) / f32(n) - 0.5) * zoom + center_x
    y = (f32(i) / f32(n) - 0.5) * zoom + center_y

    c = wp.vec2(x, y)
    z = wp.vec2(0.0, 0.0)

    iteration = f32(0.0)
    max_iter = f32(200.0)

    # 迭代计算
    while wp.length(z) < 2.0 and iteration < max_iter:
        z = mandelbrot_func(z, c)
        iteration += 1.0

    # 平滑着色
    smooth_iter = max_iter
    if wp.length(z) > 2.0:
        smooth_iter = iteration + 1.0 - wp.log2(wp.log2(wp.length(z)))
    p[i, j] = smooth_iter / max_iter

def main():
    t = 0.0
    while True:
        wp.launch(paint, dim=(n, n), inputs=[pixel, t], device='cuda:0')

        # 创建彩色效果
        np_pixel = pixel.numpy()
        # 使用更丰富的颜色映射
        colored = cv2.applyColorMap(
            (np_pixel * 255).astype(np.uint8),
            cv2.COLORMAP_MAGMA
        )

        cv2.imshow("Mandelbrot Set", colored)

        t += 0.01
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

if __name__ == "__main__":
    main()

@wp.kernel
def particle_update(
    positions: wp.array(dtype=wp.vec3),
    velocities: wp.array(dtype=wp.vec3),
    dt: float
):
    tid = wp.tid()
    # Warp 级优化：同一 Warp 内的线程处理连续数据
    # 线程索引分组本身不构成同步；每个线程独立更新自己的粒子。
    velocities[tid] += wp.vec3(0, -9.8, 0) * dt
    positions[tid] += velocities[tid] * dt