Closed
[Requirement|需求建议]: stateless_random_normal_v2/v3、stateless_uniform 算子新增 ascend960dt/ascend960pr 平台支持 #3657
吴成文Add on  22 days agoClosed at  22 days ago
吴成文
吴成文
22 days ago Create

一、背景信息 (必填)

stateless_random_normal_v2、stateless_random_normal_v3、stateless_uniform 三个随机数算子当前仅支持 ascend950/ascend350,需新增 ascend960dt/ascend960pr 平台支持,使三个算子可在新一代芯片上使用。

二、价值/作用 (必填)

补齐随机数算子在新一代 Ascend 芯片(Ascend 960)上的能力,保证随机数生成类算子在 960 平台可直接通过 aclnn 直调/图方式使用,与主线其他已完成 960 适配的算子保持一致。

三、设计方案 (必填)

逐文件比对确认三个算子的 kernel 逻辑与 arch35(ascend950)版本完全等价,tiling 仅切分对齐粒度存在架构差异,属"少量修改",故按主线既有约定走架构隔离、复用 arch35 目录,不新建 arch92 目录。

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)

ACLNN 直调(CMakeLists.txt 中三个算子均为 ACLNNTYPE aclnn_exclude);同时注册 kernel 二进制编译配置。

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

沿用 ascend950 既有支持(见各算子 op_host/*_def.cpp 的输出 dtype 定义),本次不改变数据类型集合。

3.2.2 host侧设计
  • 算子注册:三个算子 def.cpp 增加 AddConfig("ascend960dt")/AddConfig("ascend960pr")(保留原有 ascend950/ascend350);CMakeLists.txt 的 SUPPORT_COMPUTE_UNIT 增加两颗芯片,SUPPORT_TILING_DIR 一一对应 arch35。
  • tiling 架构隔离(按 NpuArch 等值判断 DAV_9201/DAV_9202):
    • stateless_random_normal_v2:新增 IsArch92();arch92 核间按 32 元素对齐、UB 按 256 元素向下对齐,arch35 保持原 512 字节核间对齐、32 字节 UB 向上对齐;CompileInfo 增加 isArch92 标记,覆盖 platformInfo 不可用而回退 CompileInfo 通路的判定。
    • stateless_random_normal_v3:新增 BeforeProcess() 覆写,arch92 下将 coreAlignSize 由 512 改为 32;platformInfo 为空时保持默认值。
    • stateless_uniform:tiling 与 arch35 完全一致,无需隔离。
  • op_api:aclnn_normal_out.cpp 抽取 IsNormalV3SupportedArch(),Normal V3 兼容通路由仅 DAV_3510 扩展为 DAV_3510/DAV_9201/DAV_9202。
  • 编译配置:scripts/kernel/binary_config/ascendc_config.json 为三个算子注册 ascend960dt/ascend960pr 的 compute_units 与 compile_options(-mllvm -cce-aicore-dcci-before-kernel-end=false)。
3.2.3 kernel侧设计

复用 arch35 kernel;stateless_random_normal_v3 kernel 入口用 #if defined(__DAV_VEC__) 包裹(该算子为 AIV_ONLY 且依赖仅 Vector 核提供的 RegBase VF 接口)。

3.3 支持硬件

ascend960dt、ascend960pr(同时保留 ascend950、ascend350)。

3.4 算子约束限制

沿用既有约束(不新增/不放松):stateless_random_normal_v2 的 alg 仅支持 RNG_ALG_PHILOX;stateless_random_normal_v3 的 mean/stdev 不支持广播(shapeSize 需等于输出或为 1)。

💡 备注(选填)
芯片判断一律使用 NpuArch 等值匹配(SocVersion 枚举中 960 追加在末尾,不能用区间比较)。
新增 stateless_uniform 的 host shim(tests/ut/op_kernel/stateless_uniform_tiling.h)以修复其 op_kernel UT 因缺 __mul_i32toi64 而无法链接的问题。

关联 PR:https://gitcode.com/cann/ops-math/pull/6175

likedislike
吴成文吴成文
22 days ago Added label: requirement
吴成文
吴成文
22 days ago Comments:
likedislike
CANN-robotCANN-robotCollaborators
22 days ago Assign wuchengwen_oh as responsible person
CANN-robotCANN-robotCollaborators
22 days ago Close issue
CANN-robotCANN-robotCollaborators
22 days ago Added label: resolved