昇腾社区首页
中文
注册

aclnnGather

产品支持情况

产品 是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]
[object Object]Atlas 200I/500 A2 推理产品[object Object] ×
[object Object]Atlas 推理系列产品 [object Object]
[object Object]Atlas 训练系列产品[object Object]

功能说明

  • 算子功能:对输入tensor中指定的维度dim进行数据聚集。

  • 计算公式: 给定张量selfself,维度dd,和一个索引张量indexindex,定义nnselfself的维度,idi_d表示维度dd的索引,indexidindex_{i_d}表示索引张量indexindex在维度dd上的第idi_d个索引值。对指定维度d的gather功能可以用如下的数学公式表示:

    gather(X,index,d)i0,i1,,id1,id+1,,in1=selfi0,i1,,id1,indexid,id+1,,in1gather(X,index,d)_{i_0,i_1,\cdots,i_{d-1},i_{d+1},\cdots,i_{n-1}} = self_{i_0,i_1,\cdots,i_{d-1},index_{i_d},i_{d+1},\cdots,i_{n-1}}
  • 示例:

    • 示例1: 假设输入张量self=[123456789]self=\begin{bmatrix}1 & 2 & 3\\ 4 & 5 & 6\\ 7 & 8 & 9\end{bmatrix},索引张量index=[0210]index=\begin{bmatrix}0 & 2\\ 1 & 0\end{bmatrix}dim=0dim = 0,那么输出张量out=[1842]out=\begin{bmatrix}1 & 8\\ 4 & 2\end{bmatrix},具体计算过程如下:out0,0=selfindex0,0,0=self0,0=1out0,1=selfindex0,1,1=self2,1=8out1,0=selfindex1,0,0=self1,0=4out1,1=selfindex1,1,1=self0,1=2\begin{aligned} out_{0,0}&=self_{index_{0,0}, 0}=self_{0,0}=1 \\ out_{0,1}&=self_{index_{0,1}, 1}=self_{2,1}=8 \\ out_{1,0}&=self_{index_{1,0}, 0}=self_{1,0}=4 \\ out_{1,1}&=self_{index_{1,1}, 1}=self_{0,1}=2 \end{aligned}
    • 示例2: 假设输入张量self=[123456789]self=\begin{bmatrix}1 & 2 & 3\\ 4 & 5 & 6\\ 7 & 8 & 9\end{bmatrix},索引张量index=[0210]index=\begin{bmatrix}0 & 2\\ 1 & 0\end{bmatrix}dim=1dim = 1,那么输出张量out=[1354]out=\begin{bmatrix}1 & 3\\ 5 & 4\end{bmatrix},具体计算过程如下:out0,0=self0,index0,0=self0,0=1out0,1=self0,index0,1=self0,2=3out1,0=self1,index1,0=self1,1=5out1,1=self1,index1,1=self1,0=4\begin{aligned} out_{0,0}&=self_{0, index_{0,0}}=self_{0,0}=1 \\ out_{0,1}&=self_{0, index_{0,1}}=self_{0,2}=3 \\ out_{1,0}&=self_{1, index_{1,0}}=self_{1,1}=5 \\ out_{1,1}&=self_{1, index_{1,1}}=self_{1,0}=4 \end{aligned}

函数原型

每个算子分为undefined,必须先调用“aclnnGatherGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用“aclnnGather”接口执行计算。

  • aclnnStatus aclnnGatherGetWorkspaceSize(const aclTensor* self, const int64_t dim, const aclTensor* index, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)
  • aclnnStatus aclnnGather(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)

aclnnGatherGetWorkspaceSize

  • 参数说明:

    • self(aclTensor*,计算输入):公式中的self,Device侧的aclTensor,数据类型需要与out一致,shape支持0-8维,维度数需要与index一致,支持undefinedundefined支持ND。
      • [object Object]Atlas 推理系列产品[object Object]、[object Object]Atlas 训练系列产品[object Object]:数据类型支持DOUBLE、FLOAT16、FLOAT32、INT32、UINT32、INT64、UINT64、INT16、UINT16、INT8、UINT8、BOOL。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]:数据类型支持DOUBLE、FLOAT16、BFLOAT16、FLOAT32、INT32、UINT32、INT64、UINT64、INT16、UINT16、INT8、UINT8、BOOL。
    • dim(int64_t,计算输入):公式中的d,Host侧的整型,取值范围[-self.dim(), self.dim() - 1]。
    • index(aclTensor*,计算输入):公式中的index,Device侧的aclTensor,数据类型支持INT32、INT64,shape支持0-8维,维度数需要与self一致,且shape需要与out一致,除dim指定的维度外,其他维度的size需要小于等于self对应维度的size,支持undefinedundefined支持ND,index中的具体数值代表self对应dim的下标,取值范围[-self.shape[dim], self.shape[dim] - 1],index中的索引数据不支持越界。
    • out(aclTensor*,计算输出):Device侧的aclTensor,数据类型需要与self一致,shape支持0-8维,且shape需要与index一致,支持undefinedundefined支持ND。
      • [object Object]Atlas 推理系列产品[object Object]、[object Object]Atlas 训练系列产品[object Object]:数据类型支持DOUBLE、FLOAT16、FLOAT32、INT32、UINT32、INT64、UINT64、INT16、UINT16、INT8、UINT8、BOOL。
      • [object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]、[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]:数据类型支持DOUBLE、FLOAT16、BFLOAT16、FLOAT32、INT32、UINT32、INT64、UINT64、INT16、UINT16、INT8、UINT8、BOOL。
    • workspaceSize(uint64_t*,出参):返回需要在Device侧申请的workspace大小。
    • executor(aclOpExecutor**,出参):返回op执行器,包含了算子计算流程。
  • 返回值:

    aclnnStatus:返回状态码,具体参见undefined

    [object Object]

aclnnGather

  • 参数说明:

    • workspace(void*, 入参): 在Device侧申请的workspace内存地址。
    • workspaceSize(uint64_t, 入参): 在Device侧申请的workspace大小,由第一段接口aclnnGatherGetWorkspaceSize获取。
    • executor(aclOpExecutor*, 入参): op执行器,包含了算子计算流程。
    • stream(aclrtStream, 入参): 指定执行任务的Stream。
  • 返回值:

    aclnnStatus:返回状态码,具体参见undefined

约束说明

调用示例

示例代码如下,仅供参考,具体编译和执行过程请参考undefined

[object Object]