
pyasc 算子开发指南load_data_with_transpose 实现带转置的 2D 数据加载【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc在昇腾 AI 处理器上开发算子时本地存储A1/B1中的二维分形数据经常需要以转置布局送入 A2/B2 供向量计算使用。asc.language.basic.load_data_with_transpose是 pyasc 中对应 Ascend CLoadDataWithTranspose接口的 Python 实现支持通过 V1/V2 两套参数结构体完成带转置的 2D 格式数据搬运。读完本文你将掌握该接口的函数原型、全部参数取值范围与约束条件、可复制的 kernel 调用示例并能从 pyasc 源码角度理解一次 Python 调用是如何逐级映射为 IR 操作并最终发射出 Ascend C 代码的。接口概览与对应的 Ascend C 原型该接口实现带转置的 2D 格式数据从 A1/B1 到 A2/B2 的加载提供两个重载asc.language.basic.load_data_with_transpose(dst: LocalTensor, src: LocalTensor, params: LoadData2dTransposeParams) - None asc.language.basic.load_data_with_transpose(dst: LocalTensor, src: LocalTensor, params: LoadData2dTransposeParamsV2) - NoneLocalTensor 表示算子核内的本地内存张量是 pyasc 本地存储数据建模的核心类型。对应的 Ascend C 函数原型如下template typename T __aicore__ inline void LoadDataWithTranspose(const LocalTensorT dst, const LocalTensorT src, const LoadData2dTransposeParams loadDataParams)template typename T __aicore__ inline void LoadDataWithTranspose(const LocalTensorT dst, const LocalTensorT src, const LoadData2dTransposeParamsV2 loadDataParams)从源码结构看Python 接口与 Ascend C 原型一一对应src、dst均为LocalTensor参数结构体按 V1/V2 两套提供。这也解释了为什么该接口只支持 Local → Local 搬运如 A1/B1 → A2/B2不支持以GlobalTensor作为源——跨存储层级GM → Local的搬运需要走load_data等 G2L 类接口。参数说明dst 与 src 约束dst目的操作数类型为LocalTensor用于接收转置后的二维数据存储位置需属于VECIN/VECCALC/VECOUT中的一种起始地址需满足32 字节对齐要求。src源操作数类型为LocalTensor作为 2D 输入块的提供者仅支持 Local → LocalA1/B1 → A2/B2不支持GlobalTensor数据类型必须与 dst 一致。LoadData2dTransposeParamsV1 参数结构体字段含义取值范围默认值start_index方块矩阵 ID搬运起始位置为源操作数中第几个方块矩阵0 表示源操作数中第 1 个方块矩阵[0, 65535]0repeat_times迭代次数[0, 255]0src_stride相邻迭代间源操作数前一个分形与后一个分形起始地址的间隔单位为拼接后的方块矩阵大小[0, 65535]0dst_gap相邻迭代间目的操作数前一个迭代第一个分形的结束地址到下一个迭代第一个分形起始地址的间隔单位 512B[0, 65535]0dst_frac_gap每个迭代内目的操作数转置前一个分形结束地址与后一个分形起始地址的间隔单位 512B仅在数据类型为 float/int32_t/uint32_t/uint8_t/int8_t/int4b_t 时有效[0, 65535]0addr_mode预留参数—0LoadData2dTransposeParamsV2V2 参数结构体V2 在 V1 基础上新增了src_frac_gap字段addr_mode也变为可用的地址模式字段字段含义取值范围默认值start_index方块矩阵 ID搬运起始位置为源操作数中第几个方块矩阵0 表示第 1 个方块矩阵[0, 65535]0repeat_times迭代次数[0, 255]0src_stride源操作数步长[0, 65535]0dst_gap目的操作数间隔[0, 65535]0dst_frac_gap分形间隔[0, 65535]0src_frac_gap源分形间隔V2 新增[0, 65535]0addr_mode地址模式[0, 255]0两个参数结构体均可通过位置参数或关键字参数构造。从源码 types.py 和 types.py 可以看到构造时各字段会被物化为对应位宽的无符号整型 IR 值start_index/src_stride/dst_gap/dst_frac_gap/src_frac_gap为 uint16repeat_times/addr_mode为 uint8这与上表取值范围的上界 65535/255 完全吻合。一个需要注意的差异是文档标注repeat_times的字段默认值为 0而 Python 参数构造器的函数签名中 V1 版本该参数默认值为 1、V2 版本默认值为 0见 types.py 与 types.py。因此建议在实际 kernel 中总是显式写出repeat_times避免因默认值理解不一致而产生行为差异。调用示例V1 版本asc.jit def kernel_load_data_with_transpose(x: asc.GlobalAddress) - None: x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) y_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size512) params asc.LoadData2dTransposeParams(0, 4, 0, 0, 0, 0) asc.load_data_with_transpose(y_local, x_local, params)示例中参数含义依次为start_index0、repeat_times4、src_stride0、dst_gap0、dst_frac_gap0、addr_mode0从源操作数的第 1 个方块矩阵开始连续执行 4 次转置加载迭代。V2 版本asc.jit def kernel_load_data_with_transpose_v2(x: asc.GlobalAddress) - None: x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512) y_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size512) params_v2 asc.LoadData2dTransposeParamsV2(0, 4, 0, 0, 0, 0, 0) asc.load_data_with_transpose(y_local, x_local, params_v2)V2 示例在 V1 的 6 个位置参数之后多传入一个src_frac_gap0末尾的addr_mode仍为 0。两种写法可以出现在同一个 kernel 中接口会根据传入参数对象的类型自动选择对应的 IR 操作见下文源码解析。约束说明repeat_times为 0 时不执行搬运操作调用不会报错但也不产生任何数据移动开发者需要保证目的操作数转置后的分形之间没有重叠接口本身不校验重叠性操作数地址对齐要求需遵循《Ascend C 算子开发接口》中的通用说明和约束 - 通用地址对齐约束dst 起始地址至少 32 字节对齐从源码约束看src与dst的数据类型必须一致且src不能是GlobalTensor这两点由接口类型签名直接保证两个操作数都标注为LocalTensor。源码纵深从 Python 调用到 Ascend C 代码的完整链路1. 接口入口基于参数类型的重载分发接口定义在 mm.py通过overload声明两个类型签名实际实现使用OverloadDispatcher按参数类型自动分发overload def load_data_with_transpose(dst: LocalTensor, src: LocalTensor, params: LoadData2dTransposeParams) - None: ... overload def load_data_with_transpose(dst: LocalTensor, src: LocalTensor, params: LoadData2dTransposeParamsV2) - None: ... require_jit set_common_docstring(api_nameload_data_with_transpose) def load_data_with_transpose(dst: BaseTensor, src: BaseTensor, *args, **kwargs) - None: dispatcher OverloadDispatcher(__name__) builder global_builder.get_ir_builder() dispatcher.register_auto def _(params: LoadData2dTransposeParams): builder.create_asc_LoadDataWithTransposeOp( dst.to_ir(), src.to_ir(), params.to_ir()) dispatcher.register_auto def _(params: LoadData2dTransposeParamsV2): builder.create_asc_LoadDataWithTransposeV2Op( dst.to_ir(), src.to_ir(), params.to_ir()) dispatcher(*args, **kwargs)几个关键机制require_jit装饰器要求该接口只能在asc.jit标注的 kernel 编译上下文中调用这与所有示例代码都必须包裹在asc.jit函数内的用法一致传入 V1 参数对象时生成LoadDataWithTransposeOp传入 V2 参数对象时生成LoadDataWithTransposeV2Op两个重载因此可以在同一 kernel 中自由混用params.to_ir()触发参数结构体到 IR 常量值的物化即前文所述的create_asc_ConstructOp构造过程。2. IR 操作定义OpMm.td 中的 TableGen 声明两个操作在 Asc 方言中以 TableGen 形式定义于 OpMm.tddef AscendC_LoadDataWithTransposeOp : CopyToL0Opload_data_with_transpose, LoadDataWithTranspose, [AscFunc] { let description Perform 2D LoadData with transpose between tensors; let arguments (ins AscendC_LocalTensor:$dst, AscendC_LocalTensor:$src, AscendC_LoadData2dTransposeParams:$loadDataParams); ... } def AscendC_LoadDataWithTransposeV2Op : CopyToL0Opload_data_with_transpose_v2, LoadDataWithTranspose, [AscFunc] { let description Perform 2D LoadData with transpose (V2 params) between tensors; let arguments (ins AscendC_LocalTensor:$dst, AscendC_LocalTensor:$src, AscendC_LoadData2dTransposeParamsV2:$loadDataParams); ... }操作继承了CopyToL0Op基类Copy 到本地存储 0 层的通用类别并实现了AscFunc接口——意味着它只能出现在算子核函数体内。同时 IR 层面将dst/src都约束为AscendC_LocalTensor类型从机制上落实了仅支持 Local → Local的文档约束。3. 流水映射由 MTE1 搬运引擎执行在 Utils.cpp 中可以看到CopyToL0Op整类操作含本接口的两个变体被统一映射到Pipe::PIPE_MTE1.Case([](CopyToL0Op) { return Pipe::PIPE_MTE1; })这说明load_data_with_transpose在硬件上由 MTE1 搬运引擎执行与load_data、fixpipe等本地存储搬运/写入接口同属一个数据通路类别。4. 代码发射验证mlir 级测试Target 层的发射测试 mm.mlir 验证了 IR 到 Ascend C 的转换结果// CHECK-LABEL: void emit_load_data_with_transpose(AscendC::LocalTensorint32_t v1, ..., // AscendC::LoadData2dTransposeParams v3, AscendC::LoadData2dTransposeParamsV2 v4) func.func emit_load_data_with_transpose(...) { ascendc.load_data_with_transpose %v1, %v2, %v3 : ... ascendc.load_data_with_transpose_v2 %v1, %v2, %v4 : ... }即ascendc.load_data_with_transpose/ascendc.load_data_with_transpose_v2这两条 IR 指令会被发射为 Ascend C 中对应LoadDataWithTranspose模板函数的调用。5. 单元测试mock 启动方式下的用法基线单元测试 test_mm.py 提供了一个可参考的最小 kernel 模板与官方示例一致并通过kernel[1]()指定 1 个 AI Core 启动、断言 launcher 只被调用一次def test_load_data_with_transpose(mock_launcher_run): asc.jit def kernel_load_data_transpose() - None: x_local asc.LocalTensor( dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size512, ) y_local asc.LocalTensor( dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size512, ) params_v1 asc.LoadData2dTransposeParams(0, 4, 0, 0, 0, 0) params_v2 asc.LoadData2dTransposeParamsV2(0, 4, 0, 0, 0, 0, 0) asc.load_data_with_transpose(y_local, x_local, params_v1) asc.load_data_with_transpose(y_local, x_local, params_v2) kernel_load_data_transpose[1]() assert mock_launcher_run.call_count 1测试同时覆盖了 V1/V2 两个重载在真实 kernel 编译流程中的构建路径可以作为接口可用性的回归基线。小结与相关资源load_data_with_transpose是 pyasc 本地存储搬运族接口中专门处理转置 分形迭代场景的算子V1 参数结构体覆盖常规的源/目的步长与分形间隔控制V2 在此基础上增加了源分形间隔src_frac_gap与可用的地址模式字段。使用时的三个核心要点是——src/dst必须同为LocalTensor且类型一致、repeat_times0时不执行搬运、目的侧转置分形不得重叠。可继续深入的相关路径接口实现python/asc/language/basic/mm.py参数类型定义python/asc/language/core/types.pyIR 操作 TableGen 定义include/ascir/Dialect/Asc/IR/Basic/OpMm.td流水映射lib/Dialect/Asc/Utils/Utils.cpp发射测试test/Target/AscendC/basic/mm.mlir单元测试python/test/unit/language/basic/test_mm.py接口导出load_data_with_transpose由 python/asc/language/basic/init.py 从.mm模块统一导出同族接口普通 2D/3D 加载、稀疏加载load_data、load_data_with_sparse均位于 mm.py【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考