前言

在昇腾CANN生态中，算子开发与加速一直是构建高性能AI应用的核心环节。PyPTO作为一套面向昇腾NPU的Python端算子加速库，在CANN工具链中承担着关键角色：它让算法工程师能够用Python直接编写高性能算子，而无需接触底层的Ascend C语言或硬件指令集。PyPTO的全称是Parallel Tensor/Tile Operation，一套基于Tile计算模型的高性能编程框架。它的设计目标是在保持Python开发体验的同时，将计算密集型算子高效地映射到昇腾NPU的达芬奇架构上执行。

PyPTO的价值在于它重新定义了算子开发的抽象层次。传统的算子开发流程中，算法开发者需要用Ascend C编写算子核心逻辑，再通过框架适配器集成到训练或推理流程中。这种模式要求开发者同时理解硬件架构、编程语言和框架接口三个层面的知识。PyPTO将这个过程简化为一套Python API，开发者只需要描述计算逻辑，编译链路自动完成从高层Tensor图到底层可执行代码的转换。

这个框架自2025年12月首次上线以来，经历了多个版本的快速迭代。0.1.0版本在2026年1月发布，提供了基础的Python编程接口和编译能力。0.1.1版本引入了新前端表达方式，扩充了API覆盖范围。0.1.2版本增加了集群训练场景支持，优化了编译性能。到2026年4月的0.2.0版本，前端表达方法做了较大变更，整体功能和性能都得到了完善。

从架构角度看，PyPTO处于CANN五层架构的上层，更准确地说是连接应用开发层和编译层的桥梁。它接收用户通过Python API构建的计算图，经过多层中间表示转换和编译优化，最终生成可在昇腾NPU上执行的目标代码。这个定位决定了PyPTO的设计必须同时兼顾三方面的需求：表达能力要足够丰富以覆盖各类算子场景，编译效率要足够高以支持快速迭代，运行时性能要足够接近手写Ascend C算子的水平。

Tile计算模型的核心设计

PyPTO最核心的设计决策是采用Tile计算模型。Tile是硬件感知的数据块，它的尺寸与昇腾NPU上的Cube单元和Vector单元的处理粒度对齐。这种设计让编译器能够精确控制数据在寄存器和各级缓存之间的流动，充分利用硬件并行计算能力和内存层次结构。

传统Python算子库通常将整个张量作为操作对象，编译器只能做粗粒度的优化。PyPTO的做法是将大张量切分成多个Tile，每个Tile独立映射到计算单元上执行。这种设计让编译器能够精确控制每个Tile的数据搬运、计算和写回时机，避免无谓的全局内存访问。

在PyPTO中，一个典型的Tile遍历操作沿着张量的空间维度展开。
每个Tile包含足够的数据以充分利用NPU上的计算单元，
同时尺寸控制在片上缓存的容量范围内，避免数据溢出到主存。

这个模型的实际效果体现在两个方面。从内存带宽利用率来看，Tile化的数据流让每个数据在被加载到寄存器后都能被充分计算，提高了计算密度。从并行度来看，多个Tile可以在不同的计算单元上并行执行，且不需要显式同步，因为Tile之间没有数据依赖关系。

多层中间表示编译管线

PyPTO的编译系统使用多层中间表示(IR)来逐步降低计算描述的抽象层次。完整的转换路径从Tensor Graph开始，经过Tile Graph、Block Graph，到达Execution Graph，每一步都施加针对性的编译优化。

Tensor Graph是最高层的抽象，它直接对应开发者编写的Python代码。开发者在Python中创建的每个Tensor操作都被表示为图中的一个节点，Tensor之间的数据依赖关系构成图的边。这个阶段的信息密度最低，但可读性最强，适合做算子融合和内存规划等高层优化。

编译Pass将Tensor Graph转换为Tile Graph的过程，是整个系统中最关键的一步。这个阶段的核心工作是张量分片——将大张量拆分为多个尺寸一致的Tile，同时确定每个Tile所需的计算和数据依赖关系。Tile Graph保留了完整的Tile级数据流信息，是后续优化的基础。

