math/ones_like 算子当前 AICore 路径仅支持 FLOAT/FLOAT16/BFLOAT16/INT8/INT32/UINT8/BOOL 共 7 种 dtype。int64 输入时 l0op::OnesLike 的 IsAiCoreSupport 检查不通过,自动回退到 AICPU 路径执行,性能明显低于 AIV 向量核,无法充分发挥 ascend950/350/960dt 硬件算力。
math/ones_like
l0op::OnesLike
IsAiCoreSupport
在大模型训练场景中,int64 张量常用于索引、计数(如 bincount、argmax 类下游消费)、 损失掩码占位等,torch.ones_like(x) 对 int64 张量的调用频率不低,当前均为 AICPU 执行。
torch.ones_like(x)
torch.ones_like
conversion/zeros_like
Vec::Duplicate<int64_t>
Vec::CopyOut<int64_t>
Aclnn 直调(aclnnInplaceOne),无需框架侧改动;PyTorch 侧经 dtype view 消费。
aclnnInplaceOne
新增 INT64:x/y 的 DataType 列表各追加 ge::DT_INT64(ND 格式,与既有 dtype 保持一致), 共支持 FLOAT、FLOAT16、BFLOAT16、INT8、INT32、UINT8、BOOL、INT64。
ge::DT_INT64
ones_like_def.cpp
ones_like_tiling_key.h
ONES_LIKE_TPL_INT64
1+(8<<1)
ones_like_tiling_arch35.cpp
SetTilingData
CalcInputDtype
RunTiling
ElewiseBaseTiling.DoTiling<OnesLikeDAG<int64_t>::OpDag>
op_api/ones_like.cpp
ones_like_binary.json
OnesLike_INT64
零改动。kernel 采用模板化 DAG(MAKE_CONST(T,1) + Vec::Duplicate<T> + Vec::CopyOut<T>), 编译 int64 二进制时 DTYPE_X=int64_t 自动实例化 OnesLikeDAG<int64_t>,与 conversion/zeros_like 的 int64 实现方式完全一致。
MAKE_CONST(T,1)
Vec::Duplicate<T>
Vec::CopyOut<T>
DTYPE_X=int64_t
OnesLikeDAG<int64_t>
ascend950(Ascend 950PR/950DT)、ascend350、ascend960dt(AICore 路径); 老 910 系列保持 AICPU 回退不变。
无新增约束。输入输出 dtype 需一致、shape 需一致的既有校验规则不变; 老 910 芯片 int64 仍走 AICPU(保守处理,后续另行评估)。
一、背景信息 (必填)
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 执行。二、价值/作用 (必填)
torch.ones_like的 dtype 对齐缺口(PyTorch 支持 int64);转 int64 后同样受益于 AICore 路径;
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(保守处理,后续另行评估)。