Warp 1.5.0 引入基于图块的编程以增强 GPU 效率

最新版本的 Warp 1.5.0 引入了基于图块的编程原语，有望提高 GPU 效率和生产力。据 NVIDIA 称，新工具利用 cuBLASDx 和 cuFFTDx，可在 Python 内核中实现高效的矩阵乘法和傅里叶变换。这一进步对于加速模拟和科学计算尤其重要。

GPU 编程的演变

在过去的十年里，GPU硬件已经从纯粹的SIMT（单指令、多线程）执行模型转变为严重依赖协作操作的执行模型，从而提高了效率。随着 Tensor Core 数学单元成为 GPU 计算不可或缺的一部分，对其进行高效编程至关重要。 BLAS 等传统高级 API 虽然提供了广泛的抽象，但在与用户程序交互时通常在集成和效率方面存在不足。

Warp 中基于图块的编程

基于图块的编程模型（例如 Warp 1.5.0 中引入的模型）允许开发人员在多个线程可以协作执行的图块上表达操作。该模型扩展了 Warp 基于内核的编程，以包含基于图块的操作，从而实现从 SIMT 到基于图块的执行的无缝过渡。它减少了对手动索引和共享内存管理的需求，同时支持训练的自动区分。

扭曲平铺基元

Warp 的新图块基元包括构造、加载/存储、线性代数和映射/归约操作。这些原语自然地扩展了 Warp 现有的基于内核的编程模型。可以使用 NumPy 式操作在 Warp 内核内部构建图块，从而实现跨 CUDA 块的数据的高效管理。

增强型矩阵乘法

基于图块的编程的主要优点之一是能够执行协作矩阵乘法。 Warp 1.5.0 引入了 wp.tile_matmul() 原语，它利用 cuBLASDx 调度适当的 Tensor Core MMA 指令以获得最佳性能。这一进步可以显着提高性能，对于较大的矩阵实现大约 70-80% 的 cuBLAS 性能。

案例研究和应用

Warp 中基于图块的编程对于需要密集线性代数的应用非常有益，例如机器人模拟和信号处理。例如，在机器人模拟中，Warp 的图块基元可以有效地计算前向动力学所需的矩阵乘积，通过减少全局内存往返和启动开销，优于 Torch 等传统框架。

未来发展

Warp 和 MathDx 的未来版本将包括对行式归约运算符、从 lambda 函数创建图块、改进的 GEMM 运算性能以及新的线性代数基元的额外支持。这些增强功能将继续优化 GPU 编程效率。

欲了解更多详情，请访问 NVIDIA 官方博客。

图片来源：Shutterstock