pyasc 算子开发指南：load_data_with_transpose 实现带转置的 2D 数据加载

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) -> None

LocalTensor 表示算子核内的本地内存张量，是 pyasc 本地存储数据建模的核心类型。

对应的 Ascend C 函数原型如下：

template <typename T> __aicore__ inline void LoadDataWithTranspose(const LocalTensor<T>& dst, const LocalTensor<T>& src, const LoadData2dTransposeParams& loadDataParams)

template <typename T> __aicore__ inline void LoadDataWithTranspose(const LocalTensor<T>& dst, const LocalTensor<T>& 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 → Local（A1/B1 → A2/B2），不支持GlobalTensor；
  • 数据类型必须与 dst 一致。

LoadData2dTransposeParams（V1 参数结构体）

LoadData2dTransposeParamsV2（V2 参数结构体）

V2 在 V1 基础上新增了src_frac_gap字段，addr_mode也变为可用的地址模式字段：

两个参数结构体均可通过位置参数或关键字参数构造。从源码 types.py 和 types.py 可以看到，构造时各字段会被物化为对应位宽的无符号整型 IR 值（start_index/src_stride/dst_gap/dst_frac_gap/src_frac_gap为 uint16，repeat_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(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=512) y_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=512) params = asc.LoadData2dTransposeParams(0, 4, 0, 0, 0, 0) asc.load_data_with_transpose(y_local, x_local, params)

示例中参数含义依次为start_index=0、repeat_times=4、src_stride=0、dst_gap=0、dst_frac_gap=0、addr_mode=0：从源操作数的第 1 个方块矩阵开始，连续执行 4 次转置加载迭代。

V2 版本

@asc.jit def kernel_load_data_with_transpose_v2(x: asc.GlobalAddress) -> None: x_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=512) y_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=512) 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_gap=0，末尾的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_name="load_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.td：

def AscendC_LoadDataWithTransposeOp : CopyToL0Op<"load_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 : CopyToL0Op<"load_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::LocalTensor<int32_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( dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=512, ) y_local = asc.LocalTensor( dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=512, ) 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_times=0时不执行搬运、目的侧转置分形不得重叠。

可继续深入的相关路径：

• 接口实现：python/asc/language/basic/mm.py
• 参数类型定义：python/asc/language/core/types.py
• IR 操作 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