前言

在昇腾NPU上进行深度学习算子开发，传统方案要求开发者同时具备算法设计能力和底层硬件知识。CCE（CANN Compute Engine）提供的是一套基于C++的算子开发接口，开发者需要手工处理数据布局、切分策略、指令调度等细节，开发周期长、调试成本高。伴随着大模型对融合算子需求的快速增长，业界对一种既能保持高性能、又能降低开发门槛的编程范式提出了迫切需求。

CANN社区推出的PyPTO（发音为pai p-t-o，全称Parallel Tensor/Tile Operation）正是在这一背景下诞生的产物。PyPTO是一款面向AI加速器的高性能编程框架，采用PTO编程范式，以基于Tile的编程模型为核心设计理念，通过多层次的中间表示（IR）系统，将用户通过Python API构建的AI模型应用从高层次的Tensor图逐步编译成CCE硬件指令，最终生成可在昇腾NPU上高效执行的可执行代码。整个编译管线的每一级都有明确的语义边界：Python AST解析生成Tensor Graph，Tensor Graph经TACO模块转换为Tile Graph，Tile Graph经分区生成Block Graph，Block Graph编排为Execute Graph，Execute Graph经CodeGen生成PTO虚拟指令，PTO虚拟指令经编译器编译为CCE指令与二进制。

本文以手把手实战为导向，剖析PyPTO的四层架构设计，随后深入走读编译管线的每一个环节，包括Python AST解析、TACO调度、PTO IR生成、CCE代码编译，然后通过一个完整的Softmax算子实例演示从代码编写到运行的全部步骤，最后给出CCE C++与PyPTO两种开发方式的生产效率对比。

一、PyPTO四层架构全景

PyPTO框架的分层设计是其架构可维护性和可扩展性的核心保障。从用户编写的Python代码到最终在昇腾NPU上运行的二进制指令，整个链路自上而下分为四个层次：用户API层、计算图编译层、代码生成层、调度执行层。每一层都有明确的职责边界，层与层之间通过标准化的中间表示（IR）传递信息。

1.1 用户API层

用户API层是PyPTO框架与开发者交互的接口层，提供Python友好的编程接口，使开发者能够以直观的方式表达计算逻辑，而无需深入了解底层硬件实现细节。当前版本开放的Tensor层次编程是最常用和推荐的编程方式，开发者直接使用Tensor和Tensor Operation构建计算图，无需关心底层的Tile切分和硬件细节。框架还提供了@pypto.frontend.jit装饰器，在第一次调用该Kernel时触发JIT编译，将计算图转换为硬件指令。

对于性能敏感场景，PyPTO同样暴露了Tile层次和Block层次的编程接口：Tile层次编程以Tile和Tile Operation显式体现访存与依赖，适合性能专家进行深度调优；Block层次编程定义单个处理器核执行的计算图，并通过多次实例化实现整体计算，适合系统开发者进行底层对接和工具链开发。

1.2 计算图编译层

计算图编译层负责将Tensor Graph转换为Tile Graph、Block Graph和Execute Graph，每一步都包含一系列模块化的Pass优化流程。该层是PyPTO编译管线的核心所在，也是理解PTO编译器行为的关键。

Tensor Graph阶段对应高层次抽象，所有计算操作以Tensor为基本单位进行表达，与硬件实现无关。在该阶段，框架执行冗余操作消除、类型转换优化、内存冲突推断等与硬件无关的图优化。这一阶段保留了最大化的优化空间，包括内存布局优化、数据搬运优化和多算子联合优化。

Tile Graph阶段根据开发者指定的TileShape进行Tile展开，实现Tile级别的优化，包括内存类型分配（UB、L1等）、移动操作生成、子图切分等。开发者可以通过pypto.set_vec_tile_shapes()或pypto.set_cube_tile_shapes()接口指定切分策略，框架据此将Tensor操作转换为硬件感知的Tile操作。

Block Graph阶段将Tile图分区为计算子图，进行Block级别的优化，包括乱序调度、内存重用和同步点插入等。框架检测同构子图并规范化Block Graph，追踪依赖关系。

Execute Graph阶段整合计算子图信息，构建最终的执行图。该阶段分析Block Graph之间的依赖关系，规划全局资源，生成调度提示。

