对于Vector计算,一般是用Unified Buffer去存放数据,再进行计算,所以整体数据流应该是从Global Memory>Unified Buffer>Global Memory。TIK提供了data_move接口实现Global Memory和Unified Buffer间的数据搬运,函数原型为:
data_move(dst, src, sid, nburst, burst, src_stride, dst_stride, *args, **argv)
在data_move的函数原型中,用户需要着重关注dst、src、nburst、burst、src_stride、dst_stride等6个参数。其中:
dst为目的操作数,也就是数据搬运的目的地址;src为源操作数,也就是数据搬运的起始地址;nburst表示需要执行的搬运次数;burst表示一次搬运的数据片段长度(单位为32Bytes);src_stride、dst_stride则分别代表源数据与目的数据的数据片段间隔(即前 burst 尾与后 burst 头的间隔)。
data_move支持连续地址与间隔地址两种搬运模式。
连续地址搬运是Tik算子开发中,最常见的数据搬运方式。
from tbe import tik # 实例化tik_instance对象 tik_instance = tik.Tik() # 定义一个在gm域的Tensor data_input_gm = tik_instance.Tensor("int32", (256,), name="data_input_gm", scope=tik.scope_gm) # 定义一个在ub域的Tensor data_input_ub = tik_instance.Tensor("int32", (256,), name="data_input_ub", scope=tik.scope_ubuf) # 使用data_move操作将输入的Tensor从gm搬到ub tik_instance.data_move(data_input_ub, data_input_gm, 0, 1, 32, 0, 0) # 对于ub进行一系列指令操作 ............. # 后续的搬出操作
在上述的案例中,我们首先分别在gm和ub空间开辟了一个256长度的,数据类型为int32的Tensor,然后我们执行从gm搬到ub的操作,数据搬运指令中几个参数的实际含义解释如下:
tik_instance.data_move(data_input_ub, data_input_gm, 0, 1, 32, 0, 0)
如下的数据搬运图所示,每个方块表示一个32Byte的Block,其中存放了8个int32,所以相当于是256个int32点对点进行搬运。
若输入数据较大,超过了UB的大小限制,此时需要将GM中的数据分多次搬入到UB进行计算,同样分多次搬出到GM。假设UB的可用存储空间为248KB,代码示例如下所示:
from tbe import tik tik_instance = tik.Tik() src_gm = tik_instance.Tensor("float16", (126976, 2), name="src_gm", scope=tik.scope_gm) dst_gm = tik_instance.Tensor("float16", (126976, 2), name="dst_gm", scope=tik.scope_gm) dst_ub = tik_instance.Tensor("float16", (126976, ), name="dst_ub", scope=tik.scope_ubuf) with tik_instance.for_range(0, 2) as i: # gm数据超过ub最大内存,先搬运一部分到ub进行计算,完成之后再搬回gm,可多次重复 tik_instance.data_move(dst_ub, src_gm[i*126976], 0, 1, 7936, 0, 0) with tik_instance.for_range(0, 3) as j: # repeat_times最大为255 一次无法全部计算完,可以分多次计算,此处为节省空间src和dst 为同一个ub tik_instance.vec_add(128, dst_ub[j*128*255], dst_ub[j*128*255], dst_ub[j*128*255], 255, 8, 8, 8) tik_instance.vec_add(128, dst_ub[3 * 128 * 255], dst_ub[3 * 128 * 255], dst_ub[3 * 128 * 255], 227, 8, 8, 8) # 将计算好的数据搬回至gm,再计算另外剩余的数据 tik_instance.data_move(dst_gm[i*126976], dst_ub, 0, 1, 7936, 0, 0) tik_instance.BuildCCE(kernel_name="data_move", inputs=[src_gm], outputs=[dst_gm])
with tik_instance.for_range(0, 2) as i
tik_instance.data_move(dst_ub, src_gm[i*126976], 0, 1, 7936, 0, 0)
with tik_instance.for_range(0, 3) as j: # 前三次计算,每次计算都处理256*255Byte的数据,即128*255个float16的数据。此处为节省空间src和dst 为同一个ub tik_instance.vec_add(128, dst_ub[j*128*255], dst_ub[j*128*255], dst_ub[j*128*255], 255, 8, 8, 8) # 最后一次计算,处理剩余数据,剩余数据需要重复227迭代进行计算 tik_instance.vec_add(128, dst_ub[3 * 128 * 255], dst_ub[3 * 128 * 255], dst_ub[3 * 128 * 255], 227, 8, 8, 8)
上述表示了连续的数据搬运,其中src_stride和dst_stride都是0,非连续的数据搬运的场景则稍显复杂。
from tbe import tik tik_instance = tik.Tik() data_input_gm = tik_instance.Tensor("int32", (256,), name="data_input_gm", scope=tik.scope_gm) data_input_ub = tik_instance.Tensor("int32", (176,), name="data_input_ub", scope=tik.scope_ubuf) # 非连续搬运的案例 tik_instance.data_move(data_input_ub, data_input_gm, 0, 4, 4, 4, 2) .............
如下的数据搬运图可以直观地表示这样的搬运过程:
间隔地址搬运场景比较少见,经常使用的还是连续地址搬运的场景,平常写算子的时候,也不会在burst上直接写数字,而是用:element_size_to_move * DATA_TYPE_SIZE / BLOCK_SIZE_BYTE 去表示:
tik_instance.data_move(data_input_ub, data_input_gm, SID, DEFAULT_NBURST, element_size_to_move * DATA_TYPE_SIZE // BLOCK_SIZE_BYTE, STRIDE_ZERO, STRIDE_ZERO)
其中DEFAULT_NBURST=1,BLOCK_SIZE_BYTE=32,STRIDE_ZERO=0。这样其他用户在阅读TIK算子代码的时候,其中参数的含义会比单纯的数字更加简明易懂。
在实际的开发过程中也会出现从Tensor某一位置搬运的场景,此时可以有如下的data_move使用方式:
from tbe import tik tik_instance = tik.Tik() data_input_gm = tik_instance.Tensor("int32", (256,), name="data_input_gm", scope=tik.scope_gm) data_input_ub = tik_instance.Tensor("int32", (256,), name="data_input_ub", scope=tik.scope_ubuf) # 有offset的连续数据搬运 tik_instance.data_move(data_input_ub[8], data_input_gm[16], 0, 1, 30, 0, 0) .............
由数据搬运图可见,总共需要搬运30个Block,但是src是从第3个Block开始搬运,即第16个int32开始搬运,然后将30个Block依次搬到dst第2个开始的Block中,即第8个int32开始填入。
如图2所示,输入数据为23个类型是float16的数据。
from tbe import tik tik_instance = tik.Tik() src_gm = tik_instance.Tensor("float16", (23, ), name="src_gm", scope=tik.scope_gm) dst_gm = tik_instance.Tensor("float16", (23, ), name="dst_gm", scope=tik.scope_gm) src_ub = tik_instance.Tensor("float16", (32, ), name="src_ub", scope=tik.scope_ubuf) dst_ub = tik_instance.Tensor("float16", (32, ), name="dst_ub", scope=tik.scope_ubuf) tik_instance.vec_dup(32, src_ub, 0, 1, 1) tik_instance.vec_dup(32, dst_ub, 0, 1, 1) with tik_instance.for_range(0, 2) as i: # 可以进行两次搬运,第一次32Byte对齐搬运,第二次数据前移再按照32Byte对齐方式搬运到ub tik_instance.data_move(src_ub[i*16], src_gm[i*(23-16)], 0, 1, 1, 0, 0) tik_instance.vec_add(32, dst_ub, src_ub, src_ub, 1, 1, 1, 1) # ub -> gm 采取相同搬运方式,第一次搬运32Byte的数据到gm;第二次将gm进行地址回退,满足32Byte对齐后,存储剩余ub中的数据 with tik_instance.for_range(0, 2) as i: tik_instance.data_move(dst_gm[i*(23-16)], dst_ub[i*16], 0, 1, 1, 0, 0) tik_instance.BuildCCE(kernel_name="data_move", inputs=[src_gm], outputs=[dst_gm])
请完成一个简单的TIK算子以实现基本的数据搬入搬出功能,要求如下:
注:假设ub的起始地址需要32B对齐。
【参考答案】:
from tbe import tik tik_instance = tik.Tik() data_input_gm = tik_instance.Tensor("float16", (146,), name="data_input_gm", scope=tik.scope_gm) data_input_ub = tik_instance.Tensor("float16", (272,), name="data_input_ub", scope=tik.scope_ubuf) tik_instance.data_move(data_input_ub, data_input_gm[2], 0, 9, 1, 0, 1) ............. data_output_gm = tik_instance.Tensor("int32", (128,), name="data_output_gm", scope=tik.scope_gm) data_output_ub = tik_instance.Tensor("int32", (384,), name="data_output_ub", scope=tik.scope_ubuf) tik_instance.data_move(data_output_gm, data_output_ub[32], 0, 8, 2, 4, 0)
【解析】:
gm搬运到ub的场景:因为129个fp16,每个Block能搬运16个fp16,所以需要搬运9个Block,因为gm从索引为2的数据开始连续搬运,且没有起始地址对齐的限制,所以gm的大小为9 * 16 + 2 = 146;因为ub是搬16个空16个,最后一个空的16个可以不分配空间,所以ub大小为9 * (16+16) - 16 = 272。每次搬1个Block,burst=1,且总共搬9个Block所以需要搬运9次,nburst=9。src端连续,src_stride=0,dst端前尾后头空16个fp16也就是1个Block,dst_stride=1,所以data_move中的参数分别是9, 1, 0, 1。
ub搬运到gm的场景:因为127个int32,每个Block能搬运8个int32,所以需要搬运16个Block,gm大小直接给成16 * 8 = 128即可;因为ub需要搬16个空32个,然后又从32地址开始搬,所以ub大小为8 * (16+32) - 32 + 32= 384。每次搬16个int32,所以每次搬2个Block,burst=2,且总共搬16个Block所以需要搬8次,nburst=8。src端前尾后头空32个int32也就是4个Block,src_stride=4,dst端连续,dst_stride=0,所以data_move中的参数分别是8, 2, 4, 0。