Related Resources
模型/库/框架

Warp 1.5.0 引入图块化编程

2024年 12月 14日
By , and

借助最新版本的 Warp 1.5.0 ，开发者现在可以使用 Python 中基于图块的新编程基元。这些新工具利用 cuBLASDx 和 cuFFTDx ，在 Python 内核中为开发者提供高效的矩阵乘法和 Fourier 变换，从而加速仿真和科学计算。在这篇博文中，我们将介绍这些新功能，并展示如何使用它们来优化应用。Warp 1.5.0 中提供的基于图块的编程模型目前处于预览阶段，在即将推出的版本中，性能和 APIs 可能会发生变化。

简介 

在过去十年中，GPU 硬件已从单纯的 SIMT (单指令多线程) 执行模型发展为高度依赖协作操作来提高效率的模型。随着 Tensor Core 数学单元在整体 GPU 计算中的作用越来越大，高效且高效的编程变得越来越重要。高级 API 如 BLAS 提供的抽象概念可以面向各种高性能低级指令。但是，这些 API 通常难以与用户程序集成，并且可能会在库调用之间强制将结果传回全局内存，从而降低效率。相反，直接在 C++/CUDA 级别对 Tensor Core 进行编程非常复杂，需要仔细管理单元之间的数据流。

为解决这些问题，我们开发了基于图块的编程模型，例如 OpenAI Triton 和 C++ AMP 中模型。与纯 SIMT 模型不同，基于图块的编程允许开发者在图块上表示多个线程可以协同执行的操作，从而提高效率和生产力。

在 Warp 1.5.0 版本中，我们扩展了 Warp 基于内核的编程模型，将基于平铺的操作包括在内，旨在让 Warp 开发者能够充分利用现代 GPU 硬件的全部功能。这些扩展程序：

  • 提供编程模型，使开发者能够从 SIMT 顺利过渡到基于图块的执行。
  • 减少对手动索引、共享内存管理和指针运算的需求。
  • 支持反向传播和训练的自动微分。

此外，Warp 还利用 cuBLASDx 和 cuFFTDx 执行矩阵乘法和 Fast Fourier Transform (FFT) 图块运算。这些 NVIDIA 设备端数学库与 Warp 的图块编程模型相结合，可在单个内核中无缝融合 Tensor Core 加速的 GEMM、FFT 和其他图块运算，从而减少内存 I/O 和内核启动用度，同时更大限度地提高算术强度。借助这种方法，对于需要密集线性代数的应用 (例如机器人前向动力学)，我们的性能可以比传统的线性代数或张量框架高出 4 倍。

扭曲图块基元 

Warp 中的新平铺基元包括构建、加载/存储、线性代数和映射/归约运算，从而自然扩展了现有基于内核的编程模型。

建筑 

可以使用 NumPy 式操作在 Warp 内核中构建图块，如下所示：

import warp as wp

@wp.kernel
def compute():


    # construct a 16x16 tile of zeroed 32-bit floats
    a = wp.tile_zeros(m=16, n=16, dtype=wp.float32)


    # construct a 16x16 tile of 16-bit floats initialized to 1.0
    b = wp.tile_ones(m=16, n=16, dtype=wp.float16)

在 Warp 中，tiles 是二维数组，可能包含标量、向量、矩阵或结构化数据类型作为元素。与 Warp 数组或 PyTorch 张量不同，其中的维度是动态的并在运行时指定，tile 的维度 (例如，上述示例中的 16×16) 必须是编译时已知的常量。此外，与线程的本地 SIMT 数据不同，tile 数据存储在整个 CUDA 块中，既可以存储在寄存器中，也可以存储在共享内存中。有关 tile 构建例程的完整列表，请单击此处访问 GitHub。

加载/存储 

Warp 为全局内存中的平铺数据提供显式加载/存储操作。块中的所有线程协同执行这些操作，确保全局内存与共享或寄存器内存之间的高效数据传输。在以下示例中，系统从全局内存加载两个数据图块，并对其求和，然后将结果存储回全局内存。用户无需明确管理共享内存分配或存储。

