并行编程实战—CUDA中的Tile编程之三Tile空间加载与存储

一、数据处理在CUDA tile的编程中最重要的是从IT转到DT数据为王。这几乎已经成为了共识。虽然现在已经不再提什么DT了但并不是它不重要了而是成为了一个基础建设。从所谓的大数据分析、挖掘到现在的AI海量数据的学习。都离不开数据的处理人们已经相当依赖数据处理的结果来进行决策。所以能够如何更好的高效、便捷的处理数据就成为了一个重要的问题。cuda tile就是用来作为更好的处理数据的一个技术出现了。cuda tile的编程模型中对数据的处理抽象出来就是两大类Array数组和Tile瓦片即数据块。数组是多维的、全局可见的容器它对所有Tile内核块是可见的。而Tile可以认为是数组上的切片go语言中的Slice它当然也是一种容器它通常是数组的一个子集所以也是多维的容器。不过它是局部的仅对单个的cuda tile代码块可见。二、加载和存储CUDA Tile的编程模型中Tile的加载和存储通常有两种形式Tile空间加载和存储这种方式使用Tile空间的索引来进行加载和存储通过视力对象规定数组元素到Tile的映射模式。如果硬件支持的话它可以由编译器降级到张量内存加速器TMA操作这比gather要快得多使用Gather和Scatter使用索引Tile或指针Tile来指示加载或存储时数组元素的源或目标的加载和存储需要说明的是开发者必须自行控制越界元素加载的时的默认值。而越界写入则会默认被自动抛弃三、Tile空间加载和存储Tile空间的加载时会创建一个视图对象指定数组如何划分为Tile大小的区域的网格。这种视图对象的映射称为Tile空间Tile内核可以通过Tile空间索引一次加载或存储一个区域。下面看一下官网提供的图理解了这个图基本就理解了Tile的划分说明图中的的数组的形状是10*16即一个10行16列的数组空间Tile的形状定义为24即切出一个2行4列的块10/25,16/44,恰好整除。即形成了一个Tile Grid54。Tile空间加载的核心就是Tile化视图(分区视图)通过上面的图片可以看到数组元素是如何映射到指定大小的Tile的。图中显示了一个分区视图是一个非重叠、无间隙的Tile视图。这也意味着如果数组的维度不能恰好的划分Tile时在多个不同的维度可能会跨越数组边界的Tile会被部分填充这是由开发者自行确定的。分区视图的加载和存储往往需要使用结构化的Tile空间它是用来在全局内存和Tile间进行数据移动的首选方式。内核会先创建一个定义Tile空间的视图对象然后通过Tile的空间索引一次加载或存储一个Tile。Python中会使用Array.tiled_view(tile_shape)来得到一个TiledView它映射的是将数组划分成指定形状的Tile。可以通过这个视图对象的.load(index) / .store(index, tile)方法来接受Tile的空间索引。下面看一个官网简单的Tile的例子ct.kerneldefvec_add(a,b,c,TILE:ct.Constant[int]):a_viewa.tiled_view((TILE,))b_viewb.tiled_view((TILE,))c_viewc.tiled_view((TILE,))bidct.bid(0)a_tilea_view.load((bid,))b_tileb_view.load((bid,))c_view.store((bid,),a_tileb_tile)另外在Python中还提供单次调用加载和存储的机制其接口如下ct.load(array,index,shape)# 在给定的Tile空间索引处读取指定形状的Tilect.store(array,index,tile)# 对应的写入操作#示例ct.kerneldefvec_add(a,b,c,TILE:ct.Constant[int]):bidct.bid(0)# 此块沿轴 0 的 Tile 空间索引a_tilect.load(a,index(bid,),shape(TILE,))# (index, shape) 取 a 的第 bid 个 TILE 大小的区域b_tilect.load(b,index(bid,),shape(TILE,))ct.store(c,index(bid,),tilea_tileb_tile)# 将 Tile 写回 c 的第 bid 个区域它与视图加载的与存储的不同在于Tile形状存储的位置不一样。前者是Tile形状绑定到视图对象而后者则是Tile形状每次调用时提供。当同一分区多次加载存储使用时推荐使用视图机制而单次加载或存储时考虑使用ct.load/ct.store当然前面提到的Tile空间边界的处理Python中也给出了相关的接口。ct.load中有一个padding_mode参数用来控制越界元素的值。两种常用模式是PaddingMode.ZERO用零填充越界元素PaddingMode.UNDETERMINED默认越界元素值由实现决定即开发者确定Tile完全在边界内时对于存储ct.store会自动丢弃对越界位置的写入不需要padding_mode参数。它同样适用于tiled_view即在视图创建时确定其padding_modect.kerneldefedge_safe(arr_in,arr_out,TILE:ct.Constant[int]):bidct.bid(0)tilect.load(arr_in,index(bid,),shape(TILE,),padding_modect.PaddingMode.ZERO)# 部分边缘Tile的越界通道变为 0ct.store(arr_out,index(bid,),tiletile)# 越界写入静默丢弃四、Gather和Scatter数据的处理往往是各种情况都有当数据空间齐整、分配完美的情况下一般是使用Tile的空间加载或存储机制。但有时候访问的数据是不规则的或有相互依赖的。就需要使用gather和scatter方式了。它允许从数组的非连续或非均匀的元素加载或存储到Tile。Python提供了接口ct.gather() / ct.scatter()用来整数索引Tile内置边界检查。前者的边界检查默认是开启的默认返回0后者越界会自动丢弃。下面看官网的一个例子ct.kerneldefvec_add_gather(a,b,c,TILE:ct.Constant[int]):bidct.bid(0)indicesbid*TILEct.arange(TILE,dtypect.int32)# 每个通道一个元素索引a_tilect.gather(a,indices)# 每个通道加载 a[indices[i]]b_tilect.gather(b,indices)ct.scatter(c,indices,a_tileb_tile)# 每个索引存储一个值到 c这种机制中默认检查是安全的。当确认索引不越界时可以使用check_boundsFalse禁用边界的检查。五、例程下面根据上面的示例代码实现一个完整的例程#!/usr/bin/env python3from__future__importannotationsimportargparseimportcupyascpimportcuda.tileasct DEFAULT_TILE256DEFAULT_SIZE1024EDGE_SIZE1003ct.kerneldefvec_add_tiled(a,b,c,tile:ct.Constant[int]):blockct.bid(0)a_viewa.tiled_view((tile,))b_viewb.tiled_view((tile,))c_viewc.tiled_view((tile,))c_view.store((block,),a_view.load((block,))b_view.load((block,)))ct.kerneldefvec_add_load_store(a,b,c,tile:ct.Constant[int]):blockct.bid(0)a_tilect.load(a,index(block,),shape(tile,))b_tilect.load(b,index(block,),shape(tile,))ct.store(c,index(block,),tilea_tileb_tile)ct.kerneldefedge_safe_copy(arr_in,arr_out,tile:ct.Constant[int]):blockct.bid(0)tile_valuect.load(arr_in,index(block,),shape(tile,),padding_modect.PaddingMode.ZERO,)ct.store(arr_out,index(block,),tiletile_value)ct.kerneldefvec_add_gather_scatter(a,b,c,tile:ct.Constant[int]):blockct.bid(0)indicesblock*tilect.arange(tile,dtypect.int32)a_tilect.gather(a,indices)b_tilect.gather(b,indices)ct.scatter(c,indices,a_tileb_tile)defpositive_int(value:str)-int:numberint(value)ifnumber0:raiseargparse.ArgumentTypeError(value must be greater than zero)returnnumberdefparse_args()-argparse.Namespace:parserargparse.ArgumentParser(description__doc__)parser.add_argument(--size,typepositive_int,defaultDEFAULT_SIZE)parser.add_argument(--tile,typepositive_int,defaultDEFAULT_TILE)parser.add_argument(--device,typeint,default0)returnparser.parse_args()deflaunch_and_check(label:str,kernel,args:tuple,grid:tuple[int,...],actual:cp.ndarray,expected:cp.ndarray,stream,)-None:ct.launch(stream,grid,kernel,args)stream.synchronize()max_errorfloat(cp.max(cp.abs(actual-expected)).get())ifmax_error!0:raiseRuntimeError(f{label}failed: max error is{max_error:g})print(f[OK]{label}: grid{grid}, max error{max_error:g})defmain()-int:argsparse_args()ifargs.device0:raiseSystemExit(--device must be non-negative)ifargs.tile(args.tile-1):raiseSystemExit(--tile must be a power of two)ifargs.size%args.tile!0:raiseSystemExit(f--size ({args.size}) must be divisible by --tile ({args.tile}) for the unpadded examples)cp.cuda.Device(args.device).use()propertiescp.cuda.runtime.getDeviceProperties(args.device)nameproperties[name]ifisinstance(name,bytes):namename.decode()print(fGPU:{name}(device{args.device}))print(fCUDA Tile:{ct.__version__}; CuPy:{cp.__version__})streamcp.cuda.get_current_stream()blocksargs.size//args.tile rngcp.random.default_rng(0)arng.random(args.size,dtypecp.float32)brng.random(args.size,dtypecp.float32)expectedabprint(Running examples from demo.txt...)ccp.empty_like(a)launch_and_check(vec_add_tiled (TiledView),vec_add_tiled,(a,b,c,args.tile),(blocks,),c,expected,stream,)c.fill(0)launch_and_check(vec_add_load_store,vec_add_load_store,(a,b,c,args.tile),(blocks,),c,expected,stream,)edge_inputcp.arange(EDGE_SIZE,dtypecp.float32)edge_outputcp.empty_like(edge_input)launch_and_check(fedge_safe_copy ({EDGE_SIZE}elements),edge_safe_copy,(edge_input,edge_output,args.tile),((EDGE_SIZEargs.tile-1)//args.tile,),edge_output,edge_input,stream,)gather_arng.random(EDGE_SIZE,dtypecp.float32)gather_brng.random(EDGE_SIZE,dtypecp.float32)gather_outputcp.empty_like(gather_a)launch_and_check(fvec_add_gather_scatter ({EDGE_SIZE}elements),vec_add_gather_scatter,(gather_a,gather_b,gather_output,args.tile),((EDGE_SIZEargs.tile-1)//args.tile,),gather_output,gather_agather_b,stream,)print(All CUDA Tile examples OKed.)return0if__name____main__:raiseSystemExit(main())运行结果python ~/cuda_tile_project/cuda_tile_demo.py GPU: NVIDIA RTX PRO 4000 Blackwell (device 0) CUDA Tile: 9.9.99; CuPy: 14.1.1 Running examples from demo.txt... [OK] vec_add_tiled (TiledView): grid(4,), max error0 [OK] vec_add_load_store: grid(4,), max error0 [OK] edge_safe_copy (1003 elements): grid(4,), max error0 [OK] vec_add_gather_scatter (1003 elements): grid(4,), max error0 All CUDA Tile examples OKed.注意代码中对解析参数时Tile为2的幂以及整除等的判断方法这些都属于一些小技巧其它其实没有什么可分析的和前面Tile相关的内容一致。六、总结CUDA Tile的加载与存储也是数据操作的基础应用。不管数据如何处理都需要处理好数据的加载和处理后的存储。既要数据来源的准确又要保证数据结果的安全。这才是数据使用中的重点。