Closed
[Requirement|需求建议]: 【社区任务】新增 Bincount 算子AscendC实现贡献 #2069
DDPlysAdd on  Jun 26Closed at  Jul 30
DDPlys
DDPlys
Jun 26 Create

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

一、背景信息 (必填)

ops-math 仓当前缺少与 PyTorch torch.bincount 功能一致的 Ascend C 算子。Bincount 属于直方图 / 散射累加(scatter-add)类算子:统计一维非负整数张量中每个值出现的次数(无 weights 时计数;有 weights 时按权重累加),输出长度 L = max(max(self)+1, minlength),第 k 个元素即值 k 的统计量。

CANN 现有内置 aclnnBincount 的主计算 kernel 跑在 AI_CPU,单次调用总耗时约 663~1446us,性能较低。本需求在 Ascend C 上以 AI Core 单 kernel 重新实现,解决"内置算子走 AI_CPU、性能低"的问题。功能语义对标 torch.bincount,性能对标原 aclnnBincount(任务书要求整体达成其 10X)。

二、价值/作用 (必填)

  1. 补齐生态:为 ops-math 提供与 torch.bincount 一致的 Ascend C Bincount,覆盖计数与加权两种模式及多种 dtype 组合。
  2. 应用场景:标签 / 类别频次统计、加权直方图、混淆矩阵计数、数据分布分析等,是机器学习与数据处理的基础算子。
  3. 性能提升:主体计算从 AI_CPU 迁移到 AI Core 并单 kernel 化。实测:本算子 AI Core 单 kernel ~6~22us,相对官方 AI_CPU 基线加速比中位 ~110×、最高 ~434×,绝大多数用例远超任务 10X 目标。

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架)
  • Aclnn 直调:对外提供 aclnnBincountGetWorkspaceSize / aclnnBincount 两段式接口,参数 self / weights / minlength / out(torch.bincount 风格),可由用户直接调用,并提供 examples/test_aclnn_bincount.cpp 调用样例。
  • 采用与官方 math/bincount 一致的两层接口结构:底层 def 注册 GE 风格原型 array / size / weights / bins(进入 AICore 二进制);L2 把 host 侧 minlength 构造为 INT64 [1] 的 size 张量、任意 dtype 的 weights 经 l0op::Cast 统一为 float 后下发底层原型。
3.2 总体设计

整体为单遍散射:在散射累加的同一遍循环里顺带按核求局部 min 做负数检查。AscendC 实现流程如下:

flowchart TD
    A([aclnn 入口: self / weights / minlength / out]) --> B["L2: minlength 构造 size 张量;<br/>weights 统一 Cast 成 float;<br/>下发底层 array/size/weights/bins"]
    B --> C{"Host Tiling 分核判定"}
    C -->|"N ≤ 4096 或 L > 2048"| P1["单核快路径 coreNum=1"]
    C -->|"大输入 且 小输出"| P2["多核并行散射"]
    C -->|"L 超 UB"| P3["largeL: GM 直接散射回退"]

    P1 --> D1["整块 CopyIn self/weights → UB;<br/>私有直方图散射 桶 value 累加 w;<br/>同遍跟踪 localMin"]
    P2 --> D2["各核 UB 私有直方图散射 + localMin →<br/>min/直方图写 workspace →<br/>单道 SyncAll → 0 号核归并"]
    P3 --> D3["求 min → 清零 GM 输出 →<br/>逐元素读改写累加"]

    D1 --> N{"全局 min ≥ 0 ?"}
    D2 --> N
    D3 --> N
    N -->|"否, 含负数"| ERR([" ascendc_assert → AI Core 异常 507015<br/>报错中止, 不产出结果"])
    N -->|"是"| OUT["写出 out: 非 double 整段 DataCopy;<br/>double 经 IEEE-754 位拼接"]
    OUT --> Z([" out, 长度 L = max(vmax+1, minlength) "])
3.2.1 算子支持的数据类型
参数 数据类型 格式 维度
array INT8、INT16、INT32、INT64、UINT8 ND 1
weights FLOAT、FLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8、BOOL(可空,空为计数) ND 1
size INT64 - -
bins INT32、INT64、FLOAT、DOUBLE ND 1
3.2.2 host 侧设计
  • InferShape:依据 self 值依赖 + minlength 推导 out 长度(self 标记 ValueDepend)。
  • Tiling 分核:先判定是否走单核——N≤4096或 L>2048时 coreNum=1;否则满核均分,按 N 裁剪空转核。
  • UB 分块:私有直方图常驻 UB,tileDataNum 由 UB 容量反解;L 超 UB 时置 largeL=1 走 GM 直接散射回退路径。
3.2.3 kernel 侧设计
  • 单核快路径:coreNum=1 时一个核在 UB 私有直方图算完,不写 workspace、不 SyncAll、不归并,直接写出。占绝大多数用例、是达标主因。
  • 多核路径:各核片上私有直方图并行散射,min 与直方图一并写 workspace,单道 SyncAll 后 0 号核 DataCopy 读回逐桶归并写出。
  • 整块 CopyIn:self/weights 经 MTE 批量搬入 UB,消除逐元素 GM 标量读。
  • DOUBLE 输出:910B AI Core 无原生 fp64,kernel 以 float 累加后经 IEEE-754 位拼接产出 double。
  • 负数检查:kernel 求全局 min,min<0 触发 ascendc_assert → AI Core 异常(507015)报错中止。
3.3 支持硬件

Atlas A2 训练系列产品 / Ascend 910B。

3.4 算子约束限制

💡 备注(选填)

likedislike
DDPlysDDPlys
Jun 26 Modify description of issue
DDPlysDDPlys
Jun 26 Change title to "[Requirement|需求建议]: 【社区任务】新增 Bincount 算子AscendC实现贡献", originally titled "[Requirement|需求建议]: 新增 Bincount 算子AscendC实现贡献".
DDPlys
DDPlys
Jun 26 Comments:

/assign

likedislike
CANN-robotCANN-robotCollaborators
Jun 26 Assign ddplys as responsible person
DDPlysDDPlys
Jun 26 Modify description of issue
DDPlysDDPlys
Jun 26 Modify description of issue
DDPlysDDPlys
Jun 26 Modify description of issue
DDPlysDDPlys
Jun 27 Modify description of issue
DDPlysDDPlys
Jun 27 Associated pull request
DDPlys
DDPlys
Jun 29 Comments:

/assign

likedislike
CANN-robot
CANN-robotCollaborators
Jun 29 Comments:

Notice

This issue is already assigned to ddplys. Please do not assign repeatedly.

likedislike
DDPlysDDPlys
Jun 30 Associated pull request
llimwangCollaborators
Jul 30 Comments:

该 Issue 已超过 20 天未更新,先行关闭以便归档;如仍需跟进请评论或重新打开。

likedislike
LllimwangCollaborators
Jul 30 Close issue