资讯动态

PyPTO Pro SIMD API:transpose 转置指令的完整用法与源码级解析

发布时间:2026/10/13 14:59:45 来源:尧图企业网站定制
人工智能编译器模型编译深度学习高性能计算CANNAscend【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载PyPTOPyPTO ProplDSL面向 Ascend 950PR/950DT 系列提供的pl.transpose指令可在片上 UBUnified Buffer内交换二维 Tile 的两个轴直接完成矩阵转置免去经 GM 中转的额外拷贝。本文基于仓库 API 文档 transpose.md完整覆盖产品支持、参数约束、可运行调用示例与实测输出并结合 语言层声明、IR 构建器、代码生成 与 单元测试、系统测试 逐层印证其实现细节读完后可在 kernel 中正确使用、排错并理解其底层调用链。产品支持情况pl.transpose属于 PyPTO Pro SIMD向量API 中的 Tile 计算类指令产品支持情况如下来自 transpose.md产品系列支持情况Ascend 950PR 950DT 系列产品支持Atlas A3 系列产品不支持Atlas A2 系列产品不支持因此在使用前需确认目标设备为 950 系列仓库中对应的系统测试也使用pytest.mark.soc(950)标记并做了设备检查见 test_transpose_scalar_access.py 中_require_a5(device)的调用。功能说明与函数原型pl.transpose交换二维 Tile 的两个轴实现矩阵转置若src的 shape 为[64, 128]则out的 shape 为[128, 64]且满足out[i, j] src[j, i]。函数原型pypto_pro.language.transpose(out: Tile, src: Tile) - None需要说明的是这些pl.xxx声明并不在 Python 运行时真正执行。语言层声明文件 的模块 docstring 明确写道这些声明的目的是让 IDE 跳转定义、拼写检查与类型检查可用在 PyPTO kernel 中AST parser 会在 Python 执行前拦截每一次pl.xxx调用并转入 IR 构建流程在 kernel 之外直接调用声明函数会抛出RuntimeError。transpose的声明本体位于 _api.py 的 B7 矩阵运算区块_api_decl def transpose(out: Tile, src: Tile) - None: Transpose Tile by swapping the last two dimensions. Args: out: Destination Tile src: Source Tile 从源码结构看它被归入 B7. Matrix ops矩阵运算分组与matmul、matmul_acc、matmul_mx等 Cube/UB 矩阵类指令相邻体现了它作为矩阵重排类指令的定位。参数说明与约束参数输入/输出说明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 字节对齐。上述约束并非只存在于文档中IR 构建器_ir_transpose在解析阶段就会强制检查def _ir_transpose(out: Expr, src: Expr, *, span: Span | None None) - Expr: dt getattr(out.type, dtype, None) if dt is not None and dt.get_bit() not in (8, 16, 32): raise InvalidType(ftranspose: unsupported dtype {dt}, supported: b8/b16/b32) src_dt getattr(src.type, dtype, None) if src_dt is not None and dt is not None and src_dt.get_bit() ! dt.get_bit(): raise InvalidType(ftranspose: dtype size mismatch between dst ({dt}) and src ({src_dt})) return _ir_core.create_op_call( block_ir_op(transpose), [out, src], {}, span or _span(), )由此可以确认两条与文档一致的硬性约束位宽限制out的 dtype 必须是 8/16/32 bit64 bit 元素如 FP64/INT64在解析期即被拒绝报错信息为transpose: unsupported dtype ...源/目标位宽一致src与out的 dtype 位宽必须相同否则抛出dtype size mismatch错误。另外操作数角色在 操作流水线注册表 中登记为transpose: [W, R]即第一个参数out是写Write角色第二个参数src是读Read角色。该表用于多核/自动互斥场景下推导 Tile 的读写角色并插入跨核同步也解释了文档中不可与src为同一 Tile的原因——原地转置同一 Tile 既是写又是读不在支持范围内必须准备独立的目的地 Tile通常分配在不同的 UB 地址段。调用示例完整的转置 kernel文档给出的基本用法完整可运行。整体流程为创建两个 UB Tile Group源与目的使用不同地址、不同互斥 ID进入向量 sectionload从 GM 读入源 Tile执行transpose最后store回 GMimport 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])几个关键点的解读结合 _api.py 中 load/store 声明pl.TileType(shape[64, 64], dtypepl.DT_FP16, target_memorypl.MemorySpace.Vec)声明一个位于 UBVec 内存空间的 64×64 FP16 Tile 类型。pl.transpose要求操作数是二维 UB即target_memory为MemorySpace.Vec的 Tile。pl.make_tile_group(typett, addrs0x0000, mutex_ids[0])以 Tile Group多缓冲区组形式分配 UB 地址源 Tile 在0x0000、目的 Tile 在0x400064×64×2B8KB地址不重叠并分别绑定互斥 ID 0 和 1配合pl.jit(auto_mutexTrue)让框架自动推导加锁/同步逻辑。pl.load(cur_a, a, [0, 0])load(dst_tile, src_tensor, /, offsets, *, orderNone)按元素坐标从 GM Tensor 读入片上 Tileoffsets为每轴元素级偏移示例中从(0, 0)开始读取整个 64×64 块。pl.transpose(cur_out, cur_a)在 UB 内完成转置无需 GM 往返。pl.store(out, cur_out, [0, 0])store(dst_tensor, src_tile, /, offsets, ...)将转置结果写回 GM。文档同时给出了实测结果示例输入是一个元素按行号*16 列偏移*0.25规律递增的 64×64 矩阵输出矩阵各行变为输入矩阵的对应列输入数据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 ...], ...]可以看到第 0 行输出[1, 17, 33, 49, ...]正好是输入矩阵第 0 列的首元素序列符合out[i, j] src[j, i]的语义。该示例对应的可执行验证即 test_transpose_scalar_access.py 中的transpose_kernel与test_transpose用例用torch.rand(64, 64, fp16)作为输入运行后以torch.testing.assert_close(out, a.t().contiguous(), rtol1e-2, atol1e-2)断言结果与 PyTorch 的a.t()一致。编译与执行流程从pl.transpose到 TileOp 调用结合仓库源码一次pl.transpose调用会经历如下链路AST 拦截与角色登记parser 在 kernel 内拦截pl.transpose(out, src)查 _op_pipeline.py 的_BLOCK_OP_TILE_ROLES得到参数角色[W, R]供自动互斥auto_mutex推导同步关系。IR 构建与类型检查分发到register_table中注册的transpose: OpSpec(builder_ir_transpose)由_ir_transpose完成 8/16/32 bit 位宽检查与源/目标 dtype 位宽一致性检查然后通过_ir_core.create_op_call(block_ir_op(transpose), [out, src], {}, span)生成 IR op 调用。单元测试test_manual_transpose验证了这一点对 shape[16, 32]的src与[32, 16]的dst调用pl.transpose(dst, src)后解析出的程序 IR 文本中包含block.transpose节点。向量化代码生成在 NPU codegen 阶段codegen_vector.cpp读取 op 属性中的shapetranspose 的轴描述并输出对应 TileOp 模板调用std::vectorint64_t transposeAxis AnyCaststd::vectorint64_t(opAttrs.at(OP_ATTR_PREFIX shape)); oss tileOpName (transposeAxis[0] correctionAxis) , (transposeAxis[1] correctionAxis)从这段代码结构看生成的 CCE 代码形如transposeaxis0, axis1...的 TileOp 模板实例化转置轴由 IR 属性携带并加上维度校正量后输出。这印证了交换二维 Tile 的两个轴在底层对应一条带轴参数的向量 TileOp 指令。 4.运行时执行kernel 由pl.jit编译并在向量流水V 流水上执行 UB 内转置load/store则经由 DMAMTE完成 GM 与 UB 之间的数据搬运与转置指令本身解耦。相关用法与 load 的order转置搬运的区别在使用pl.transpose时值得注意一个替代路径数据搬运指令pl.load支持order参数——升序如[0, 1]为不转置加载降序如[1, 0]则直接以转置方式加载DN 排布pl.load(l1, x, [b, s0, h, 0, k0], order[1, 3, 4]) # NZ [S_tile*G, D_tile] pl.load(l1, x, [b, s0, h, 0, k0], order[4, 3, 1]) # ZN [D_tile, S_tile*G]从源码注释看load的order降序形式DN layout与pl.transpose的适用场景不同前者在 GM→L1/UB 搬运阶段由 DMA 一次完成排布变化对 matmul 的 ZN 操作数尤其高效same bytes, costs nothing后者是 UB 内两个已有 Tile 之间的计算转置。若转置只发生在从 GM 取出时优先考虑load/store的order降序写法若转置发生在片上计算中途如 attention 中 QK^T 结果的 ND/DN 排布互换则使用pl.transpose。仓库中 test_matmul_softmax.py 即在 attention 流程中多次调用pl.transpose(cur_qk_dn, cur_qk_nd)等完成片上排布转换属于典型用例。验证方式与测试参考仓库中与该指令相关的可查证测试包括类型路径验证内容单元测试IRtest_block_ops.pytest_manual_transpose[16,32]→[32,16]转置后程序 IR 含block.transpose节点系统测试端到端test_transpose_scalar_access.py64×64 FP16 转置 kernel与torch的a.t()数值比对前端组合测试test_abs_addc_and.pypl.transpose与 abs/addc/and 等指令的组合调用链代码生成测试test_cce_tensor_order_variants.py含转置语义的 load/order 变体在 CCE 代码生成上的正确性执行这些测试的前提是具备 950 系列设备stest 通过pytest.mark.soc(950)过滤UT 层面的 IR 断言则不依赖硬件。小结pl.transpose(out, src)是 PyPTO Pro SIMD API 中面向 Ascend 950PR/950DT 的 UB 内二维 Tile 转置指令交换两个轴outshape 为src的转置结果两者不可为同一 Tile仅支持 8/16/32 bit dtypeDT_INT8/UINT8/INT16/UINT16/FP16/BF16/INT32/UINT32/FP32源/目标位宽必须一致且主维长度×元素字节数须 32 字节对齐这些约束在_ir_transpose中于解析期强制执行典型用法是load → transpose → store三步用make_tile_group为源/目的 Tile 分配不同 UB 地址与互斥 ID编译链路上pl.transpose经角色表[W, R]与 IR 注册表落到block.transposeop最终由codegen_vector.cpp生成带轴参数的 TileOp 模板调用若转置可合并进 GM 搬运阶段可优先考虑pl.load/pl.store的order降序参数以避免额外的 UB 转置与地址开销。参考文档入口SIMD-API 总目录、转置与元素读写目录、transpose 指令文档。赞分享人工智能编译器模型编译深度学习高性能计算CANNAscend【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PyPTO SIMD-API 高级计算指令详解interleave、de_interleave、pack、unpack 与 unsqueezePyPTO SIMD API 高级计算指令详解interleave、de_interleave、pack、unpack 与 unsqueeze 本文基于 Py人工智能编译器模型编译深度学习高性能计算CANNAscendPHP-CS-Fixer 完整指南安装、配置、命令用法与源码级原理解析PHP CS Fixer 完整指南安装、配置、命令用法与源码级原理解析 PHP Coding Standards FixerPHP CS Fixer是一款开发工具代码质量静态分析Lint格式化PyPTO vf.ge 逐元素大于等于比较指令详解Ascend 950 SIMD-API 掩码生成与实战用法PyPTO vf.ge 逐元素大于等于比较指令详解Ascend 950 SIMD API 掩码生成与实战用法 导读 vf.ge 是 CANN PyPTO 项目人工智能编译器模型编译深度学习高性能计算CANNAscend创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