文档
注册

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/300/500 推理产品

Atlas 训练系列产品

Atlas推理系列产品(Ascend 310P处理器)AI Core

Atlas推理系列产品(Ascend 310P处理器)Vector Core

Atlas A2训练系列产品

Atlas 200/500 A2推理产品

注意事项

Atlas 200/300/500 推理产品,Expr中不能含有float16或float32类型的scalar。

Atlas 训练系列产品,Expr中不能含有float16类型的scalar。

Atlas推理系列产品(Ascend 310P处理器)AI Core,Expr中不能含有float16类型的scalar。

Atlas推理系列产品(Ascend 310P处理器)Vector Core,Expr中不能含有float16类型的scalar。

Atlas A2训练系列产品,Expr中不能含有float16类型的scalar。

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]
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词