1.3 代码生成层

代码生成层将优化后的Execute Graph转换为目标平台的可执行代码，分为两个子步骤：CodeGen模块从Execute Graph生成PTO虚拟指令代码（PTO Virtual Instructions），这是一种与硬件无关的中间指令集表达。其次，通过编译器将PTO虚拟指令编译为目标NPU平台的CCE指令，即最终可在昇腾硬件上执行的机器码。

PTO虚拟指令的设计屏蔽了不同硬件平台的差异，使得上层编译优化逻辑可以在不考虑具体硬件指令集的情况下进行，同时也便于工具链对编译过程进行可视化和调试。CodeGen模块根据目标昇腾NPU的硬件特性（如AIC向量计算能力、AIV张量计算能力、Cube矩阵乘法单元等）生成对应的CCE指令序列，处理指令调度、寄存器分配和指令重排。

1.4 调度执行层

调度执行层负责将可执行代码在昇腾NPU设备上调度执行。PyPTO基于MPMD（Multiple Program Multiple Data）执行模型，与传统的SPMD（Single Program Multiple Data）模型形成鲜明对比。在SPMD模型下，用户需要编写单一内核逻辑并实例化到多个处理器核上运行，这不可避免地引入全局同步开销和性能瓶颈。而MPMD模型将计算抽象为一组异构任务，任务之间通过依赖关系组织，运行时调度器根据依赖关系将任务分配到合适的执行单元，避免了全局同步限制，提升了整体利用率与效率。

可执行代码被加载到昇腾NPU设备侧后，通过MPMD方式调度到设备上的处理器核（包括AIC核和AIV核），实现细粒度的并行计算。开发者可以通过PyPTO Toolkit可视化工具观测每个核上的任务执行情况，识别性能瓶颈。

二、PTO编程范式与DSL语法

2.1 PTO范式的核心思想

PTO（Parallel Tensor/Tile Operation）编程范式的核心设计理念建立在四个支柱之上：第一，Tensor级别抽象，以Tensor而非单个元素描述计算，贴近算法设计者的数学表达式；第二，声明式编程，开发者只需描述"做什么"，框架自动处理"怎么做"；第三，基于Tile的计算，所有计算最终都基于Tile（硬件感知的数据块）进行，充分利用硬件并行计算能力；第四，计算图驱动，通过构建计算图，框架可以自动进行优化、调度和执行。

在PyPTO中，所有运算都以Tensor为输入或输出，形成可追溯的计算图结构。这种设计使得开发者能够专注于算法逻辑本身，而将性能优化的工作交给编译器完成。以矩阵乘法为例，传统CCE C++开发需要手动处理数据在UB和GM之间的搬运、Tile级别的循环切分、以及Cube指令的调度，而使用PyPTO只需三行核心代码即可完成等效的计算描述。

2.2 核心数据结构

Tensor是PyPTO中最基本的数据结构，表示一个多维数组，包含以下核心属性：数据类型（dtype）如FP32、FP16、INT32、BOOL等；形状（shape）用一个整型数组描述各维度长度，支持动态维度（用-1表示）；格式（format）描述数据在内存中的排布方式；名称（name）用于在计算图中标识该Tensor。

Tile是Tensor的子区间，通过Tiling（切分）将大的Tensor切分为多个子块。Tile的设计目的是使其能够存放在处理器核内私有缓存（如UB、L1）中，以提升数据局部性，同时充分利用硬件并行计算能力。在Tensor层次编程中，Tiling由框架自动完成，开发者只需通过配置接口指定TileShape，框架会自动进行切分。

View与Assemble提供了对子Tensor的视图和组合操作，在处理动态Shape和循环计算中非常有用。View允许在不复制数据的情况下访问Tensor的子区间，Assemble则将多个子Tensor组合成一个更大的Tensor。

2.3 @pypto.frontend.jit装饰器与Kernel定义

PyPTO通过@pypto.frontend.jit装饰器将用户定义的Python函数转换为可编译的Kernel。在首次调用该Kernel时，框架会触发JIT编译流程，自动解析函数体内的Python AST，构建计算图，并经过编译管线生成可执行代码。以下代码展示了使用装饰器定义一个向量加法算子的完整流程：

