已关闭
[Requirement|需求建议]: GatherElements算子AscendC实现贡献 #202
cc创建于  2025年12月3日关闭于  2025年12月25日
cc
cc
2025年12月3日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud(背景信息)

使用AscendC对GatherElements算子进行实现,实现了AscendC实现的GatherElements算子对Atlas 200/500 A2推理产品和Atlas 800I/T A2硬件的适配。

Origin(信息来源)

哈工大算子团队

Benefit / Necessity (价值/作用)

GatherElements 算子核心功能是按照给定的索引张量,在指定轴上“逐元素”抽取源张量对应位置的数据,输出张量与索引张量形状完全一致。该算子弥补了 Gather、IndexSelect 等“整段 / 整行”采样算子的不足,能够实现“点对点”的精细化取值,是 高频使用的索引算子。在深度学习模型中,GatherElements 广泛应用于
多标签分类中按目标位置抽取 logits、序列到序列模型中根据 beam-index 重组隐藏状态、图像分割后处理阶段按最大索引回取置信度和对比学习及稀疏特征路由中按动态下标提取嵌入向量等中。用户仅需给出索引张量与指定轴,即可一次性完成任意形状的高维采样,无需手写循环,显著提升开发效率。

Design(设计方案)

host侧设计:
为简化分块逻辑,首先将输入 X 和 Indices 的维度沿着采集维度 dim 展平为三个逻辑部分:PreDim (并行行数)、GatherDim(采集长度)、PostDim (Tiling 粒度)。Host 侧的分块目标是沿着 PreDim * GatherDim 定义的总行数 totalRows 进行核心分配,并沿着 PostDim 进行单核 Tiling 切分。
1)分核策略
分核策略基于需要处理的总任务行数 totalRows = idxPreDim * idxGatherDim 来确定实际使用的核心数 usedCores,遵循数据量优先加速的原则。策略首先进行加速判断:当 totalRows大于或等于平台总核心数 coreNum 时,为了实现最大化并行,将采用所有核心 进行满核加速;而当 totalRows 小于coreNum 时,则仅使用与任务行数相同的核心数,以避免闲置核心带来的调度开销。
2) 单core内切分策略
使用双缓冲,并充分使用UB空间的原则。由于没有借助临时空间,且输入输出的形状相同,将UB等分为2等分,每一等分对应一个输入或者输出。
以块为单位,充分利用每一等分的UB空间。
kernel侧设计:
进行Init和Process两个阶段,其中Process包括数据搬入(CopyIn)、计算(Compute)、搬出(CopyOut)三个阶段。

  1. 初始化阶段(Init):
    从 tiling 数据中获取输入输出张量的维度信息,包括 xPreDim、xGatherDim、xPostDim 以及对应的索引张量维度 idxPreDim、idxGatherDim、idxPostDim。
    获取当前核 ID(coreId_)和总核数(coreNum_),为后续任务分配提供依据。
    设置全局张量(GlobalTensor)的起始地址,包括输入张量 xGm_、索引张量 idxGm_(数据类型为 int32_t)和输出张量 yGm_。
    初始化管道(TPipe)并分配 UB 缓冲区,为索引输入队列 idxInQue_ 和输出队列 yOutQue_ 分配按 32 字节对齐的内存空间,确保符合硬件访问优化要求。
  2. 主处理阶段(Process) :
    任务分配:每个核按照跨步(stride)方式处理总行数(totalRows),通过循环变量 rowId = coreId_ + i × coreNum_ 实现负载均衡,确保各核工作量分布均匀。
    行列坐标计算:对于每个处理行,计算其在 PreDim 和 GatherDim 维度上的坐标(preIdx, gatherIdx),进而确定索引数据在全局内存中的基础偏移量。
    分块循环:在 PostDim 维度上以 tileSize 为单位进行循环切分,每次处理一个数据块,通过双缓冲实现数据搬运与计算的流水线并行。
  3. 数据搬入(CopyIn):
    使用 DataCopyPad 接口将索引数据从全局内存搬移至 UB 中的本地张量。通过配置 DataCopyExtParams 参数,设置数据拷贝的 burst 长度和 stride 模式,确保高效的内存访问。
    支持非对齐数据的搬运,通过 pad 参数处理边界情况,保证数据完整性。搬入完成后将本地张量压入索引输入队列,为计算阶段提供数据。
  4. 计算(Compute):
    从索引输入队列中获取索引张量,遍历每个索引值并进行边界处理:对负索引进行规范化(加上 xGatherDim_ 使其为正),确保索引值在有效范围内。
    根据索引值计算输入张量中的对应数据位置:xRealOffset = xBase + indexVal × xPostDim_ + (postStart + i)。
    使用 GetValue 接口从输入张量中读取对应位置的数据,并通过 SetValue 接口写入输出本地张量的相应位置,实现点对点的数据采集。
  5. 数据搬出(CopyOut):
    使用 DataCopyPad 接口将计算结果从 UB 搬出到全局内存的输出张量中。配置合适的 DataCopyExtParams 参数,确保数据高效回写。
likedislike
cc
cc
2025年12月3日 评论:

/assign

likedislike
CANN-robotCANN-robot成员
2025年12月3日 将 c15503545287 设为负责人
cccc
2025年12月3日 关联了pull request:GatherElements算子AscendC实现贡献
张喻翔
张喻翔
2025年12月3日 评论:

感谢您的反馈,可以关联上您的PR,在验证后由commiter审视合入。

likedislike
Ppeihaobo
2025年12月3日 关联了pull request:个人-AscendC实现HardSwish算子贡献
Ppeihaobo
2025年12月3日 删除了关联的pull request:个人-AscendC实现HardSwish算子贡献
cccc
2025年12月15日 修改了issue 的描述
cccc
2025年12月15日 修改了issue 的描述
Cchenqi317成员
2025年12月25日 issue状态由 进行中 改变为 已确认
CANN-robotCANN-robot成员
2025年12月25日 添加了label:Accepted
CANN-robotCANN-robot成员
2025年12月25日 关闭了 issue
CANN-robotCANN-robot成员
3月20日 添加了label:resolved