瓦片

警告

Warp 中的基于瓦片的操作处于预览阶段，API 可能会发生变化。

基于块的编程模型（例如 OpenAI Triton 中的模型）已被证明是表达高性能内核的有效方法，这些内核可以利用现代 GPU 上的协同操作。通过 Warp 1.5.0 [1]，开发人员现在可以在 Warp 内核中访问新的基于瓦片的编程原语。利用 cuBLASDx 和 cuFFTDx，这些新工具为开发人员提供了高效的矩阵乘法和傅立叶变换，用于加速模拟和科学计算。

要求

基于瓦片的操作在针对 CUDA 12 运行时构建的 Warp 版本上得到完全支持。大多数线性代数瓦片操作（如 tile_cholesky()、tile_fft() 和 tile_matmul()）需要 Warp 使用 MathDx 库构建，CUDA 11 不支持该库。 有关本地构建 Warp 以支持线性代数瓦片操作的更多详细信息，请参阅使用 MathDx 构建。

执行模型

Warp 的执行模型允许用户指定最多具有 4 个维度的逻辑线程网格，以便在启动时执行内核。 随着瓦片原语的引入，用户现在可以指定内核启动的块大小，该块大小将线程网格划分为较小的线程集，这些线程集在单个计算单元上执行。

在内核内部，瓦片操作跨每个线程块协同执行，使它们能够利用高效的内存访问、本地内存和专用硬件单元，如 Tensor Cores。

在以下示例中，我们启动一个线程网格，其中每个块负责从 2D 数组加载一行数据并计算其总和

TILE_SIZE = wp.constant(256)
TILE_THREADS = 64

@wp.kernel
def compute(a: wp.array2d(dtype=float), b: wp.array2d(dtype=float)):

    # obtain our block index
    i = wp.tid()

    # load a row from global memory
    t = wp.tile_load(a[i], TILE_SIZE)

    # cooperatively compute the sum of the tile elements; s is a single element tile
    s = wp.tile_sum(t)

    # store s in global memory
    wp.tile_store(b[i], s)

N = 10

a_np = np.arange(N).reshape(-1, 1) * np.ones((1, 256), dtype=float)
a = wp.array(a_np, dtype=float)
b = wp.zeros((N,1), dtype=float)

wp.launch_tiled(compute, dim=[a.shape[0]], inputs=[a, b], block_dim=TILE_THREADS)

print(f"b = {b[:,0]}")

b = [   0.  256.  512.  768. 1024. 1280. 1536. 1792. 2048. 2304.]

在这里，我们使用了新的 warp.launch_tiled() 函数，该函数将 TILE_THREADS 线程分配给启动网格中的每个元素。 然后，每个 TILE_THREADS 线程块从全局内存数组加载一整行 256 个值，并（协同地）计算其总和。

瓦片属性

在 Warp 中，瓦片对象是数据数组，其中瓦片元素可以是标量、向量、矩阵或用户定义的结构。 瓦片最多可以有四个维度。 我们可以直接从 2D 全局内存数组加载 2D 瓦片，如下所示

TILE_M = wp.constant(16)
TILE_N = wp.constant(16)
TILE_THREADS = 64

@wp.kernel
def compute(a: array2d(dtype=float)):

    # obtain our 2d block index
    i, j = wp.tid()

    # load a 2d tile from global memory
    t = wp.tile_load(array, shape=(TILE_M, TILE_N), offset=(i*TILE_M, j*TILE_N))
    s = wp.tile_sum(t)
    ...

wp.launch_tiled(compute, dim=[a.shape[0]/TILE_M, a.shape[1]/TILE_N], inputs=[a], block_dim=TILE_THREADS)

在这里，我们将数组 a 分成形状为 16 x 16 的 2D 瓦片。 每个块协同地从输入数组加载一个瓦片并计算其总和。

瓦片存储

创建瓦片时，它们会放置在寄存器或共享内存中。 通常，Warp 会尝试确定瓦片的最佳存储位置。 默认情况下，瓦片分配在寄存器存储中，但某些操作（如矩阵乘法）可能会根据需要将数据从寄存器迁移到共享存储。

寄存器瓦片

寄存器瓦片中的值存储在整个块中。 例如，如果启动时的块维度设置为 64，则 shape=(1, 256) 的寄存器瓦片将导致每个线程存储 4 个元素。 基于寄存器的存储是大多数硬件上最快的存储，但单个线程无法有效地随机访问分配给另一个线程的数据，因为瓦片存储分散在块中的线程中。 因此，对瓦片的操作倾向于表示为更高级别的映射、缩减和重塑操作，这些操作可能会通过共享内存传输值。

共享内存瓦片

某些操作（如矩阵乘法）需要访问整个瓦片的值。 在这种情况下，瓦片数据可以存储在共享内存中，这允许高效的随机访问。 Warp 将根据需要自动将瓦片迁移到共享内存以进行特定操作。 共享内存是一种有限的资源，因此必须适当地设置瓦片大小，以避免超出硬件限制。 否则，内核编译可能会失败。

示例：通用矩阵乘法 (GEMM)

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=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)):

    # output tile index
    i, j = wp.tid()

    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)

    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)
    C = np.zeros((M, N), dtype=np.float32)

    A_wp = wp.array(A)
    B_wp = wp.array(B)
    C_wp = wp.array(C)

    with wp.Tape() as tape:
        wp.launch_tiled(
            tile_gemm,
            dim=(int(M / TILE_M), int(N / TILE_N)),
            inputs=[A_wp, B_wp, C_wp],
            block_dim=TILE_THREADS)

    assert(np.allclose(C_wp.numpy(), A@B))

    print("Example matrix multiplication passed")