import pypto

# 配置Tile切分策略，指定向量化操作的Tile维度
# WHY: TileShape配置直接影响数据局部性和并行度，需根据硬件特征调优
pypto.set_vec_tile_shapes(64)

# 使用@pypto.frontend.jit装饰器定义Kernel函数
# WHY: 装饰器在首次调用时触发JIT编译，将Python AST转换为计算图
@pypto.frontend.jit
def vector_add(
    a: pypto.Tensor(shape, dtype),
    b: pypto.Tensor(shape, dtype),
    output: pypto.Tensor(shape, dtype)
):
    # Tensor级别的加法操作，直接使用Python运算符重载
    # WHY: 运算符重载使得计算表达式与数学写法高度一致，降低算法迁移成本
    result = a + b
    output[:] = result

2.4 符号化编程与动态Shape支持

PyPTO支持符号化标量（SymbolicScalar），用于支持动态Shape Tensor的表达和处理，使得框架可以在编译时进行Shape推断和优化。通过pypto.symbolic_scalar()接口，开发者可以在编译期创建一个符号化标量，其具体数值在运行时从Tensor的动态维度中获取。这一能力对于处理可变Batch Size等场景至关重要。以下代码演示了动态Shape的处理方式：

import pypto

def softmax_core(x: pypto.Tensor) -> pypto.Tensor:
    # 计算行最大值，用于数值稳定性优化
    row_max = pypto.amax(x, dim=-1, keepdim=True)
    # 值归一化：减去行最大值防止指数溢出
    sub = x - row_max
    # 指数运算
    exp = pypto.exp(sub)
    # 按行求和
    esum = pypto.sum(exp, dim=-1, keepdim=True)
    # 最终Softmax归一化
    return exp / esum

@pypto.frontend.jit
def dynamic_softmax(
    input_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),
    output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),
):
    bs, seqlen, head, dim = input_tensor.shape
    tile_b = 1
    b_loop = bs // tile_b
    # 配置向量TileShape，指定Tile在每个维度上的切分大小
    # WHY: 动态Shape场景下TileShape需与运行时维度对齐，否则编译器无法完成静态推导
    pypto.set_vec_tile_shapes(1, 4, 1, 64)
    # 使用pypto.loop定义循环控制流，替代Python原生for循环
    # WHY: pypto.loop在计算图中生成显式循环节点，使编译器能够进行循环级优化
    for idx in pypto.loop(0, b_loop, 1, name="LOOP_BATCH", idx_name="idx"):
        b_offset = idx * tile_b
        b_offset_end = (idx + 1) * tile_b
        input_view = input_tensor[b_offset:b_offset_end, :seqlen, :head, :dim]
        softmax_out = softmax_core(input_view)
        output_tensor[b_offset:, ...] = softmax_out

三、编译管线手把手走读

3.1 Python AST解析阶段

PyPTO编译管线的起点是Python AST（Abstract Syntax Tree，抽象语法树）解析。当开发者调用一个被@pypto.frontend.jit装饰的Kernel函数时，框架通过Python内置的ast模块将函数体的源代码解析为AST节点树。这一步骤的核心价值在于将动态的解释型Python代码转换为结构化的语法树表达，从而为后续的分析和转换提供基础。

在AST解析阶段，框架遍历语法树节点，识别以下关键元素：函数调用节点（对应PyPTO的Tensor Operation）、赋值节点（对应计算图的边和数据流）、索引/切片节点（对应Tensor的视图操作）、循环节点（对应pypto.loop控制流）以及条件分支节点（对应pypto.cond控制流）。通过这一遍历过程，框架将Python代码的语义完整地提取出来，构建出与硬件无关的Tensor Graph表示。

值得强调的是，AST解析阶段直接处理的是Python源码而非字节码，这是因为PyPTO需要保留变量名、注释等源码级别的信息用于计算图构建和调试。每个pypto.xxx函数调用在这一阶段被识别为一个Tensor Op节点，其输入输出张量作为节点的数据依赖连接起来。

3.2 TACO调度与Tile Graph生成