import warp as wp

@wp.kernel
def compute(A: wp.array2d(dtype=float),
            B: wp.array2d(dtype=float),
            C: wp.array2d(dtype=float)):
    
    # cooperatively load input tiles
    a = wp.tile_load(A, i=0, j=0, m=16, n=16)
    b = wp.tile_load(B, i=0, j=0, m=16, n=16)

    # compute sum
    c = a + b

    # cooperatively store sum to global memory
    wp.tile_store(C, i=0, j=0, t=c)

A = wp.ones((16,16), dtype=float)
B = wp.ones((16,16), dtype=float)
C = wp.empty((16,16), dtype=float)

wp.launch_tiled(compute, dim=1, inputs=[A, B, C], device="cuda:0", block_dim=64)

除了加载/存储操作外，Warp 还支持 wp.tile_atomic_add() 等原子操作。有关内存操作的完整列表，请参阅以下 文档 。

矩阵乘法 

基于图块的编程的主要优势之一是能够执行协作矩阵乘法。Warp 1.5.0 引入了通用乘积累加基元 wp.tile_matmul()，允许开发者执行协同矩阵乘法。这在幕后利用 cuBLASDx ，根据元素类型、矩阵大小和数据布局，cuBLASDx 将自动发送适当的 Tensor Core MMA 指令，以实现最佳性能。

我们来看看在 Warp 中使用基于图块的编程来执行矩阵乘法的示例：

import warp as wp

TILE_M = wp.constant(32)
TILE_N = wp.constant(64)
TILE_K = wp.constant(64)

@wp.kernel
def gemm_tiled(A: wp.array2d(dtype=float), 
               B: wp.array2d(dtype=float), 
               C: wp.array2d(dtype=float)):

    i, j = wp.tid()

    # allocate output tile
    sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=float)
    count = int(K / TILE_K)

    # iterate over inner dimension
    for k in range(count):
        a = wp.tile_load(A, i, k, m=TILE_M, n=TILE_K)
        b = wp.tile_load(B, k, j, m=TILE_K, n=TILE_N)
 
        # perform gemm + accumulate
        wp.tile_matmul(a, b, sum)

    # store result
    wp.tile_store(C, i, j, sum)


# test with 1024^2 inputs
M, N, K = 1024, 1024, 1024

A = wp.ones((M, K), dtype=float)
B = wp.ones((K, N), dtype=float)
C = wp.empty((M, N), dtype=float)