瓦片操作

构造

• tile_zeros()

• tile_ones()

• tile_arange()

• tile()

• untile()

• tile_view()

• tile_broadcast()

加载/存储

• tile_load()

• tile_store()

• tile_atomic_add()

映射/缩减

• tile_map()

• tile_reduce()

• tile_sum()

• tile_min()

• tile_max()

线性代数

• tile_matmul()

• tile_transpose()

• tile_fft()

• tile_ifft()

• tile_cholesky()

• tile_cholesky_solve()

• tile_diag_add()

瓦片和 SIMT 代码

传统上，Warp 内核主要用 SIMT 编程模型编写，其中每个线程的执行都是独立发生的。 另一方面，瓦片允许线程协同工作以执行操作。 Warp 公开了 warp.tile() 和 warp.untile() 方法，用于在每个线程的值类型和等效的瓦片表示形式之间转换数据。 例如

TILE_THREADS = 64

@wp.kernel
def compute():
    i = wp.tid()

    # perform some per-thread computation
    x = i*2.0 + wp.sin(float(i))

    # tile the value x across the block
    # returns a tile with shape=(1, TILE_THREADS)
    t = wp.tile(x)
    ...

# launch as regular SIMT kernel
wp.launch(compute, dim=[N], inputs=[], block_dim=TILE_THREADS)

在此示例中，我们使用 wp.launch() 启动了一个具有 N 个逻辑线程的常规 SIMT 网格。 内核执行一些每个线程的计算，然后使用 warp.tile() 将标量 x 值转换为瓦片对象。 此函数将单个值作为输入，并返回一个瓦片，其维度与块中线程数相同。 从这里，瓦片可以用于其他常规协同操作，如缩减、GEMM 等。

同样，我们可以将瓦片对象取消瓦片化回其每个线程的标量等效值。

注意

块中的所有线程都必须执行瓦片操作，但瓦片操作周围的代码可能包含任意条件逻辑。

示例：使用瓦片加速数组范围内的缩减

在 Warp 中添加瓦片支持之前，数组范围内的缩减通常在单个内核中使用内置原子函数（如 wp.atomic_add()）执行。 与 cub::BlockReduce 等优化机制相比，这可能非常低效。 考虑对数组进行以下平方和缩减

import  numpy as np

import warp as wp

@wp.kernel
def reduce_array_simt(
    a: wp.array2d(dtype=wp.float64),
    result: wp.array(dtype=wp.float64),
):
    i, j = wp.tid()

    local_val = a[i, j]*a[i, j]

    wp.atomic_add(result, 0, local_val)

rng = np.random.default_rng(42)

data = wp.array(rng.random((4096, 4096), dtype=np.float64), dtype=wp.float64)
result = wp.zeros((1,), dtype=wp.float64)

wp.launch(reduce_array_simt, (4096, 4096), inputs=[data], outputs=[result])

以上内核在 Warp 1.6.1 中在 RTX 3090 GPU 上运行大约需要 27.5 毫秒。 我们可以使用瓦片来加速此缩减，首先从标量 local_val 创建一个瓦片，然后使用 wp.tile_sum() 使用共享内存协同计算瓦片总和。 然后，我们可以使用 wp.tile_atomic_add() 将缩减瓦片的结果累积到全局内存中

import  numpy as np

import warp as wp

@wp.kernel
def reduce_array_tile(
    a: wp.array2d(dtype=wp.float64),
    result: wp.array(dtype=wp.float64),
):
    i, j = wp.tid()

    local_val = a[i, j]*a[i, j]

    t = wp.tile(local_val)
    s = wp.tile_sum(t)

    wp.tile_atomic_add(result, s)

rng = np.random.default_rng(42)

data = wp.array(rng.random((4096, 4096), dtype=np.float64), dtype=wp.float64)
result = wp.zeros((1,), dtype=wp.float64)

wp.launch(reduce_array_tile, (4096, 4096), inputs=[data], outputs=[result])

使用瓦片的缩减内核运行大约需要 0.528 毫秒，比原始内核提高了 52 倍。

可以通过试验不同的块大小来进一步提高速度。 如果我们通过将 block_dim=128 添加到 wp.launch() 中，将此示例中的块大小从默认的 256 个线程减少到 128 个线程，则内核仅需要大约 0.436 毫秒即可完成，而纯 SIMT 内核相对不受影响。

自动微分

Warp 可以自动生成基于瓦片的程序的向后版本。 通常，瓦片程序必须遵守与常规 Warp 程序的自动微分相同的规则，例如避免原地操作等。 有关更多详细信息，请参阅 可微性 部分。

使用 MathDx 构建

线性代数中描述的大多数瓦片操作都需要 Warp 使用 MathDx 库构建。 从 Warp 1.5.0 开始，PyPI 分发将开箱即用地支持利用 MathDx API 的瓦片操作。

当使用 build_lib.py 在本地构建 Warp 时，该脚本将尝试自动从 cuBLASDx 下载页面 下载 libmathdx。也可以在使用 build_lib.py 时使用 --libmathdx_path 选项，或者在 LIBMATHDX_HOME 环境变量中定义路径，来指定现有 libmathdx 安装的路径。

[1]

技术博客：在 Warp 1.5.0 中引入基于 Tile 的编程