if_scope

功能说明

创建一个TIK的条件判断语句，当条件满足时，则执行if_scope代码块。

伪代码如下：

# 当满足cond1时，执行if_scope代码块。
if_scope(cond1)

函数原型

if_scope(cond)

参数说明

表1 参数说明
参数名称	输入/输出	含义
cond	输入	代表判断条件。支持InputScalar，Scalar，立即数(int，float)， Expr，bool，以及any、all、negate函数。其中： Scalar、InputScalar类型的参数会自动转换为表达式，例如若参数名称为a，则转换为：a != 0。如果Expr为数值表达式，则会自动判断是否!=0，例如参数为scalar // 2,则转换为 scalar // 2 !=0 须知： Expr中支持如下运算符：比较运算符：>，<，!=，==，>=，<= 算术运算符：+，-，*，/，//，% 位运算符：&，\|，<<，>> 可以用支持的运算符组成复杂表达式，但是不支持比较运算符的连用，例如，若想表达n>1，且n<4：不支持如下表达方式 1 < n < 4 可以使用如下表达方式 tik.all((1 < n) , (n < 4))

支持的芯片型号

Atlas 200/500 A2推理产品

注意事项

Atlas 200/500 A2推理产品，Expr中不能含有float16类型的scalar。

返回值

TikWithScope对象。

继承TVM的WithScope。

调用示例

tik_instance = tik.Tik()

cond = tik_instance.InputScalar(dtype="int16", name="cond")
src_gm = tik_instance.Tensor(dtype="int64", shape=(16, ),
                             scope=tik.scope_gm, name="src_gm")
dst_ub = tik_instance.Tensor(dtype="int64", shape=(4, ),
                             scope=tik.scope_ubuf, name="dst_ub")
dst_gm = tik_instance.Tensor(dtype="int64", shape=(4, ),
                             scope=tik.scope_gm, name="dst_gm")
# 将dst初始化为 src_gm[0:3]
tik_instance.data_move(dst_ub, src_gm[0:3,], 0, 1, 1, 0, 0)

# 根据cond的值，将src的不同切片搬运到dst中
# 当满足cond==1时，执行if_scope内代码块
with tik_instance.if_scope(cond == 1):
    tik_instance.data_move(dst_ub, src_gm[4:7,], 0, 1, 1, 0, 0)

# 将dst搬运到gm，用以输出
tik_instance.data_move(dst_gm, dst_ub, 0, 1, 1, 0, 0)

tik_instance.BuildCCE(kernel_name="if_scope", inputs=[src_gm, cond], outputs=[dst_gm])

结果示例

输入数据(src_gm):
[0, 1, 2,..., 15]

输入数据(cond):
1

输出数据(dst_gm):
[4, 5, 6, 7]

父主题： 程序控制