【TBE算子开发】TIk算子开发总结
TIk算子开发总结
tik算子开发时可以采用当前我们DSL的开发框架进行，当变量过多时可以以class形式来开发
tik算子开发只需要在算子接口处进行shape、dtype等的验证，但tik不需要定义tvm.placeholder，直接定义TIK的DSL容器，语句如下
1 2 | import tik tik_instacnce = tik.Tik(tik.Dprofile("v100", "cloud")) |
目前只支持v100下cloud和mini，其中tik.Tik用于创建TIK的DSL的容器，tik.Dprofile用于指定编程时针对的编程环境
一旦创建了TIK容器，所有的语句都是基于tik_instacnce进行语句开发，不能出现普通的python语句
当前tik语句支持的数据类型有：int8,uint8, int16, uint16, int32, uint32, float16, float32，uint1（bool），比如在cloud形式下，除了输入输出的gm类型局限在fp32、fp16、int8、int32、uint8，tik中的数据类型都可以采用上述的数据类型；
定义Scalar方法：
1 | x = tik_instance.Scalar("int64") |
相当于定义了一个uint64的Scalar变量x；
定义tensor方法：
1 | x = tik_instance.Tensor(name = "X", dtype="float32", shape=[1,], scope=tik.scope_gm) |
其中scope表示tensor定义的区域；
Scalar和tensor的生命周期：在定义语句时创建，到定义语句所在代码块（if，else，for）的结束并释放，只有在存在期间可以被访问,，所以要根据需要去确定定义的范围
分支结构：采用if_scope进行分支判断(else部分为可选)，实现如下：
1 2 3 4 | with tik_instance.if_scope(x>1): #some methodwith tik_instance.else_scope: #some method |
For语句：
1 | tik_instance.for_range(begint,endt,name="i", thread_num=1,block_num=1,dtype="int32") as i |
等同于 for i in range(begint, endt)，里面的begint, endt可以使用常量，不一定要tik定义的scalar
thread_num可以选择1或者2，默认为1，当选择2时，为开启pingpong，在该for循环语句下，每个buf会double
block_num：分核数，默认值为1，只能在一个for循环中设置该值，建议在最外循环外设置该值；
tik.Tik.BuildCCE(kernel_name, inputs,outpus)：用于将DSL容器中语句编译成CCEC代码
inputs:List型变量，List中存放程序的输入Tensor（需要是global memory的存储类型）
outputs，List型变量，List中存放程序的输出Tensor（需要是global memory的存储类型）
以上为tik开发最基本的语句
下面给出一个简单的例子：使用tik语句实现a+b：
a：fp32 [1024, 1024，1024]
b : fp32 [1024, 1024，1024]
add属于vector操作，需要把a和b的数据从gm搬到ub中，在然后进行add操作，最后再从搬出来
考虑到a+b的结果可以复写到a中，因此考虑在ub中定义两个tensor，分别用于存放a和b；另外由于a和b的shape比较大，在考虑分核的同时，还需要考虑DB使能，因此a和b可以分到的UB为UB总大小的1/4（记为a_ub和b_ub）.
UB大小为256K，a_ub=b_ub=64k，相当于64*1024/4=16*1024个fp32数据，因此每次可以从gm往ub中搬运16*1024个数据
vector一拍可以处理64个fp32数据，因此repeat=1024*16/64=256(repeat 上限是255，考虑分两次执行，每次128)
具体实现如下：
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | # 定义tik容器，这里定义为cloud环境下tik_instance = tik.Tik(tik.Dprofile("v100", "cloud"))# 定义输入输出的gminput_x_gm = tik_instance.Tensor("float32", [1024, 1024, 1024], name="input_x", scope=tik.scope_gm)input_y_gm = tik_instance.Tensor("float32", [1024, 1024，1024], name="input_y", scope=tik.scope_gm)output_gm = tik_instance.Tensor("float32", [1024, 1024，1024], name="ioutput", scope=tik.scope_gm)# 最外面循环进行分核，cloud下最大为32核with tik_instance.for_range(0,32, block_num=32) as out: # db循环，循环次数为1024^3/32/16/1024=2048，该循环内定义的ub都会自动double with tik_instance.for_range(0, 2048, thread_num=2) as db_out: # 定义ub_a和ub_b a_ub = tik_instance.Tensor("float32", (16,1024), name="a_ub", scope=tik.scope_ubuf) b_ub = tik_instance.Tensor("float32", (16,1024), name="b_ub", scope=tik.scope_ubuf) # 搬运数据 tik_instance.data_move(a_ub, input_x_gm[out*32 + db_out*16%1024, db_out*16%1024, 0], 0, 1, 16*1024*4//32, 0, 0) tik_instance.data_move(y_ub, input_y_gm[out*32 + db_out*16%1024, db_out*16%1024, 0], 0, 1, 16*1024*4//32, 0, 0) # 数据运算 tik_instance.vadd(64, a_ub, a_ub, b_ub, 128, 1, 1, 1, 8, 8, 8) tik_instance.vadd(64, a_ub[8, 0], a_ub[8, 0], b_ub[8, 0], 128, 1, 1, 1, 8, 8, 8) # 搬运数据 tik_instance.data_move(input_x_gm[out*32 + db_out*16%1024, db_out*16%1024, 0], a_ub, 0, 1, 16*1024*4//32, 0, 0)tik_instance.BuildCCE(kernel_name="add", inputs=[input_x_gm, input_y_gm], outputs=[output_gm,]) |
- 点赞
- 收藏
- 关注作者
评论（0）