从Tile Graph到Block Graph的转换，引入了硬件执行单元的概念。Block代表一个在NPU计算核上可独立执行的计算单元，包含完整的指令序列和数据搬运描述。多个Block之间通过MPMD（Multiple Program Multiple Data）模型调度执行。

Execution Graph是最终的可执行表示，每个节点对应一个具体的硬件操作序列。这个阶段的IR已经高度细化，包含寄存器分配、指令排序、内存地址计算等底层信息。

PyPTO的编译Pass系统采用模块化设计，每个Pass独立执行一种优化任务。Pass之间的执行顺序由编译器根据当前IR的类型和目标硬件的特性动态决定。这种设计让编译管线具有可扩展性——新的优化Pass可以插入到管线中任意位置，而不影响其他Pass的实现。

import pypto as pt

def fused_gemm_relu(x, w):
    a = pt.matmul(x, w)
    b = pt.relu(a)
    return b
# The compiler fuses matmul and relu into a single Tile operation,
# avoiding intermediate tensor materialization between the two ops.

上例中，开发者用Python API描述了一个矩阵乘法加ReLU激活的融合算子。PyPTO编译器在Tensor Graph阶段识别出matmul和relu之间的数据流关系，在Tile Graph阶段将两个操作合并为单个Tile级操作，避免了对中间张量a的显式存储和读取。这种融合策略减少了内存带宽消耗，并且消除了从Vector单元到Cube单元的数据搬运开销。

def tile_gemm(x_tile, w_tile):
    block = pt.Block()
    a_local = block.load(x_tile)
    b_local = block.load(w_tile)
    c_local = block.mma(a_local, b_local)
    out = block.store(c_local)
    return out
# This exposes the tile-level data movement explicitly.
# block.mma invokes the Cube unit's matrix multiply-accumulate,
# which runs at peak throughput only when tiles fit in local memory.

在Tile层次，开发者可以控制数据的显式搬运策略。block.load操作将Tile数据从全局内存加载到核间共享内存或寄存器中，block.mma操作调用Cube单元的矩阵乘加指令执行计算，block.store操作将结果写回。这个层次适合性能敏感的场景，开发者可以通过调整Tile尺寸和搬运策略来优化内存访问模式。

def element_wise_kernel(x):
    cfg = pt.TileConfig(shape=(128, 128), block_size=32)
    tiles = pt.split(x, cfg)
    results = []
    for t in tiles:
        y = t * 2.0 + 1.0
        results.append(y)
    out = pt.merge(results)
    return out
# Manual tile splitting lets the developer control L1 cache reuse.
# A tile size of 128x128 matches the Vector unit's preferred working set,
# ensuring each tile stays in on-chip memory throughout the computation.

上演示了手动的Tile划分。开发者指定Tile配置为128x128，每个Block处理32x32的子块。split操作将输入张量按配置切分，每个Tile独立计算，最终通过merge操作重新组装。这种模式在处理超大张量时特别有效，避免了因张量尺寸超出片上缓存而导致的频繁换入换出。

Python API的设计层次

PyPTO的Python API采用分层抽象设计，对不同技术背景的开发者暴露不同粒度的接口。这种设计让团队中的算法工程师和性能工程师可以在同一个框架下协作，而不用各自维护不同的工具链。

Tensor层次是面向算法开发者的默认编程接口。这个层次的API直接对标PyTorch的Tensor操作，支持常见的张量运算、索引、广播和shape推导。开发者用这个层次写代码，感觉就像在用高级框架写模型前向计算，完全不需要考虑硬件细节。

自动化代码生成与MPMD执行

PyPTO的CodeGen模块负责将Execution Graph转换为PTO虚拟指令序列。CodeGen的工作分为三个步骤：指令选择、指令调度和寄存器分配。

指令选择阶段，CodeGen为Execution Graph中的每个操作匹配可用的虚拟指令。不同的指令对应不同的计算单元——矩阵操作映射到Cube指令，向量操作映射到Vector指令，数据搬运映射到Load/Store指令。指令选择器会根据操作的数据类型和shape选择最优的指令变体。

指令调度阶段，CodeGen根据指令之间的数据依赖关系排出执行序列。调度的目标是最小化指令停顿，让计算指令和数据搬运指令尽量重叠执行。NPU上的Cube单元和Vector单元可以并行工作，调度器会利用这一点，在Cube计算的同时调度Vector搬运下一次计算所需的数据。

