if_scope
功能说明
创建一个TIK的条件判断语句,当条件满足时,则执行if_scope代码块。
伪代码如下:
# 当满足cond1时,执行if_scope代码块。 if_scope(cond1)
函数原型
if_scope(cond)
参数说明
参数名称 |
输入/输出 |
含义 |
---|---|---|
cond |
输入 |
代表判断条件。 支持InputScalar,Scalar,立即数(int,float), Expr,bool,以及any、all、negate函数。
其中:
|
支持的芯片型号
Atlas 200/300/500 推理产品
Atlas 训练系列产品
Atlas 推理系列产品AI Core
Atlas 推理系列产品Vector Core
注意事项
Atlas 200/300/500 推理产品,Expr中不能含有float16或float32类型的scalar。
Atlas 训练系列产品,Expr中不能含有float16类型的scalar。
Atlas 推理系列产品AI Core,Expr中不能含有float16类型的scalar。
Atlas 推理系列产品Vector Core,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]
父主题: 程序控制