# launch kernel with 128 threads per-block
wp.launch_tiled(gemm_tiled, 
                dim=(int(M//TILE_M), int(N//TILE_N)), 
                inputs=[A, B, C],
                block_dim=128)

在本示例中，我们定义了执行平铺矩阵乘法的核 gemm_tiled()。内核循环处理全局内存中的 2D 数据片，将其加载到共享内存图块中，使用 wp.tile_matmul() 执行矩阵乘法，在共享内存中累加结果，然后将结果存储回全局内存。

下图显示了上述 GEMM 内核在 NVIDIA A100 80GB SXM（时钟锁定至其最大值）上与 cuBLAS 12.4 相比在各种 FP32 矩阵大小下的性能百分比。对于小问题，我们发现性能与 cuBLAS 具有竞争力，这可能是因为我们使用自动调整来为这种小尺寸找到最佳参数，而且启动开销在成本中占更大的比例。对于更大的问题，由于目前 tile 结果始终存储在共享内存中，因此性能会降低。然而，即使对于这个简单的示例，我们也会看到较大矩阵的 cuBLAS 性能约为 70–80%。未来版本的 Warp 和 cuBLASDx 将在寄存器中保留 GEMM 的输出，从而提高性能。

A barchart showing GEMM performance as a percentage of cuBLAS
图 1、GEMM 在各种问题规模下的性能占 cuBLAS 的百分比

下图中，我们看看图块大小对单个问题大小的整体性能的影响。整体性能是图块维度和块维度的函数，前者决定问题的分解方式，后者决定向每个子问题分配多少线程。在这里，我们可以看到，对于 M = N = K = 1024 问题，使用 TILE_M = 32、TILE_N = 64、TILE_K = 64 和 128 线程的平铺维度获得了最佳性能。Warp 的动态编程和运行时内核创建允许用户轻松执行超参数自动调整，如基准脚本示例所示。

A barchart showing GEMM performance for different tile and block sizes.
图 2、M=N=K=1024 下不同 tile 和 block 大小的 GEMM 性能

请参阅此 参考 资料，了解平铺线性代数基元的完整列表。

映射/归约 

Warp 1.5.0 还包含 map/reduce 基元，支持开发者在图块上执行归约和元素级运算。这些基元对于 LayerNorm 和 SoftMax 等任务至关重要，这些任务需要对不同数量进行高效归约。

以下示例展示了如何使用每行一个 CUDA 块和 wp.tile_sum() 计算数组行中所有元素的总和，以执行协作归约。

import warp as wp

@wp.kernel
def row_sum(input: wp.array2d(dtype=float),
            output: wp.array1d(dtype=float)):

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

    # load a row of 256 elements from global memory
    t = wp.tile_load(input[i], i=0, n=256)
    # cooperatively sum elements
    s = wp.tile_sum(t)
    # store sum to output
    wp.tile_store(output, i, s)

   
a = wp.ones((1024, 256), dtype=float)
b = wp.empty(1024, dtype=float)

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

Warp 还支持自定义归约运算符，在本例中，我们使用内置的 wp.tile_reduce() 和 wp.mul() 来计算因子，不过也可以使用用户定义的 @wp.func 归约运算符。

import warp as wp

@wp.kernel
def factorial():
    t = wp.tile_arange(1, 10, dtype=int)
    s = wp.tile_reduce(wp.mul, t)
    
    # prints "tile(m=1, n=1, storage=register) = [[362880]]"
    print(s)


wp.launch(factorial, dim=[16], inputs=[], block_dim=16)

此处 提供了 map/reduce 基元的完整列表。

An original image of 2D fluid simulation
Output of a 2D simulation after one minute of training
图 3、示例：使用小型融合 MLP 网络对 2D 流体模拟图像进行近似计算。上图：原始图像，下图：经过大约一分钟训练后的 4 层 MLP 输出。

案例研究 

融合神经网络 

基于图块的编程还支持高效实现融合多层感知器（MLP）。以下是在 Warp 中使用基于图块的编程的融合 MLP 示例：

import warp as wp


DIM_IN = wp.constant(4)
DIM_HID = wp.constant(32)
DIM_OUT = wp.constant(3)

@wp.kernel
def mlp_fused(weights_0: wp.array2d(dtype=wp.float16), 
              weights_1: wp.array2d(dtype=wp.float16), 
              loss: wp.array(dtype=float)):


    t = wp.tid()    

    # construct simple positional encoding
    x = wp.vec4h(wp.sin(x), 
                 wp.cos(x),
                 wp.sin(x*2.0),
                 wp.cos(x*2.0))

    # tile input across block to create feature vectors
    f = wp.tile(x)
    
    # fully connected layer 0
    w0 = wp.tile_load(weights_0, 0, 0, m=DIM_HID, n=DIM_IN)
    z = wp.tile_map(relu, wp.tile_matmul(w0, f))

    # fully connected layer 1
    w1 = wp.tile_load(weights_1, 0, 0, m=DIM_OUT, n=DIM_HID)
    z = wp.tile_map(relu, wp.tile_matmul(w1, z))

    # loss function
    l = wp.tile_sum(z)

    wp.atomic_add(loss, 0, l)

wp.launch(mlp_fused, dim=(1,), inputs=[weights_0, weights_1, loss], block_dim=128)

在此示例中 ，mlp_fused() 内核通过加载权重、执行矩阵乘法、使用 wp.tile_map() 应用激活函数以及计算损失来评估简单的两层神经网络，所有这些操作均在单个内核中完成。下图展示了使用此方法对图像进行编码的示例。由于 Warp 支持自动微分，因此我们可以直接评估和训练网络权重，以学习从图像坐标 (x,y) 映射到像素颜色 (RGB) 的函数。完整示例请参见此处。

信号处理 

线程束平铺运算集成 cuFFTDx，用于内核内的正向和反向 FFT，在数据块上提供高效的 Fourier-transform 运算。以下是在 Warp 中使用基于图块的 FFT 来使用滤波器计算卷积的示例：

import warp as wp

@wp.kernel
def conv_tiled(x: wp.array2d(dtype=wp.vec2d), 
  y: wp.array2d(dtype=wp.vec2d), 
  z: wp.array2d(dtype=wp.vec2d)):
i, j = wp.tid()

# load signal and filter
a = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)
f = wp.tile_load(y, i, j, m=TILE_M, n=TILE_N)

# compute Fourier transform of input signal
wp.tile_fft(a)

# compute filter in frequency space
c = wp.tile_map(cplx_prod, a, b)

# convert back to real
wp.tile_ifft(c)
wp.tile_store(z, i, j, c)

在此示例中，conv_tiled() 内核 (沿最后一个维度) 执行一块数据的前向 FFT，应用过滤器，然后计算反向 FFT。在幕后，cuFFTDx 用于实现。完整示例可在此处找到。下图显示了对噪声输入信号应用滤波器的输出。

A graph showing the output of applying the filter using Warp's FFT tile operations.
图 4、使用 Warp 的 FFT 平铺操作进行 1D 信号过滤

机器人前向动力学 

An image of quadruped robots in simulation.
图 5、图片来源：ETH 机器人系统实验室 (RSL)

基于图块的编程也非常有利于需要密集线性代数的模拟应用。在机器人仿真中，复合刚体算法 (CRBA) 方法用于计算关节机制的前向动力学。在 CRBA 方法中，需要以下三矩阵乘积，其中内矩阵 M 是块稀疏对角线质量矩阵：

A block-sparse matrix showing the CRBA method
图 6、CRBA 方法所需的矩阵乘积，包括块稀疏对角线质量矩阵

构建完成后，系统矩阵使用 Cholesky 分解进行分解，并使用前向和反向替换求解。我们可以使用 Warp 的 tile 基元来利用 M 的稀疏性，将此问题的批处理版本表述如下：

import warp as wp


@wp.kernel
def foward_dynamics(
    J_arr: wp.array3d(dtype=float),
    M_arr: wp.array3d(dtype=float),
    R_arr: wp.array3d(dtype=float),
    H_arr: wp.array3d(dtype=float),
    L_arr: wp.array3d(dtype=float),
):
    batch = wp.tid()

    J = wp.tile_load(J_arr[batch], 0, 0, 
  m=wp.static(6 * num_joints), n=num_dofs)
    P = wp.tile_zeros(m=wp.static(6 * num_joints), n=num_dofs, dtype=float)

    # compute P = M*J where M is a 6x6 block diagonal mass matrix
    for i in range(int(num_joints)):


        # 6x6 block matrices are on the diagonal
        M_body = wp.tile_load(M_arr[batch], i, i, m=6, n=6)

        # load a 6xN row from the Jacobian
        J_body = wp.tile_view(J, i * 6, 0, m=6, n=num_dofs)

        # compute weighted row
        P_body = wp.tile_matmul(M_body, J_body)

        # assign to the P slice
        wp.tile_assign(P, i * 6, 0, P_body)

    # compute H = J^T*P
    H = wp.tile_matmul(wp.tile_transpose(J), P)

    # cholesky L L^T = (H + diag(R))
    R = wp.tile_load(R_arr[batch], 0, 0, m=num_dofs, n=1, storage="shared")


    H += wp.tile_diag(R)
    L = wp.tile_cholesky(H)

    wp.tile_store(L_arr[batch], 0, 0, L)

# launch kernel with 64 threads per-robot
wp.launch_tiled(forward_dynamics, 
                dim=(num_robots,), 
                inputs=[J_arr, M_arr, R_arr, H_arr, L_arr], 
                block_dim=64)

在本例中，forward_dynamics() 内核执行 CRBA 方法，加载 Jacobian 矩阵和质量矩阵的图块，并计算其积以形成系统矩阵 H 及其 Cholesky 分解。在此特定用例中，Torch 需要启动十几个内核，而 Warp 实现则需要一个完全融合的内核。这可减少全局内存往返次数和启动开销，从而显著提高性能。

对于在 NVIDIA A100 80GB GPU 上运行的前向动态内核，1,024 个四足机器人的性能如下，所有计时均以毫秒为单位 (越低越好)：

A bar chart showing the performance of quadruped robots using Warp's tile primitives.
图 7、 Warp (SIMT) 使用基于其现有 SIMT 模型的实现，为使用 Warp 的平铺基元的分批机器人前向动力学提供性能。 Warp (Tile + cuBLASDx) 使用新的 tile 操作。 Torch (cuBLAS) 使用 Torch 的 bmm() 和 cholesky() 函数。请注意，Torch 实现未利用 M 的稀疏性。

这里 提供了完整的示例，用于展示 Cholesky 分解和反向替换的扩展程序即将推出。

未来发展 

Warp 和 MathDx 的未来版本将包括：

  • 对逐行归约运算符的额外支持
  • 从 Lambda 函数创建图块
  • 数据类型和布局转换
  • 提高 GEMM 运算的性能
  • 其他线性代数基元，包括各种矩阵分解算法。

了解详情 

Warp 1.5.0 中基于图块的编程提供了一种强大而灵活的 GPU 编程方法，使开发者能够显著提升其应用的性能。通过利用 cuBLASDx 和 cuFFTDx，Warp 1.5.0 可实现 GEMM 和 FFT 操作的无缝融合，从而减少内存 I/O 和内核启动用度。

要开始使用 Warp 的 Tile 操作，请在 Python 环境中使用以下命令安装 Warp：

pip install warp-lang

要运行融合的 MLP 示例，请使用以下命令：

python -m warp.examples.tile.example_mlp.py

如需详细了解 Warp 1.5.0 和 NVIDIA Math device acceleration (Dx) 库，请访问以下链接：

致谢 

感谢 Paweł Grabowski、Doris Pan、Neil Lindquist、Jakub Szuppe、Łukasz Ligowski、Sergey Maydanov 和 Łukasz Wawrzyniak 对此帖子和项目的贡献。

 

相关资源

标签

机器人 | 模型/库/框架 | 模拟/建模/设计 | CFD/CAE | CUBLAS | CUDA | cuFFT | General | HPC/科学计算 | Intermediate Technical | Models / Libraries / Frameworks | 仿真/建模/设计 | 制造业 | 功能 | 机器人学

关于作者

迈尔斯·麦克林（ Miles Macklin ）是NVIDIA 的首席工程师，致力于模拟技术。他从哥本哈根大学获得计算机科学博士学位，从事计算机图形学、基于物理学的动画和机器人学的研究。他在 ACM SIGGRAPH 期刊上发表了几篇论文，他的研究已经被整合到许多商业产品中，包括NVIDIA 的 PhysX 和 ISAAC 健身房模拟器。他最近的工作旨在为 GPU 上的可微编程开发健壮高效的框架。
Leopold Cambier 于 2021 在斯坦福大学获得了计算和数学工程博士学位。在博士期间，他专注于从理论和并行计算的角度研究大型稀疏线性系统的快速求解器。 2016 年和 2017 年，他还在NVIDIA cuDNN 团队实习。自从他在 2021 年 1 月加入全职以来，利奥波德一直在cuFFT。
Eric Shi 是 NVIDIA 的高级软件工程师，负责 Warp 库的开发。加入 NVIDIA 之前，他是劳伦斯利福摩尔国家实验室的计算物理学家。Eric 拥有普林斯顿大学等离子体物理学博士学位。他目前的兴趣是改进 Warp 对分布式计算的支持。