动态Shape编译处理策略

PyPTO对动态Shape的支持不是在运行时通过reshape或padding来实现的，而是在编译层面做了专门设计。当开发者传入的Tensor尺寸在编译时无法确定时，编译器不会退化为使用最大静态尺寸或动态内存分配，而是将Shape变量作为编译参数留到运行时绑定。

Tensor Graph在构建时会记录每个Tensor的维度符号表达式，而不是具体的整数尺寸。在Tile Graph阶段，Tile的划分方案以符号表达式描述，例如将M维度的Tile数量表示为ceil(M / tile_size)，其中M在编译时未知。这一步的关键在于Tile的边界处理逻辑需要兼容任意尺寸：当M不能被tile_size整除时，边界处的Tile包含的数据量小于标准Tile，编译器会生成一条额外的边界处理路径。

从Tile Graph到Block Graph的转换是符号表达式具体化的关键节点。如果某些维度尺寸在编译时已知，编译器在此时将符号表达式替换为整数常量，生成确定性的Block调度方案。如果所有维度均为动态，编译器生成一个包含条件分支的Block模板，运行时根据实际尺寸选择执行路径。

Python端集成的实践路径

PyPTO在Python端的集成主要以几个维度展开。对于已经使用PyTorch的开发团队，PyPTO提供了一套兼容的数据类型和算子接口。开发者可以将PyPTO的Tensor与PyTorch的Tensor互相转换，平滑地迁移现有代码中的计算密集型部分。

在推理场景中，PyPTO可以直接加载模型的权重数据，通过其编译管线生成NPU可执行代码。由于PyPTO支持动态Shape，同一个编译产物可以处理不同批次大小的输入，避免了静态Shape方案中因输入尺寸变化需要重新编译的问题。

训练场景的集成相对复杂，因为训练过程涉及前向计算、反向传播和参数更新三个步骤。PyPTO通过算子级别的支持来覆盖训练流程——开发者为前向算子编写对应的反向算子，PyPTO的编译系统自动将两者关联起来，保证反向传播的梯度计算与硬件特性对齐。

PyPTO在DeepSeek V3.2和GLM V4.5等大模型上已经有了实际落地的示例。在DeepSeek V3.2的稀疏Flash Attention量化实现中，PyPTO的Tile模型被用来处理稀疏注意力计算中的不规则数据访问模式，通过将稀疏矩阵分解为多个稠密Tile来匹配NPU的向量处理单元。在GLM V4.5的专家选择器实现中，PyPTO的动态Shape支持让每个token路由到不同专家网络的过程保持高效的Tile化计算。

使用前后的效率对比

PyPTO引入的执行模型通过在编译期规划数据布局和运算次序，在实际算子执行中产生了可衡量的效率差异。下面的表格从几个关键维度进行对比。

维度	使用前	使用后	差异来源
开发效率	需要精通Ascend C语言和硬件体系结构才能编写算子，单个算子开发周期通常以周计算	用Python Tensor API描述计算逻辑，编译管线自动生成NPU代码，开发周期缩短至天级	PyPTO的分层抽象将硬件细节隐藏在多层IR系统中，开发者只需关注算子语义
内存带宽利用率	手写算子中，内存访问模式往往缺乏系统化规划，CPU到NPU的数据搬运存在冗余	Tile化数据流将数据加载与计算紧密耦合，每个数据加载后得到充分计算	Tile计算模型将数据分片与NPU缓存层次对齐，编译器自动插入数据预取指令
算子融合深度	多个连续算子之间的融合依赖手工实现，融合策略受限于开发者的硬件理解	Tensor Graph上的融合Pass自动识别可融合的算子链，生成单Tile执行计划	多层IR系统让编译器能在高层图(Tensor Graph)进行全局优化，而非局部指令级优化
跨平台适配	为不同硬件平台编写不同的算子实现，代码维护成本随硬件版本线性增长	PTO虚拟指令作为独立于具体硬件的中间层，一次编写适配多代硬件	虚拟指令集架构在硬件抽象层和指令生成层之间插入了一个平台无关的隔离层