从Tensor Graph到Tile Graph的转换由TACO（Tiling And Coloring Optimization，切分与着色优化）模块负责。TACO模块是PyPTO编译管线中最具技术含量的环节之一，它根据开发者在set_vec_tile_shapes或set_cube_tile_shapes中指定的切分策略，将粗粒度的Tensor操作分解为细粒度的Tile操作。

在向量运算场景中，TACO根据指定的TileShape（如64），将一个长度为N的向量切分为N/64个Tile，每个Tile包含64个连续元素。这些Tile被映射到昇腾NPU的向量计算单元（AIV）上并行执行。在矩阵乘法场景中，TACO根据Cube TileShape配置，将矩阵切分为多个子块，每个子块加载到UB（Unified Buffer）中，然后调用Cube指令执行矩阵乘法。

TACO模块的调度决策包括：Tile到处理器核的映射关系、数据在不同内存层次（GM、UB、L1）之间的搬移时机、Tile之间的同步屏障插入位置、以及同构Tile的合并策略。这些决策直接影响算子在硬件上的执行效率，因此TACO模块通常需要结合硬件特征模型和运行时性能数据进行迭代优化。

3.3 PTO IR到CCE指令的编译链路

PTO IR（PTO Intermediate Representation）是PyPTO自定义的中间表示，设计目标是既保留足够的语义信息用于高级优化，又足够底层以便于最终代码生成。PTO IR以虚拟指令的形式存在，每条虚拟指令对应一个具体的计算模式（如矩阵乘法、向量运算、归约操作等），但不直接对应特定硬件的机器码。

PTO IR到CCE指令的编译过程由PyPTO的CodeGen模块完成。CodeGen遍历PTO IR中的每一条虚拟指令，根据目标昇腾NPU的硬件特性（如AIC向量计算能力、AIV张量计算能力、Cube矩阵乘法单元等）生成对应的CCE指令序列。这一过程需要处理指令调度（将指令分配到合适的执行单元）、寄存器分配（将虚拟寄存器映射到物理寄存器）以及指令重排（通过指令重排消除数据依赖并提高流水线利用率）。

最终生成的CCE指令被打包为二进制可执行文件，加载到昇腾NPU设备上后由调度执行层以MPMD方式分发到各个处理器核。

3.4 编译产物与工具链可视化

PyPTO在编译过程中会自动生成各个阶段的中间产物，开发者可以通过PyPTO Toolkit可视化工具查看这些计算图的结构。关键的中间产物文件包括：program.json，包含Execute Graph和Block Graph的汇总信息；merged_swimlane.json，包含运行时泳道图数据，展示每个AIC/AIV核上的任务执行顺序和耗时。

通过泳道图，开发者可以直观地观察到计算核之间的空闲间隔（即气泡）以及耗时较长的泳道条目，从而识别性能瓶颈并针对性地调整TileShape配置或算法逻辑。泳道条目的长度对应任务的耗时，能够直观地反映计算的密集程度。这一工具链设计使得性能调优过程从经验驱动转变为数据驱动。

四、算子注册与框架集成

4.1 PyTorch集成方式

PyPTO算子与PyTorch的集成通过框架自动处理Tensor类型转换实现。开发者在调用PyPTO Kernel时直接传入PyTorch Tensor，框架会在内部完成从PyTorch Tensor到PyPTO Tensor的隐式转换。这一设计使得存量PyTorch模型可以零改动地引入PyPTO自定义算子。以下代码展示了与PyTorch集成的完整流程：

import pypto
import torch

# 定义PyPTO Kernel，权重以Python变量形式在Kernel外部定义
# WHY: 将权重作为闭包变量引入，使框架能够识别数据依赖关系并正确处理数据传输
@pypto.frontend.jit
def my_operator(x: pypto.Tensor(in_shape, dtype), output: pypto.Tensor(out_shape, dtype)):
    result = pypto.matmul(x, weight)
    output[:] = result

# 创建PyTorch Tensor作为输入输出缓冲区
# WHY: PyPTO框架自动处理PyTorch Tensor与内部Tensor表示之间的格式转换和内存拷贝
input_torch = torch.randn(32, 128, device='npu')
output_torch = torch.zeros(32, 64, device='npu')

# 直接调用Kernel，传入PyTorch Tensor对象
my_operator(input_torch, output_torch)

