inplace_add

功能说明

inplace_add实现lhs指定行和rhs的加法计算。

例如：

res = lhs
res[ids,:] += rhs
return res

您可以在“te/lang/cce/te_compute/inplace_compute.py”查看接口定义。

约束说明

此接口暂不支持与其他TBE DSL计算接口混合使用。

Atlas 200/300/500 推理产品：支持的数据类型有float16、float32、int32

Atlas 训练系列产品：支持的数据类型有float16、float32、int32

• rhs第一维最大为7934，超过7934无法处理。
• 当rhs第一维较大时（如5000+），可能会因操作系统OS栈溢出导致core dump，可以使用ulimit -s指令把栈空间增大，例如由8192增加到81920。

函数原型

te.lang.cce.inplace_add(lhs, inplace_ids, rhs)

参数说明

• lhs：输入左tensor。
• inplace_ids：为int、list类型，数值大于等于0且不能大于lhs第一维的整型。其长度必须和输入rhs的第一维的长度相同。
• rhs：输入右tensor或scalar，除了第一维外，其他维度必须与lhs相同。当inplace_ids为int型单个数时，rhs比lhs差一个维度。例如：lhs为(10,1024)，inplace_ids为[5] ，rhs为(1,1024）；lhs为(10,1024) ，inplace_ids为5，rhs为(1024,）。

返回值

res_tensor：表示计算后的tensor。

支持的芯片型号

Atlas 200/300/500 推理产品

Atlas 训练系列产品

调用示例

import tvm
import te.lang.cce
input_dtype = "float16"
dataA = tvm.placeholder((6,1024), name="dataA", dtype=input_dtype)
dataB = tvm.placeholder((5,1024), name="dataB", dtype=input_dtype)
inplace_ids = [1,1,4,2,2]
res = te.lang.cce.inplace_add(dataA, inplace_ids, dataB)
res.shape = (6,1024)
# res[0] = dataA[0]
# res[1] = dataA[1] + dataB[0] + dataB[1]
# res[2] = dataA[2] + dataB[3] + dataB[4]
# res[3] = dataA[3]
# res[4] = dataA[4] + dataB[2]
# res[5] = dataA[5]