Closed
[Requirement|需求建议]: Transpose 算子支持 ascend960dt/ascend960pr 平台(含 HIF4/FP4 4bit 数据类型) #3627
zhanw_codingAdd on  25 days agoClosed at  24 days ago
zhanw_coding
zhanw_codingCollaborators
25 days ago Create

一、背景信息 (必填)

Transpose(permute)算子当前仅注册 ascend310p / ascend910_93 / ascend910b / ascend950 / mc62 / ascend350
等芯片,且 x 的数据类型支持列表不包含 4bit(HIF4/FP4)类型。新平台 ascend960dt、ascend960pr 需要接入
Transpose 算子,并在其架构(arch92)下新增对 4bit 数据类型的支持。

二、价值/作用 (必填)

  • 使 Transpose 算子在 ascend960dt / ascend960pr 平台上可用(图模式 op 注册 + aclnn 直调 + 二进制编译配置齐备)。
  • 新增 4bit 数据类型(DT_FLOAT4_E2M1 / DT_FLOAT4_E1M2 及 HIFLOAT4 / HIFLOAT4_SCALE)支持,满足低比特
    量化场景下转置的数据搬运需求。
  • 针对 4bit 尾轴转置场景补充专用 tiling/kernel 模板,并对 NDDMA 类模板在新架构下启用双 buffer,改善大 shape 搬运性能。

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)
  • 图模式:operator 定义(op_host/transpose_def.cpp)新增芯片配置。
  • Aclnn 直调:aclnnPermute / TransposeV2 的 dtype 校验按架构分支扩展。
  • 二进制化:op_host/config/{ascend960dt,ascend960pr}/ 下新增 binary.json 与 simplified_key.ini,
    并在 scripts/kernel/binary_config/ascendc_config.json 中登记 compute_units 与编译选项。
3.2 总体设计
3.2.1 算子支持的数据类型
  • x / y(ascend960dt、ascend960pr,arch92)新增:DT_FLOAT4_E2M1、DT_FLOAT4_E1M2,以及自定义
    DT_HIFLOAT4 = 42、DT_HIFLOAT4_SCALE = 43;原有基础类型(int8/uint8/int16/uint16/int32/uint32/
    int64/uint64/float/float16/bf16/bool/hifloat8/float8_e5m2/float8_e4m3fn 等)保持一致。
  • perm 索引类型:DT_INT32 / DT_INT64。
  • 通过 CombineDataTypes 对 x 与 perm 的 dtype 做笛卡尔组合,生成 x/perm 的 DataType 列表,避免手工枚举。
  • Aclnn 侧按架构分支:在 NpuArch::DAV_9201 / DAV_9202 下使用 ARCH92_DTYPE_SUPPORT_LIST
    (aclnn_permute)/ ARCH92_AICORE_DTYPE_SUPPORT_LIST(transpose),其余架构沿用原列表。
3.2.2 host侧设计
  • transpose_def.cpp:新增 buildConfig() 统一构造 OpAICoreConfig;新增
    AddConfig("ascend960dt", config92)、AddConfig("ascend960pr", config92)。
  • CMakeLists.txt:COMPUTE_UNIT 增加 ascend960dt、ascend960pr,TILING_DIR 均指向 arch35。
  • scripts/kernel/binary_config/ascendc_config.json:Transpose 的 compute_units 增加 ascend960dt/
    ascend960pr,compile_options 增加 -mllvm -cce-aicore-dcci-before-kernel-end=false。
  • 新增 4bit 双切分 tiling:transpose_tiling_b4_dual_cut_arch35.{h,cpp},新增
    SplitMode::B4_DOUBLE_SPLIT = 10009;触发条件为 4bit 且尾轴转置。
  • transpose_tiling_arch35:新增 CheckB4Tiling()/TryB4Tiling()/AdjustShapeForB4()/CheckB4ShapeInfo();
    4bit 非尾轴转置按 int8 view 处理(尾轴折算);4bit 尾轴转置走 B4 模板;校验 4bit 输入/输出尾轴必须为偶数。
  • 缓冲策略:新增 tiling 字段 bigDimBufferNum;NDDMA 缓冲数在 DAV_9201/DAV_9202 下取双缓冲;
    BIG_DIM 场景支持超大函数拆分后的缓冲回退(双缓冲不适用时回退单缓冲)。
3.2.3 kernel侧设计
  • 新增 op_kernel/arch35/transpose_b4.h:TransposeHIF4 双切分模板。UB 内流程为
    CopyIn(Compact,HIF4 原始字节)→ Cast fp4→bf16 → 在 bf16 域转置 → Cast bf16→fp4 → CopyOut;
    UB buffer 采用双缓冲。
  • transpose_apt.cpp:新增 TilingKey B4_DUAL_CUT = 10009 分支及 TransposeHIF4DualCutProcess;
    4bit 场景下 TENSOR_MOVE/SMALL_SHAPE/CUT_ONE_AXIS/CUT_TWO_AXIS/BIG_DIM/N_LAST/GATHER 等模板按 int8_t 实例化处理。
  • NDDMA 相关模板(big_dim / cut_one_axis / cut_two_axis / tensor_move / n_last / transdata_5hd /
    transdata_5hd_021 / with_gather)在 __DAV_920R1__ 下使用带额外参数的 DataCopy/DataCopyPad 接口,
    base.h 新增 NDDMA_BUFFER_NUM(920R1 为 2,其余为 1),并将 UB buffer 按缓冲数切分。
3.3 支持硬件
  • 新增:ascend960dt、ascend960pr(arch92 / DAV_9201、DAV_9202)。
  • 保持:ascend310p、ascend910_93、ascend910b、ascend950、mc62、ascend350。

3.4 算子约束限制

  • 4bit 类型:x 与 y 的尾轴长度必须能被 2 整除(否则 tiling 报错)。
  • B4 双切分模板(TilingKey 10009)仅在 4bit 且尾轴转置场景命中;非尾轴 4bit 转置按 int8 view 处理。
  • 4bit 专用路径与 NDDMA 双缓冲仅针对 arch92(DAV_9201/DAV_9202)生效。
  • UB 内 move_align 限制内轴数(DUAL_CUT_UB_LEN = 4)。

💡 备注(选填)

likedislike
zhanw_codingzhanw_codingCollaborators
25 days ago Added label: requirement
zhanw_coding
zhanw_codingCollaborators
25 days ago Comments:

/assign

likedislike
CANN-robotCANN-robotCollaborators
25 days ago Assign zhanw_coding as responsible person
CANN-robotCANN-robotCollaborators
24 days ago Close issue
CANN-robotCANN-robotCollaborators
24 days ago Added label: resolved