4.2 多函数组合与复杂模式

PyPTO支持多函数组合模式，开发者可以将计算逻辑拆分为多个子函数，每个子函数独立构建计算图，最后在主Kernel中组合调用。这一设计使得复杂算子可以被模块化地组织，便于代码复用和分步调试。例如，一个Transformer Attention算子可以分解为QKV投影、点积注意力、Softmax归一化和输出投影等多个子函数，每个子函数独立开发测试后组合为完整算子。

4.3 大模型实现样例

PyPTO仓库中提供了多个大模型实现样例，展示了该框架在大模型开发场景中的实际应用能力。DeepSeekV3.2 SFA（稀疏Flash Attention量化实现）和GLM V4.5 Attention（GLM注意力机制实现）等样例说明了PyPTO处理复杂融合算子的能力边界。这些样例覆盖了MoE（Mixture of Experts）、Attention变体、稀疏计算等大模型中的典型算子模式，为开发者提供了可参考的实现范式。

五、开发效率对比

下表从多个关键维度对比CCE C++传统开发方式与PyPTO新型开发方式的生产效率差异。数据基于典型融合算子开发场景的实测经验，具体数值会因算子复杂度和开发者熟练程度而有所浮动。

维度	CCE C++方式	PyPTO方式	差异来源
代码行数（典型融合算子）	800~2000行	50~200行	PyPTO自动处理Tile切分和数据搬运，无需手工编码
编译周期	10~30分钟/次	1~5分钟/次	PyPTO编译链路简化，减少了手工调试环节
动态Shape支持	需手工适配	原生支持	PyPTO符号化标量和动态Shape接口直接可用
计算图可视化	需第三方工具	内置Toolkit	PyPTO Toolkit提供泳道图、控制流图等多视图
算法开发者学习曲线	高（需理解CCE架构）	低（Tensor级别API）	PyPTO以Tensor操作为入口，贴近算法直觉
算子融合优化	需手工识别	编译器自动融合	PyPTO编译层自动识别并融合可优化的算子组合
PyTorch集成成本	高（需写适配层）	低（自动类型转换）	PyPTO Kernel直接接收PyTorch Tensor作为输入
硬件指令调优空间	完全可控	受框架抽象限制	CCE C++可精细控制每条指令，PyPTO通过TileShape间接调优

从对比数据可以看出，PyPTO在开发效率、代码可维护性和动态Shape支持方面具有明显优势，而CCE C++在硬件指令调优的精细度上保持优势。对于算法验证和快速迭代阶段，PyPTO是更优选择；对于极致性能敏感的核心算子，CCE C++仍然是不可替代的手段。两者并非互斥关系，业界常见的做法是用PyPTO完成算法验证和快速迭代，再对关键瓶颈算子使用CCE C++进行深度手工调优。

六、完整实战：Softmax算子开发

本节通过一个完整的Softmax算子开发实例，演示从环境准备到编译运行的全部步骤。读者可以按照以下步骤在自己的环境中复现整个过程。

步骤一：环境准备

需要在昇腾NPU环境中部署PyPTO。根据官方文档，环境部署主要包括软件包获取和第三方依赖安装两个环节。确保昇腾驱动和CANN运行时环境（ascend-toolkit）已正确安装，并通过source /usr/local/Ascend/ascend-toolkit/set_env.sh命令配置环境变量。安装完成后，通过python3 -c "import pypto; print(pypto.__version__)"命令验证安装是否成功。

步骤二：编写Softmax算子代码

创建softmax.py文件，实现完整的Softmax算子。定义核心计算函数softmax_core，使用PyPTO提供的amax、exp、sum等Operation接口组合实现Softmax的数学表达式：

import pypto
import torch
import numpy as np
from numpy.testing import assert_allclose

# 核心计算函数，使用PyPTO Operation接口组合实现Softmax数学表达式
# WHY: 将核心计算封装为独立函数，使计算逻辑与Kernel调度逻辑解耦，便于测试复用
def softmax_core(x: pypto.Tensor) -> pypto.Tensor:
    row_max = pypto.amax(x, dim=-1, keepdim=True)  # 计算行最大值
    sub = x - row_max  # 值归一化，防止指数溢出
    exp = pypto.exp(sub)  # 指数运算
    esum = pypto.sum(exp, dim=-1, keepdim=True)  # 按行求和
    return exp / esum  # 概率归一化

