并行编程实战—CUDA Tile编程入门之二创建和应用

一、CUDA Tile的创建

在知道了基础的Tile的处理机制及相关的核函数的定义声明后，就来到了CUDA Tile的真正的目的即处理数据块。也就是说，如何设计块的结构、大小以及其形状、元素类型等。在Tile编程中，上述的参数表现为一个固定长度的、多维的标题元素数组，并且其形状和元素类型必须在编译时确定。
这里需要对Shape——形状进行说明一下，在CUDA Tile中主要面对的数据结构就是数组（多维张量），Tile可以理解为这个数组的一个子集（想象从一个大的数据块中砍出来一大块）。形状其实就是对砍下来的这一块进行维度上的描述。这样说可能不好理解，其实就是如何从原始的数组（数据块）中以何种方式来划分出一个子集（子数据块）进行加载。并且这里有一个限制，Tile的每个维度必须是2幂。
举一个简单例子，去买豆腐，肯定要告诉人家，买几块（大小），是长条形切还是方块形切（一个维度），是从当中下刀还是从两边（又一个维度）…，而且要求必须是够两顿饭的（2的幂）。
再以Tile为例，在ct.load（）中，shape参数是一个由常量整形组成的元组，如果shape=(tr,tc),则说明加载时会从原数组中切分出一个Tile,其大小为tr行和tc列。其实就是处理数据的粒度。
Tile是值语义的，很好理解，它需要复制过程。不过由于CUDA控制了Tile的硬件内容表示方式，所以其复制的开销很低，而且这也不需要开发者主动去管理相关的内存。
在CUDA的开发中，Tile有两种方式来创建。一种是通过数组加载数据来创建；另外一种是通过生成填充指定模式的Tile工厂函数来创建。

二、创建块的方法

CUDA Tile中创建Tile的具体方法有：

1. 数据加载
   看下面的定义：

ct.load(array,index,shape,*,order='C',padding_mode=PaddingMode.UNDETERMINED,latency=None,allow_tma=None)array：要从中加载数据的源数组（如CuPy或PyTorch张量） index：一个元组，指定在Tile空间中的起始索引（而不是原始元素的绝对索引） shape：一个由编译时常量整数组成的元组，定义要加载的Tile的尺寸

2. 工厂函数
   看相关的定义：

填充0或1的 Tile ct.zeros(shape,dtype)和ct.ones(shape,dtype)具有任意常量值的 Tile ct.full(shape,fill_value,dtype)包含[0,1,...,size-1]的1维 Tile ct.arange(size,dtype=...)shape：Tile 的维度大小（必须是2的幂的编译时常量）。 fill_value：（仅ct.full）用于填充Tile的标量值。 dtypeTile元素的数据类型。

对于Tile来说，其工厂需要的shape和dtype参数都是编译期确定的，而Python也提供了字面量以及Constant注解的内核参数进行支持。

三、初步控制

在创建Tile后就可以进行相关的管理和控制了，主要包括：

1. 查询块
   这个类似于前面的SIMT中通过blockIdx和threadIdx来计算线程索引。不过对Tile来说，只需要块索引就可以了，块内线程索引由编译器自行处理。
   在Python中，可以使用ct.bid(axis)返回当前块沿指定轴（0、1或2）的索引，作为int32标量；或者使用ct.num_blocks(axis)返回沿该轴的总块数（其对边界检查和循环计数很有用）
2. 编译时常量
   在刚刚的分析中已经说明，Tile的形状等必须在编译期确定，那么就可以使用字面量或编译期常量来进行处理。如使用字面量来指定Tile 形状和数据类型：

ct.zeros((64,64),dtype=ct.float32)#和ct::tile<float,ct::shape<64,64>>。

同样也可以使用Python Constant[T]来指定参数，如TILE: ct.Constant[int]
3. 控制流
Tile是数据块就需要控制访问，Tile内核的每个块建立了单一的控制流路径，和普通程序一样，利用条件和边界限定来进行控制流的处理。内部的Tile操作则由编译自动处理。不过它并不完全等同于其它语言的流控制，至少，Tile不允许从循环内部返回（其它细节参看官方文档）。
数据访问最重要的就是遍历。所以Tile支持常见的循环，主要的方式有range()、for、while 和嵌套循环，但在实际的应用中需要注意：步长不允许为负

四、例程

下面看一个相关例程：

importcuda.tileasctimportcupyascp@ct.kerneldefcombine_kernel(a:cp.ndarray,b:cp.ndarray,tile_size:ct.Constant[int]):pid=ct.bid(0)a_tile=ct.load(a,index=(pid,),shape=(tile_size,))zero_tile=ct.zeros((tile_size,),dtype=a.dtype)full_tile=ct.full((tile_size,),5,dtype=a.dtype)result_tile=a_tile+full_tile-zero_tile# a_tile + 5ct.store(b,index=(pid,),tile=result_tile)defmain():N=1024tile_size=256assertN%tile_size==0a=cp.arange(N,dtype=cp.float32)b=cp.zeros(N,dtype=cp.float32)num_blocks=N//tile_size# CUDA Tile的启动参数依次为：stream、grid、kernel、kernel_args# Tile大小由内核参数和ct.load的shap决定，不需要单独传block维度ct.launch(cp.cuda.get_current_stream(),(num_blocks,),combine_kernel,(a,b,tile_size),)expected=a+5ifcp.allclose(b,expected):print("验证通过！")else:print("验证失败。")print("输入前10:",a[:10])print("输出前10:",b[:10])print("期望前10:",expected[:10])if__name__=="__main__":main()

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

五、总结

利用Python来学习Tile应该是一个很好的入门方式。一来语言本身相对简单，二来对Tile的封装比C++可能更高级一些。这对于开发，是一个好事。入门之后，再理解C++的开发机制，估计会更容易理解。