TIK语句

什么是TIK语句

上节我们讲过表达式的概念。TIK表达式不会生成对应的IR,而TIK语句会生成对应的IR,当TIK语句中包含了表达式时,表达式会被实例化进行相应的计算。

本节我们会介绍TIK算子编码过程中,使用TIK语句与Python语句结合的灵活性,从而实现算子不同的需求,以及在编译时确定和在运行时确定的语句特性。

语句(statement)是一个会产生影响的代码单元,例如新建一个变量或显示某个值。

TIK语句主要包含:赋值语句、循环语句、条件语句、计算语句、打印语句、容器管理语句等。

TIK语句以普通Python语言要素的形式存在,但是“能够直接或者间接对TIK DSL程序产生影响”是TIK API区别于其它部分Python程序的本质特性。

读者在写前端TIK时,除了用到TIK API外,还会用到普通Python语言要素,即Python的类与集合。也就是TIK前端包含了TIK语句和Python语句,在某些情况下可以加快开发者的编码效率,同时减少出现bug的可能性。

循环语句的妙用

我们已经知道用for_range完成TIK的循环,如果要在TIK语句中使用循环变量i,则必须要用for_range语句,同时,它也可以开启double buffer来完成性能的优化。

案例一:

如果需要编写如下的TIK算子片段,Tensor需要定义4次,vec_dup命令也要写4次。

input_ub0 = tik_instance.Tensor("int32", (64,), name="input_ub0", scope=tik.scope_ubuf)
input_ub1 = tik_instance.Tensor("int32", (64,), name="input_ub1", scope=tik.scope_ubuf)
input_ub2 = tik_instance.Tensor("int32", (64,), name="input_ub2", scope=tik.scope_ubuf)
input_ub3 = tik_instance.Tensor("int32", (64,), name="input_ub3", scope=tik.scope_ubuf)
tik_instance.vec_dup(64, input_ub0, 0, 1, 8)
tik_instance.vec_dup(64, input_ub1, 0, 1, 8)
tik_instance.vec_dup(64, input_ub2, 0, 1, 8)
tik_instance.vec_dup(64, input_ub3, 0, 1, 8)

上述定义和操作语句需要重复定义,需要修改里面的变量,极易出错,而且不满足对扩展开放,对修改关闭的原则。

如果使用Python的循环语句,就可以达到事半功倍的效果。

repeat_times = 4
input_ub_list = []
for i in range(repeat_times):
    tmp_ub = tik_instance.Tensor("int32", (64,), name="input_ub"+str(i), scope=tik.scope_ubuf)
    input_ub_list.append(tmp_ub)
    tik_instance.vec_dup(64, input_ub_list[i], 0, 1, 8)

我们先定义了一个input_ub_list列表专门用来存放Tensor变量,然后使用了Python的循环语句,让位于循环内的TIK语句循环执行多次,每执行一条TIK语句,都会生成一条对应的IR,这样就会有和重复编写代码一样的效果。

在TIK算子中,Python解释器会根据代码逻辑依次去执行语句,每执行一条TIK语句,都会生成对应的IR。

这种TIK前端写法非常方便拓展,只要修改例子中repeat_times的数值,就会有不同的效果。

TIK实际上是一个代码生成器(Code Generator),Python解释器执行到某条TIK语句的时候才会emit(发射)指令生成对应的IR,所以TIK语句跟Python语句有明显的区别。

条件语句的妙用

我们已经知道用if_scope和else_scope来完成算子内的条件分支判断。在TIK算子的编写过程中,Python的条件语句同样也会有相当大的用处。我们对案例一作为扩展。

案例二:

need_clear = False
repeat_times = 4
input_ub_list = []
for i in range(repeat_times):
    tmp_ub = tik_instance.Tensor("int32", (64,), name="input_ub"+str(i), scope=tik.scope_ubuf)
    input_ub_list.append(tmp_ub)
for i in range(repeat_times):
    if need_clear:
        tik_instance.vec_dup(64, input_ub_list[i], 0, 1, 8)
    else:
        tik_instance.vec_dup(64, input_ub_list[i], 1, 1, 8)

用户在使用算子的时候,会输入运行时参数和编译时参数,编译时参数是编译时已知的,所以可以有效地简化TIK算子编译出的CCE的代码量,而简化则需要用Python的条件语句来进行判断哪些TIK语句会被执行到,从而只生成相应的IR。

这里假设need_clear和repeat_times是编译时参数,也就是用户编译TIK算子的同时需要给出的,根据之前给出的规则,只有Python解释器执行到的TIK语句会生成相应的IR。

合理地使用Python的条件语句,根据算子的编译时参数合理地区分TIK分支,不同的分支走不同的逻辑,可以有效地减少TIK编译产生的冗余代码,从而提升算子的性能。

编译时参数和运行时参数的区别

确定编译时参数和运行时参数的阶段不同,当编写好TIK算子之后,需要把TIK算子编译成CCE算子,而编译时参数就是在这个过程中确定,而编译出的CCE算子需要进一步通过CCEC编译器编译成.o,然后输入运行时参数去跑相应的功能。需要注意的是,TIK语句所定义的操作不是在Python调用时发生,而是在程序定义完成进行编译的阶段发生,调用只是生成对应的IR。