# 使用@pypto.frontend.jit装饰器定义Kernel
# WHY: 装饰器使函数在首次调用时触发JIT编译，将计算图编译为可执行代码
@pypto.frontend.jit
def softmax_kernel(
    input_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),
    output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),
):
    bs, seqlen, head, dim = input_tensor.shape
    tile_b = 1
    b_loop = bs // tile_b
    # 配置向量TileShape，指定每个维度上的Tile切分大小
    # WHY: TileShape配置直接影响数据在UB和L1之间的搬移策略，需结合硬件特征设置
    pypto.set_vec_tile_shapes(1, 4, 1, 64)
    for idx in pypto.loop(0, b_loop, 1, name="LOOP_L0_bIdx", idx_name="idx"):
        b_offset = idx * tile_b
        b_offset_end = (idx + 1) * tile_b
        input_view = input_tensor[b_offset:b_offset_end, :seqlen, :head, :dim]
        softmax_out = softmax_core(input_view)
        output_tensor[b_offset:, ...] = softmax_out

步骤三：编写测试用例

实现test_softmax函数，使用PyTorch内置Softmax作为基准结果，验证PyPTO实现的数值正确性：

def test_softmax(device_id: int = None, run_mode: str = "npu") -> None:
    device = f'npu:{device_id}' if (run_mode == "npu" and device_id is not None) else 'cpu'
    shape = (32, 32, 1, 256)
    # 创建PyTorch Tensor作为输入输出缓冲区
    # WHY: PyPTO框架支持直接接收PyTorch Tensor，无需手动分配PyPTO Tensor
    x = torch.rand(shape, dtype=torch.float, device=device)
    y = torch.zeros(shape, dtype=torch.float, device=device)
    # 调用PyPTO Kernel执行计算
    softmax_kernel(x, y)
    # 使用PyTorch内置Softmax生成基准结果进行对比
    golden = torch.softmax(x, dim=-1).cpu()
    y = y.cpu()
    # 计算最大绝对误差
    max_diff = np.abs(y.numpy() - golden.numpy()).max()
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {y.shape}")
    print(f"Max difference: {max_diff:.6f}")
    if run_mode == "npu":
        # 允许一定容差（CANN浮点运算的特性）
        assert_allclose(np.array(y), np.array(golden), rtol=3e-3, atol=3e-3)
    print("Softmax test passed")

if __name__ == "__main__":
    test_softmax(run_mode="npu")

步骤四：执行与验证

配置环境变量并执行脚本。设置设备ID后运行python3 softmax.py，程序将在${work_path}/output/output_*/目录下生成编译和执行结果文件，包括计算图（program.json）和泳道图数据（merged_swimlane.json）。开发者可以通过PyPTO Toolkit打开program.json文件逐层钻取查看Execute Graph和Block Graph结构，通过泳道图分析各处理器核的任务执行情况。

结尾

PyPTO作为CANN生态中面向昇腾NPU的Python编程框架，通过PTO编程范式和四层架构设计，在保持高性能计算潜力的前提下大幅压缩了算子开发的技术门槛。从Python AST解析到PTO IR生成、从TACO Tile切分调度到CCE指令编译，PyPTO的编译管线每一级都有明确的语义边界和模块化设计，使得不同角色的开发者可以在各自擅长的抽象层次上工作。算法开发者使用Tensor级别的API快速实现算法验证，性能专家通过TileShape调优追求极致性能，系统开发者则在PTO虚拟指令层次上进行框架集成和工具链开发。

在架构设计层面，PyPTO的四层架构体现了分层抽象的核心原则：用户API层屏蔽了硬件细节，计算图编译层通过多级Pass实现渐进式优化，代码生成层将高层计算描述翻译为底层指令，调度执行层则以MPMD模型实现高效的核间并行。这种分层设计的直接收益是编译管线的每一级都可以独立演进而不影响其他层——当昇腾NPU引入新的硬件特性时，只需在代码生成层增加对新指令的支持，而无需修改上层的计算图构建逻辑。

---

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