ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PyPTO `pypto_pro.language.transpose` 详解:二维 Tile 转置 API 的用法、约束与底层实现

PyPTO `pypto_pro.language.transpose` 详解:二维 Tile 转置 API 的用法、约束与底层实现 人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载导读pypto_pro.language.transpose是 PyPTOParallel Tensor/Tile Operation 编程范式中用于交换二维 Tile 两个轴、实现矩阵转置的核心向量指令 API。本文以 transpose.md 文档为主体完整讲解其产品支持情况、函数原型、参数与对齐约束、典型调用示例并结合仓库中 Python API 声明、NPU 代码生成实现 与 ST 测试用例帮助读者既能直接上手编写转置 kernel又能理解转置在编译器后端的真实映射路径。产品支持情况该 API 对不同昇腾平台的支持情况如下以当前仓库文档声明为准平台支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持可以看到transpose属于 Ascend 950 系列新增/启用的向量能力A2/A3 系列产品不提供该指令支持。对应的 ST 测试用例同样通过pytest.mark.soc(950)标注了运行平台见 test_abs_addc_and.py与文档保持一致。在实际工程中使用前应先确认目标设备属于 Ascend 950 系列。功能说明transpose的功能为交换二维 Tile 的两个轴实现矩阵转置。即对于src中坐标为(i, j)的元素转置后写入out的(j, i)位置。它属于向量Vector域的片内计算原语与通过 GM 往返的软件转置方案不同它直接在片上完成数据排布变换。从 Python 层 API 的 docstring 可以看到其语义被精确定义为 Transpose Tile by swapping the last two dimensions交换最后两个维度见 _api.py与文档中的交换二维 Tile 的两个轴表述一致。函数原型pypto_pro.language.transpose(out: Tile, src: Tile) - Nonetranspose接收两个 Tile 类型参数原地改写out不返回任何值。它是同步的向量指令 API必须在pl.section_vector()向量节section内调用通常与pl.load/pl.store配合完成GM → UB → 转置 → UB → GM的完整数据流。参数说明参数输入/输出说明out输出目的操作数Tile 类型存放转置结果。dtype 与 src 一致shape 为 src 转置后的结果如 src 为 [64, 128] 时 out 为 [128, 64]不可与 src 为同一 Tile。src输入源操作数Tile 类型二维 UB。dtype 支持 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_FP16、DT_BF16、DT_INT32、DT_UINT32、DT_FP32不支持 64 bit 元素。源、目标主维长度乘元素字节数须 32 字节对齐。围绕参数说明有几点值得展开dtype 家族支持 1/2/4 字节宽度的整型与浮点类型INT8/UINT8、INT16/UINT16、FP16/BF16、INT32/UINT32、FP32但不支持 64 bit 元素如 FP64/INT64。原因在于转置的片上搬移以 32 字节为基本对齐单元64 bit 元素在该对齐策略下无法高效编排。Python API 声明中同样没有对 dtype 做额外限制合法性与否完全由硬件指令决定见 _api.py。对齐约束源、目标主维长度乘元素字节数须 32 字节对齐——即对形状[M, N]的 Tile要求N × sizeof(dtype)为 32 的整数倍。这是transpose能否被静态代码生成路径而非退化的非对齐路径处理的关键判据详见下文约束与后端分支小节。out 与 src 不可同一 Tile转置是轴交换若原地操作会破坏尚未读到的源数据因此必须在 UB 中为out单独分配地址空间。约束说明文档明确无除参数说明中已列出的 dtype 与对齐要求外无其他约束。补充说明虽然 API 本身无额外约束但在实际 kernel 编写中仍要遵守 PyPTO 向量编程的通用约定——转置前需通过pl.load将数据从 GM 载入 UB转置后通过pl.store写回跨流水如 MTE2→V→MTE3的依赖需要在向量节内用同步原语显式建立具体可参考下文测试示例。返回值说明无返回值- None。转置结果直接写入out指向的 UB 空间属于典型的就地写目的操作数型向量 API。调用示例基本用法官方示例import pypto_pro.language as pl pl.jit(auto_mutexTrue) def transpose_kernel( a: pl.Tensor[[64, 64], pl.DT_FP16], out: pl.Tensor[[64, 64], pl.DT_FP16], ): tt pl.TileType(shape[64, 64], dtypepl.DT_FP16, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_out pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) with pl.section_vector(): cur_a tile_a.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.transpose(cur_out, cur_a) pl.store(out, cur_out, [0, 0])实测结果示例如下64×64 FP16 数据输入数据a[[1 1.25 1.5 1.75 2 2.25 2.5 2.75 ...], [17 17.25 17.5 17.75 18 18.25 18.5 18.75 ...], [33 33.25 33.5 33.75 34 34.25 34.5 34.75 ...], [49 49.25 49.5 49.75 50 50.25 50.5 50.75 ...], ...] 输出数据out[[1 17 33 49 65 81 97 113 ...], [1.25 17.25 33.25 49.25 65.25 81.25 97.25 113.25 ...], [1.5 17.5 33.5 49.5 65.5 81.5 97.5 113.5 ...], [1.75 17.75 33.75 49.75 65.75 81.75 97.75 113.75 ...], ...]对示例代码做逐行拆解pl.TileType(shape[64, 64], dtypepl.DT_FP16, target_memorypl.MemorySpace.Vec)声明 64×64 FP16、位于 UBVec 内存空间的 Tile 类型。主维长度 64 × 元素 2 字节 128 字节满足 32 字节对齐要求。pl.make_tile_group(typett, addrs0x0000, mutex_ids[0])/addrs0x4000在 UB 中为源、目的分配两个不重叠的地址区间0x0000 与 0x4000后者恰为 64×64×2 字节避免 out 与 src 复用同一 Tilemutex_ids用于auto_mutexTrue时自动生成互斥同步。pl.section_vector()进入向量指令节transpose在硬件向量流水Vector pipeline上执行。pl.load→pl.transpose→pl.store构成GM 载入 UB → 片内轴交换 → UB 写回 GM的完整数据流。实测输出中out第 0 行为a第 0 列1, 17, 33, 49, …正是标准的矩阵转置语义。带显式流水同步的写法ST 测试参考仓库 ST 测试 test_abs_addc_and.py 给出了更贴近实际工程、显式管理流水依赖的写法FP32 与 INT16 两个用例pl.jit() def kernel_transpose_fp32( a: pl.Tensor[[DYN, DYN], pl.DT_FP32], out: pl.Tensor[[DYN, DYN], pl.DT_FP32], ): tf pl.TileType(shape[T_M, T_N], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) ta pl.make_tile(tf, addr0) tc pl.make_tile(tf, addrT_M * T_N * 4) with pl.section_vector(): pl.load(ta, a, [0, 0]) pl.system.sync_src(set_pipepl.PipeType.MTE2, wait_pipepl.PipeType.V, event_id0) pl.system.sync_dst(set_pipepl.PipeType.MTE2, wait_pipepl.PipeType.V, event_id0) pl.transpose(tc, ta) pl.system.sync_src(set_pipepl.PipeType.V, wait_pipepl.PipeType.MTE3, event_id1) pl.system.sync_dst(set_pipepl.PipeType.V, wait_pipepl.PipeType.MTE3, event_id1) pl.store(out, tc, [0, 0])pl.system.sync_src/sync_dst在loadMTE2 流水与transposeV 流水、transpose与storeMTE3 流水之间显式建立依赖保证转置读到的是完整的载入数据、store 写出的转置结果未被覆盖。测试断言使用torch.testing.assert_close(out, a.t(), rtol1e-5, atol1e-5)见 test_abs_addc_and.py即与 PyTorch 的.t()转置结果逐元素比对验证片上转置与框架语义完全一致。动态形状用法ST 测试中的 kernel 均以pl.Tensor[[DYN, DYN], ...]声明动态形状输入DYN 为动态维度标记Tile 形状在编译期固定为[T_M, T_N]如 8×8、16×16运行时通过pl.load(ta, a, [0, 0])从 GM 的[0, 0]偏移处取块。这说明transpose既可配合静态形状使用也适用于GM 动态、Tile 静态分块的典型 tiling 场景后端会根据是否为动态函数自动选择不同的代码生成分支见下节。源码级实现transpose 在编译器后端的映射transpose并非 Python 层的普通函数而是经由_api_decl声明的内建 API 原语见 _api.py并被注册进语言命名空间见 language/init.py 与 #L384。编译期它会被翻译为 tile 图上的转置算子最终由 NPU 代码生成器产出底层指令。算子到指令的分发在 NPU 代码生成器的算子分发表中转置被映射为两个方向的片内数据搬移指令见 codegen_op_npu.cpp{Opcode::OP_TRANSPOSE_MOVEOUT, [this]() { return GenTransposeDataMove(); }}, {Opcode::OP_TRANSPOSE_MOVEIN, [this]() { return GenTransposeDataMove(); }},OP_TRANSPOSE_MOVEIN与OP_TRANSPOSE_MOVEOUT分别对应转置过程中 GM↔UB 两个方向的数据搬移共用同一个生成函数GenTransposeDataMove。此外还有OP_TRANSPOSE_VNCHWCONVV 流水上的 NCHW 卷积相关转置见同文件第 149 行属于另一类用途。转置代码生成的四种分支GenTransposeDataMove先根据算子方向确定 GM 与本地UB操作数的索引isCopyLocalToGM判断是UB→GM还是GM→UB见 codegen_vector.cpp随后依据硬件能力与函数形态选择具体打印策略见 codegen_vector.cpp分支触发条件特点PrintTransposeDataMoveLayoutisSupportTileTensor支持 Tile Tensor 布局生成带Coord坐标的模板指令转置轴以模板参数(axis0corr), (axis1corr)直接编码进指令见 codegen_vector.cppPrintTransposeDataMoveDynamicisSupportDynamicAligned支持动态对齐形状、偏移均以运行时宏PREFIX_STR_RAW_SHAPE/PREFIX_STR_OFFSET参数化适用于动态形状场景见 codegen_vector.cppPrintTransposeDataMoveDynamicUnalignedisDynamicFunction且不支持动态对齐动态形状的非对齐退路路径PrintTransposeDataMoveStatic静态形状默认路径形状在编译期已知直接以模板参数展开GM 指针用__gm__空间修饰、UB 指针用__ubuf__修饰见 codegen_vector.cpp四类分支的模板参数中都包含dtype 字符串、归一化到 4 维的 shape、以及经correctionAxisSHAPE_DIM4 - 实际维度修正后的转置轴索引——从源码结构可以推断这是为了将二维转置统一映射到底层 4 维数据搬移指令的参数规范。无论哪个分支最终产出的都是带dtype, shape..., axis...模板参数的转置搬移指令调用。与参数说明的印证代码生成中出现的srcShape/gmShape归一化NormalizeShape(..., SHAPE_DIM4)与源、目标主维长度乘元素字节数须 32 字节对齐的约束相呼应只有满足对齐的静态形状才能走高效的模板展开路径非对齐/动态场景则退化为带运行时参数的分支。生成指令中 GM 端使用(__gm__ dtype*)、UB 端使用(__ubuf__ dtype*)显式区分地址空间见 codegen_vector.cpp印证了src必须位于 UB二维 UB的文档约束。常见使用场景与注意事项对齐检查先行使用前确认主维长度 × 元素字节数 % 32 0。例如 8×8 FP328×432 字节与 16×16 INT1616×232 字节均满足这也是 ST 测试选取这两个用例的原因。内存布局规划out与src不能是同一 Tile需在 UB 中分配不重叠地址。官方示例使用 0x0000/0x4000测试用例则通过addr0与addrT_M * T_N * 4计算偏移均可参考。流水同步不可省略在向量节内loadMTE2→ transposeV→ storeMTE3之间的依赖建议用pl.system.sync_src/sync_dst显式同步或借助auto_mutexTrue自动互斥否则可能出现读未载入、写被覆盖的竞态。仅限 Ascend 950 系列A2/A3 平台不支持该指令跨平台代码需做能力判断或改用 GM 往返转置方案。总结pypto_pro.language.transpose提供了交换二维 Tile 两个轴的片上转置能力官方文档给出了完整的原型、参数、对齐约束与可运行示例仓库源码进一步印证了其实现路径Python 层经_api_decl注册为内建 APItile 图编译期翻译为OP_TRANSPOSE_MOVEIN/OP_TRANSPOSE_MOVEOUT算子NPU 代码生成器根据 Tile Tensor 支持、动态对齐与静态形状等条件选择四种搬移指令打印分支最终产出带轴参数的转置数据搬移指令。配套的 ST 测试test_abs_addc_and.py通过torch.testing.assert_close(out, a.t())验证了该 API 与标准转置语义的一致性。读者可对照本文示例直接编写转置 kernel并在需要时深入 codegen_vector.cpp 与 codegen_op_npu.cpp 追踪更底层的指令生成细节。赞分享人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PyPTO Tensor API 详解pypto.isnan 逐元素 NaN 检测算子的用法、约束与底层实现PyPTO Tensor API 详解pypto.isnan 逐元素 NaN 检测算子的用法、约束与底层实现 导读 pypto.isnan 是 CANN/Py人工智能编译器模型编译高性能计算深度学习CANNPyPTO pypto.atan2 详解从 API 调用、约束原理到 Tile 级反正切实现PyPTO pypto.atan2 详解从 API 调用、约束原理到 Tile 级反正切实现 pypto.atan2 是 PyPTO 提供的逐元素双变量反正切人工智能编译器模型编译高性能计算深度学习CANNPyPTO ceil_div 整数上取整除法API 约束、TileShape 配置与底层实现原理PyPTO ceil_div 整数上取整除法API 约束、TileShape 配置与底层实现原理 本文围绕 PyPTO Tensor API 中的 ceil_人工智能编译器模型编译高性能计算深度学习CANN上一篇LeRobot机器人学习框架3步跑通策略训练下一篇claude-code-ide.el未来发展路线图即将推出的令人兴奋的新功能创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表