并行编程实战——CUDA Tile编程入门之一内核和函数

一、CUDA Tile编程

在前面的CUDA13.1的更新中有过说明，在新的版本中，CUDA Tile是一个技术亮点。以前是没有开发环境，所以也无法展开相关的代码开发。随着CUDA13.2的安装成功，配套的相关CUDA Tile安装也安装完成并通过了初步的测试。这样就可以直接在CUDA中开发相关的Tile编程了。
不过，官方的版本说明中，CUDA13.2只支持Python环境下的Tile开发，要想使用C++版本的，则需要升级到CUDA13.3。由于目前硬件和驱动的原因，尚无法升级，所以先从Python版本的CUDA Tile开始学习。

二、编程思想

每一次重要技术的升级，其实本质是一种编程思想的变化。CUDA Tile也是如此，它将SIMT模型中需要对每个线程进行处理相同任务（数据）转变为对以Block为单位处理Tile（瓦片，其实就是数据块）。即编译器通过定义好的对Tile的操作自动进行展开并行化并分配到相关的线程。
这样做的好处非常明显，极大的减轻了开发者的并行的压力（扩大了设计开发上的并行粒度，但并未降低运行时的粒度），能够更好的“压榨”硬件的极限。
这也再一次验证了前面的分析，技术总是由复杂到简单发展的。不过需要说明的是，这种简单和复杂是相对的，不是绝对的。
Tile的本质是提供了NVIDIA GPU更高级的操作接口，或者说更简单的操作方式。特别是对于复杂的内存加速和张量处理。需要提醒的是，Tile只是思想层面的转变，CUDA底层的运行逻辑并未发生本质的变化。可以这样理解，CUDA Tile的思维角度下的代码是SIMT思维角度下代码的进一步抽象。它仍然需要SIMT作为底层的技术支持。
改变的是思想及其指导下的代码，不改变的是底层的运行逻辑和工作顺序。

三、基本入门

在Python的开发环境中，Tile API位于cuda.tile模块中，一般来说，其别名都会定义为ct。Tile的内核是引导GPU操作的入口点（有点类似于C++中main函数），Tile中的函数可以在内部互相调用。它会在启动Grid中的每个Block（块）上执行一次（类似于Init）。
不过需要说明的是，Tile内核与SIMT内核一样，仍然无法从主机代码中直接调用，必须通过Launch方式调用执行。Python中，必须使用@ct.kernel装饰器将函数标记为Tile内核入口点，使用@ct.function装饰器将函数标记为可从Tile内核或其他Tile函数中调用。
不过在实际的应用中，为了方便，只要从内核中调用的函数都会自动编译为Tile的代码，这也意味着@ct.function装饰器是一个可选项或者说内核中调用的函数默认都有@ct.function装饰器。
对于函数参数的数组，可以使用DLPack或CUDA Array Interface的设备驻留数组，Pytorch中的张量或CuPy中的数组等等。对于普通的参数则直接传递即可。
刚刚说过，Tile内核在Tile块的网格上启动（SIMT内核在线程块的网络上启动），开发者可以从三个维度对Grid的参数进行指定。Python的ct.launch提供了四个参数用来处理相关的位置即CUDA流、指定每个维度中Tile块数量的网格元组、内核对象以及内核参数元组。
正常的开发中，一般需要启动能够覆盖整个数组的多个块，特别是在无法块数量无法整除的情况下，仍然要启动+1的数量的块（向上取整）。目的是保证对数组大小的完整容纳。

四、例程

下面看一个简单的Tile的例程：

import cupy as cp
import numpy as np
import cuda.tile as ct

# 定义Tile的大小
# 每个线程块将处理TILE_SIZE个元素
TILE_SIZE = 16

# 定义CUDA Tile核函数
# 使用 @ct.kernel装饰器标记这是一个GPU 内核
@ct.kernel
def vector_add_kernel(a, b, result):
    # 获取当前线程块的 ID (一维)
    block_id = ct.bid(0)

    # 从全局内存加载一个Tile 的数据到寄存器/共享内存
    # shape=(TILE_SIZE,) 指定了要加载的 Tile 大小
    a_tile = ct.load(a, index=(block_id,), shape=(TILE_SIZE,))
    b_tile = ct.load(b, index=(block_id,), shape=(TILE_SIZE,))

    # 执行元素级加法，得到一个新的Tile
    result_tile = a_tile + b_tile

    # 将结果Tile存储回全局内存
    ct.store(result, index=(block_id,), tile=result_tile)

# 主机端启动
def vector_add(a: cp.ndarray, b: cp.ndarray, result: cp.ndarray):
    # 确保输入输出形状一致
    assert a.shape == b.shape == result.shape

    # 计算网格 (Grid) 大小: 总共需要的线程块数量
    # ct.cdiv 是向上取整的除法，用于处理向量大小不能被 TILE_SIZE 整除的情况
    grid = (ct.cdiv(a.shape[0], TILE_SIZE), 1, 1)

    # 启动内核
    # 参数: CUDA流, 网格大小, 内核函数, 内核参数
    ct.launch(cp.cuda.get_current_stream(), grid, vector_add_kernel, (a, b, result))

if __name__ == "__main__":
    # GPU上创建测试数据
    VECTOR_SIZE = 128
    rng = cp.random.default_rng()
    a = rng.random(VECTOR_SIZE, dtype=cp.float32)
    b = rng.random(VECTOR_SIZE, dtype=cp.float32)
    result = cp.zeros_like(a)

    # 调用封装好的函数启动内核
    vector_add(a, b, result)

    # 复制GPU结果到CPU
    a_np = cp.asnumpy(a)
    b_np = cp.asnumpy(b)
    result_np = cp.asnumpy(result)
    expected = a_np + b_np

    # NumPy验证
    print("结果是否一致:", np.allclose(result_np, expected, atol=1e-5))

代码来自官方文档修改过来。

五、总结

对于新技术要敏感，但是否要在工程中引入新技术，还是需要根据实际情况来决定。善于接受新事物和新技术新思想，不代表就不顾实际的强行应用，至少对于CUDA Tile来说就是如此。它还太新，无论硬件和框架本身都可能存在各种问题，而实际应用的软件也很少，可借鉴的不多。这样，对于大多数的开发者说，可以先自行掌握技术，再在实际工程小小范围的试用，评估。再考虑如何引入。
饭要一口口的吃，事要一件件的办，急不得。