这四组对比反映了PyPTO设计决策的实际效果：在降低开发门槛的同时，通过编译优化保持了接近手写算子的运行效率。其中内存带宽利用率这个维度的提升最为直接——Tile化数据流从根本上减少了全局内存访问次数，这是AI算子的性能瓶颈所在。算子融合深度的提升则是在编译层面完成的，对开发者完全透明。跨平台适配能力的改善是一个附加价值——当昇腾硬件更新换代时，PTO虚拟指令的编译后端只需要适配新的硬件接口，上层Python代码无需改动。

需要注意的是，PyPTO并非在所有场景下都优于手写算子。当算子逻辑极度非规则（例如变长循环嵌套或条件分支密集的计算）时，编译器自动生成的指令序列可能不如手工编排紧凑。这种情况下，PyPTO允许开发者下沉到Tile或Block层次手动优化，在框架的统一架构下获得手写级别的控制粒度。

与CANN工具链中其他项目的关系

PyPTO在CANN工具链中的定位是面向Python开发者的高性能编程框架，它与Ascend C、ATB、catlass等项目存在互补关系。

Ascend C是昇腾算子编程语言，提供最底层的算子开发接口。用Ascend C编写的算子可以直接访问NPU的Cube单元和Vector单元，获得最高的硬件利用率。PyPTO的编译输出实际上是通过调用Ascend C的底层接口来生成可执行代码，可以理解为PyPTO在Ascend C之上构建了一层Python抽象，将Ascend C的编程细节屏蔽在编译器的实现中。

ATB（ascend-transformer-boost）是面向Transformer模型的加速库，包含FlashAttention、MoE等大规模融合算子。ATB的使用场景是在模型层面做算子编排和融合，而PyPTO的使用场景是在算子层面做计算描述和编译优化。两者可以协同工作，ATB可以调用PyPTO生成的算子，PyPTO也可以为ATB提供定制化的算子实现。

catlass是昇腾算子模板库，提供了Tile操作和矩阵运算的模板代码。PyPTO的Tile计算模型与catlass的模板设计在技术上有交集，但定位存在差异。catlass提供的是可复用的C++模板代码，供Ascend C开发者直接使用。PyPTO提供的是Python编程框架和编译系统，让开发者不用写C++代码也可以生成高性能算子。

编译中间产物的可视化分析

PyPTO工具链的一个设计特色是编译中间产物的可视化分析。编译管线产生的各层IR和运行时性能数据可以通过集成开发环境的工具链界面展示，开发者不需要手动解析中间文件就能理解编译器的决策过程。

Tensor Graph的展示以张量节点为基本单元，开发者可以检查每个Tensor节点的数据依赖关系和算子融合决策。如果两个Tensor节点被编译器融合为一个Tile操作，可视化界面会展示融合前后的算子拓扑对比，帮助开发者判断编译器的融合策略是否符合预期。

Tile Graph的展示切换到数据分片视角，重点显示每个Tile在NPU计算核上的映射关系。对于性能敏感的应用，开发者可以在这个层面检查是否存在Tile尺寸选择不佳导致的计算单元空转现象。如果多个Tile被映射到同一个计算核上，可视化界面会标记出潜在的负载不均衡问题。

PyPTO作为昇腾CANN工具链中的Python端编程框架，通过Tile计算模型和多层中间表示编译系统，为算子开发提供了从Python API到底层硬件执行代码的完整转换路径。它的分层抽象设计让不同技术背景的开发者可以在同一个框架下协同工作，Tensor层次适配算法工程师，Tile和Block层次服务性能专家。编译管线中的Tensor Graph、Tile Graph、Block Graph和Execution Graph四层IR，配合模块化的Pass系统，实现了算子融合、内存优化和指令调度等关键优化。在DeepSeek V3.2和GLM V4.5等大模型上的实践验证了PyPTO对复杂计算模式的承载能力。使用后的开发效率、内存带宽利用率和算子融合深度均获得明显改善，同时保留了下沉到Tile层次进行手工优化的能力，以应对非规则计算场景。PyPTO与Ascend C、ATB、catlass等项目形成了清晰的定位互补，构成了CANN工具链的多层次算子开发体系。

---

仓库地址：https://atomgit.com/cann/pypto