Closed
[Requirement|需求建议]:OnesLike 算子 AICore 路径支持 int64 dtype #4546
RuiWang_Add on  5 days agoClosed at  4 days ago
RuiWang_Collaborators
5 days ago Create

一、背景信息 (必填)

math/ones_like 算子当前 AICore 路径仅支持 FLOAT/FLOAT16/BFLOAT16/INT8/INT32/UINT8/BOOL
共 7 种 dtype。int64 输入时 l0op::OnesLike 的 IsAiCoreSupport 检查不通过,自动回退到
AICPU 路径执行,性能明显低于 AIV 向量核,无法充分发挥 ascend950/350/960dt 硬件算力。

在大模型训练场景中,int64 张量常用于索引、计数(如 bincount、argmax 类下游消费)、
损失掩码占位等,torch.ones_like(x) 对 int64 张量的调用频率不低,当前均为 AICPU 执行。

二、价值/作用 (必填)

  1. int64 ones_like 由 AICPU 提升到 AICore(AIV)执行,消除回退带来的性能损失;
  2. 补齐与 PyTorch torch.ones_like 的 dtype 对齐缺口(PyTorch 支持 int64);
  3. 与 master 上已合入的 uint32/uint64 dtype views 特性形成互补:uint64 经 dtype view
    转 int64 后同样受益于 AICore 路径;
  4. 同仓库 conversion/zeros_like 已用同款 DAG 框架支持 int64,本需求无技术风险,
    且可验证 Vec::Duplicate<int64_t>/Vec::CopyOut<int64_t> 在目标硬件上可用。

三、设计方案 (必填)

3.1 使能方式

Aclnn 直调(aclnnInplaceOne),无需框架侧改动;PyTorch 侧经 dtype view 消费。

3.2 总体设计
3.2.1 算子支持的数据类型

新增 INT64:x/y 的 DataType 列表各追加 ge::DT_INT64(ND 格式,与既有 dtype 保持一致),
共支持 FLOAT、FLOAT16、BFLOAT16、INT8、INT32、UINT8、BOOL、INT64。

3.2.2 host侧设计
  • ones_like_def.cpp:原型注册 dtype 列表追加 INT64;
  • ones_like_tiling_key.h:新增 ONES_LIKE_TPL_INT64(tiling key 编码值 8,key = 1+(8<<1) = 17)
    及对应 DECL/SEL 声明;
  • ones_like_tiling_arch35.cpp:SetTilingData/CalcInputDtype/RunTiling 三处增加 int64
    分支,tiling 计算复用 ElewiseBaseTiling.DoTiling<OnesLikeDAG<int64_t>::OpDag>;
  • op_api/ones_like.cpp:AICore dtype 支持列表加入 DT_INT64;
  • 三个芯片目录 ones_like_binary.json 新增 OnesLike_INT64 二进制条目。
3.2.3 kernel侧设计

零改动。kernel 采用模板化 DAG(MAKE_CONST(T,1) + Vec::Duplicate<T> + Vec::CopyOut<T>),
编译 int64 二进制时 DTYPE_X=int64_t 自动实例化 OnesLikeDAG<int64_t>,与
conversion/zeros_like 的 int64 实现方式完全一致。

3.3 支持硬件

ascend950(Ascend 950PR/950DT)、ascend350、ascend960dt(AICore 路径);
老 910 系列保持 AICPU 回退不变。

3.4 算子约束限制

无新增约束。输入输出 dtype 需一致、shape 需一致的既有校验规则不变;
老 910 芯片 int64 仍走 AICPU(保守处理,后续另行评估)。

likedislike
RRuiWang_Collaborators
5 days ago Added label: requirement
RRuiWang_Collaborators
5 days ago Associated pull request
陈思陈思Collaborators
5 days ago Assign wangrui_ as responsible person
CANN-robotCANN-robotCollaborators
4 days ago Close issue
CANN-robotCANN-robotCollaborators
4 days ago Added label: resolved