编译时参数实际上是确定的参数,我们完全可以在TIK算子中,使用Python语句直接去处理它,而区分两者其实很简单,看看BuildCCE中的inputs和outputs,就是运行时参数。

案例三:编译时参数

VLENFP32 = 64
# model_point_num的值是一个确定的值
model_point_num = 150
# model_point_repeat和model_point_residual可以在TIK转IR时直接计算出来
model_point_repeat = model_point_num // VLENFP32
model_point_residual = model_point_num % VLENFP32
input_ub = tik_instance.Tensor("float32", (256,), name="input_ub", scope=tik.scope_ubuf)
output_ub = tik_instance.Tensor("float32", (256,), name="output_ub", scope=tik.scope_ubuf)
# 使用Python的条件语句
if model_point_repeat != 0:
    tik_instance.vec_add(VLENFP32, output_ub, input_ub, output_ub, model_point_repeat, 8, 8, 8)
# 使用Python的条件语句
if model_point_residual != 0:
    tik_instance.vec_add(model_point_residual, output_ub[model_point_repeat*VLENFP32:], input_ub[model_point_repeat*VLENFP32:], input_ub[model_point_repeat*VLENFP32:], 1, 8, 8, 8)

给定算子的入口函数中,model_point_num是编译时参数,是确定的,在TIK语句转成IR的时候,就用具体的数值替代了。VLENFP32是Vector单元256B下并行的最多64个float32的个数,是给定数据类型后固定的。

这里我们用Python的条件语句很好地控制了哪些TIK语句需要执行哪些不需要执行,从而生成对应的CCE算子。

例如:

在TIK的Python程序中非TIK API部分的Python代码往往与算子的常量化密切相关,在表现形式上往往是借助普通Python变量进行配置的计算。可以看出,这样可以很好地根据不同的入参情况,编译出不同的CCE算子。

案例四:运行时参数

运行时参数,最直观的就是案例中input_ub中存放的具体float32的数值,该值在算子没有跑起来之前是无法获知的,所以说如果算子的逻辑是需要运行时参数给定的,那这块逻辑在算子中就不能省略,如下例:

VLENFP32 = 64
# model_point_ub是运行时参数
model_point_ub = tik_instance.Tensor("float32", (64,), name="input_ub", scope=tik.scope_ubuf)
# model_point_num是从model_point_ub中取出来的,编译TIK算子时并不知道它的具体数值
model_point_num = tik_instance.Scalar("int32", "model_point_num", init_value=model_point_ub[0])
# model_point_repeat和model_point_residual同样也不知道具体什么数值
model_point_repeat = model_point_num // VLENFP32
model_point_residual = model_point_num % VLENFP32
input_ub = tik_instance.Tensor("float32", (256,), name="input_ub", scope=tik.scope_ubuf)
output_ub = tik_instance.Tensor("float32", (256,), name="output_ub", scope=tik.scope_ubuf)
# 不能用Python的条件判断语句
with tik_instance.if_scope(model_point_repeat != 0):
    tik_instance.vec_add(VLENFP32, output_ub, input_ub, output_ub, model_point_repeat, 8, 8, 8)
# 不能用Python的条件判断语句
with tik_instance.if_scope(model_point_residual != 0):
    tik_instance.vec_add(model_point_residual, output_ub[model_point_repeat*VLENFP32:],input_ub[model_point_repeat*VLENFP32:], input_ub[model_point_repeat*VLENFP32:], 1, 8, 8, 8)

因为model_point_num是从运行时参数model_point_ub中取出来的,编译时无法知道它的确切数值,所以并不能判断if语句会不会运行,所以这里不能用Python的条件语句,而是要用TIK的条件语句。

按案例四的写法,最后生成的CCE代码很长,因为需要在运行时接收model_point_num的值从而判断它需要走哪个分支,这样的运行时参数更加灵活,但同时,也相应地牺牲了性能。

有读者可能会问,那我在编译时参数中使用TIK语句是否可以?答案是肯定的,但是同样会生成冗余的CCE代码。如果编译时已经确定了,就不必要把它留到运行时,我们算子还是需要达到最优性能的。

总结

Python语句是按照顺序逻辑、循环逻辑或条件判断逻辑去有序执行,当Python解释器执行到TIK语句时,就会生成对应的IR。可以这么认为,Python语句虽然不能直接生成IR,但是它又是生成IR的控制器,是生成IR的重要一环。

读者在编码环节中如果遇到选择是TIK语句还是Python语句的问题时,一定要想一想这个语句是编译阶段发挥作用还是运行阶段发挥作用的,如果是前者,则用Python语句,如果是后者,就有TIK语句。

如果读者在编写TIK算子过程中能够熟练地结合Python语句和TIK语句,一方面能够减少可能生成的bug,提高编码效率,另一方面又可以在源代码端有效地优化部分代码,提升算子的性能。

课后练习

请说明在案例三中,model_point_num满足什么样的条件两个分支都会走到,什么条件只会走第一个分支,什么条件只会走第二个分支?

【参考答案】:

  1. model_point_num>64,且不能被64整除
  2. model_point_num能被64整除,且不等于0
  3. model_point_num∈[0,64)