前言

在昇腾NPU的算子开发现实里,不同芯片代际之间的指令差异是开发者绕不开的麻烦事。Ascend 910系列用了达芬奇架构的一套流水线,后续Ascend 950又引入了新的计算单元和存储层次,直接拿底层指令写的算子换一代芯片就得重写一遍。CANN作为昇腾异构计算架构,在这个问题上给出了一套方案——PTO(Parallel Tile Operation)虚拟指令集架构。pto-isa仓库就是这套虚拟ISA的开源实现,它把tile级的计算与数据搬运抽象成90多条标准指令,让算子代码在Ascend A2、A3、A5之间具备可移植性,同时不丢掉对底层性能的控制力。这个仓库在2025年底开源,到2026年已经覆盖了计算、通信、合轴、MX、卷积、量化等多类指令,并且集成了PyPTO和TileLang Ascend两个前端框架。

tile编程模型的技术动机

昇腾NPU的计算单元结构决定了算子开发者必须以tile为单位组织计算。Cube矩阵计算单元每次处理一个固定大小的子矩阵块,Vector单元每次处理一个向量段,MTE数据搬运引擎每次搬运一块连续或步长的数据。这三个单元各有自己的吞吐上限和延迟特征,算子要跑满性能,就必须让它们各自忙起来、互相不空等。传统做法是直接在Ascend C里操控各个硬件单元的API,手动编排每个单元的启动顺序和同步关系。这种写法性能确实可控,但有一个硬伤:换了芯片代际,Cube单元的tile大小变了,MTE的搬运粒度变了,Vector单元的通道数变了,原来的编排逻辑可能完全不适用。pto-isa解决的就是这个矛盾——它定义了一套与具体硬件无关的tile操作语义,开发者用这套标准语义描述算子的计算流程,由PTO库负责把标准操作映射到具体芯片的底层实现上。tile的shape在编译期确定,运行时不变,这个设计约束保证了流水线的可预测性。编译期确定shape意味着所有的buffer分配和指令调度都可以在编译阶段完成,不需要运行时的动态判断开销。这对昇腾NPU这种资源受限的硬件环境很重要——运行时的动态分配不仅引入额外的指令开销,还可能导致buffer碎片化和不可预测的延迟。固定shape的设计要求开发者在写算子时预先确定处理数据的分块策略,根据矩阵的维度特征选择合适的tile大小和形状,再在编译阶段把这些参数固化到算子代码中。这种设计带来的另一个好处是让性能分析变得可预测——同一个算子处理同样shape的输入,每次执行的指令序列完全相同,性能表现也高度一致,便于性能回归测试和瓶颈定位。

PTO指令体系结构

pto-isa定义的指令按功能可以分为几个大类。计算类指令覆盖了矩阵乘法(TMM)、向量加法(TADD)、向量乘法(TMUL)、激活函数等常见操作。数据搬运类指令控制数据从Global Memory到Unified Buffer(UB)、从UB到寄存器、从UB到UB等各级存储之间的移动,支持连续搬运和带步长的非连续搬运。合轴类指令用于矩阵转置、reshape等布局变换操作。事件与同步类指令通过set flag和wait flag机制控制指令间的依赖关系,确保数据就绪后才启动下游计算。MX指令扩展了混合精度计算的支持,卷积类指令封装了im2col和winograd等卷积加速路径,量化类指令处理定点和浮点之间的转换。

指令的参数设计遵循tile级抽象的一致性原则。每条计算指令的输入和输出都是tile对象,tile对象的属性包括数据指针、shape、stride和mask。shape在编译期确定后不可变,mask用于标记tile中哪些位置是有效数据——这在处理非整除的矩阵维度时特别重要,比如矩阵行数不是tile大小的整数倍,末尾一行tile就只有部分有效数据。stride参数描述了数据在内存中的实际排列方式,让同一条指令可以处理连续存储和非连续存储的数据,避免了为不同存储布局写不同版本算子的麻烦。这种参数设计让指令的语义保持简洁:一条TADD就是两个tile相加写回第三个tile,开发者不需要关心底层存储的物理细节。

通信类指令是2026年新增的能力,支持核间点对点传输(TGET/TGET_ASYNC)、信号同步和集合通信,可以让计算和数据搬运在同一个kernel内深度融合。这套指令的设计思路是让每条指令对应一个完整的tile级操作,而不是把底层硬件微操作暴露给开发者。比如一条TMM指令封装了从UB读取矩阵块、送入Cube计算、写回结果到UB的全过程,开发者不需要关心Cube的流水线级数或寄存器分配细节。这种封装降低了认知负担,同时保留了tile size、tile shape、指令排序等性能敏感参数的控制权。每条指令的实现都经过了严格的CPU Simulator测试和NPU硬件验证,仓库中的测试套件覆盖了多种tile shape组合和边界条件,确保指令在各种使用场景下的正确性。

Auto与Manual双路径开发

pto-isa提供了两条开发路径。Auto模式下开发者只需要描述计算逻辑,PTO库自动分配tile buffer并插入同步指令,适合快速验证算子的正确性。Manual模式则把所有控制权交给开发者,包括buffer分配、指令顺序、事件同步等,适合做极限性能优化。实际开发中一个常见的路径是用Auto模式先把算子逻辑跑通,确认输出结果正确,再逐步把关键路径切换到Manual模式做性能优化调整。CPU Simulator是这两条路径的基础设施,它用纯软件模拟PTO指令的执行过程,不需要昇腾NPU硬件就能验证算子逻辑。这意味着开发者可以在普通的x86或ARM服务器上完成大部分功能开发工作,等逻辑稳定后再移植到NPU上采集性能数据。这套CPU Simulator支持C++20编译,通过CMake构建,可以单独跑GEMM demo、Flash Attention demo等典型算子的功能测试。

下面这个示例展示了Auto模式下用PTO指令实现一个简单的向量加法算子:

from pto import Tile, TADD

def auto_add(x, y):
    tx = Tile(x, shape=(64, 64))
    ty = Tile(y, shape=(64, 64))
    tr = Tile.empty_like(tx)
    TADD(tx, ty, tr, mode="auto")
    return tr.data

Manual模式下同一个加法算子需要开发者显式管理buffer和同步。下面的代码展示了Manual模式下的控制粒度:

from pto import Tile, TADD, SetFlag, WaitFlag

def manual_add(x, y):
    tx = Tile(x, shape=(64, 64), buf_id=0)
    ty = Tile(y, shape=(64, 64), buf_id=1)
    tr = Tile.empty_like(tx, buf_id=2)
    SetFlag(0)
    WaitFlag(0)
    TADD(tx, ty, tr, mode="manual")
    return tr.data

事件同步与流水线编排

PTO的事件同步机制是理解算子性能的关键。在昇腾NPU上,Cube、MTE、Vector三个单元可以并行工作,但它们之间有数据依赖——MTE把数据搬到UB之后,Cube才能开始计算;Cube算完之后,Vector才能做后续处理。PTO用set flag和wait flag指令来表达这种依赖关系。set flag标记某个阶段完成,wait flag阻塞等待某个flag被设置。这种机制比简单的"等待完成"更灵活,因为它允许开发者构建多级流水线——比如在Cube处理当前batch的同时,MTE已经把下一个batch的数据搬到另一个UB区域了。这就是所谓的double buffering或者triple buffering技术,在PTO里通过多组flag和多个buffer ID的组合来实现。事件同步的粒度直接影响流水线的利用率。同步点太多,各个单元会频繁空等;同步点太少,可能出现数据竞争导致计算结果错误。PTO指令集没有强制规定同步点的位置和数量,而是把这个决策留给开发者,让开发者根据具体的算子特征和硬件参数来决定最佳的同步策略。这种设计保留了足够的灵活性,但也提高了开发门槛。

下面是一个用事件同步实现双缓冲流水线的示例:

from pto import TMM, SetFlag, WaitFlag, Tile

def pipelined_gemm(a, b, batch_size=8):
    out = []
    for i in range(0, a.shape[0], batch_size):
        ta = Tile(a[i:i+batch_size], buf_id=0)
        tb = Tile(b, buf_id=1)
        to = Tile.empty_like(ta, buf_id=2)
        SetFlag(0)          # mark load complete for current batch
        WaitFlag(0)         # ensure data ready before compute
        TMM(ta, tb, to)
        out.append(to.data)
    return out

通信扩展指令集

2026年3月发布的通信扩展是pto-isa的一个重要里程碑。在分布式训练场景下,多个NPU之间需要频繁交换数据——AllReduce、AllGather、ReduceScatter等集合通信操作占据了大量的时间和带宽。传统做法是算子kernel计算完成后,调用HCCL库执行通信,通信和计算是串行的两个阶段。PTO的通信指令让开发者可以在算子kernel内部直接发起数据传输,把通信和计算重叠起来。

TGET指令执行同步的远程读操作,数据经过UB中转,适合小批量高延迟场景。TGET_ASYNC指令走异步通道,通过DMA引擎直接把数据从远端NPU搬到本地,不经过UB,适合大批量低延迟场景。这两条指令的区别在于同步开销和数据路径——同步模式需要等待远端数据完全到达UB后才能继续执行后续指令,异步模式则发出请求后立即返回,通过单独的等待指令在需要时检查数据是否就绪。在带宽测试用例中,TGET_ASYNC在大块数据传输场景下表现明显优于TGET,因为DMA引擎可以独立于计算单元持续搬运数据,不占用UB带宽。信号同步指令允许一个核通知另一个核某块数据已经准备好了,集合通信指令封装了常见的通信模式。

仓库中有一个GEMM AllReduce通算融合的示例,在同一个算子流水线里完成了矩阵乘法和AllReduce操作,让计算和通信在时间上重叠执行。这种深度融合的设计在Ascend 910B2上的Flash Attention测试中也得到了验证——短序列场景下PTO实现相对torch_npu有超过两倍的加速比,长序列场景下加速比仍然超过百分之二十。通信指令的引入让算子开发者不再需要把算子和通信当作两个独立的模块来处理,而是可以在一个统一的编程模型里同时描述计算和数据交换的逻辑,大大简化了通算融合算子的开发复杂度。

跨代可移植性的实现机制

PTO的跨代可移植性建立在两个设计决策之上。在固定tile shape的前提下,同一段PTO算子代码在不同芯片代际上产生正确的结果。开发者不需要为Ascend 910和Ascend 950分别维护两套代码,只需要编译时指定目标平台,PTO库会自动把标准指令映射到对应平台的底层实现。这种映射不是简单的指令翻译——不同平台的Cube单元可能对tile shape有不同的最优值,MTE的搬运粒度也不同,PTO库需要根据目标平台的硬件参数调整指令的内部实现,同时保持外部的tile级语义不变。在性能方面,同一份PTO代码在不同平台上通常都能获得合理的性能表现,因为PTO库已经针对每个平台做了基本的映射优化。但如果要逼近每个平台的性能极限,仍然需要手动调整tile size、tile shape、指令顺序等参数。仓库的GEMM性能示例和Flash Attention示例都展示了如何在具体平台上做这种针对性优化。性能分析工具msprof可以用来定位算子的瓶颈类型——CUBE Bound表示计算单元是瓶颈,MTE Bound表示数据搬运是瓶颈,Vector Bound表示向量处理是瓶颈。根据瓶颈类型选择不同的优化策略,是PTO算子优化调整的基本思路。

性能对比分析

下面的表格概括了PTO在不同维度上与传统开发方式的效率差异:

维度 使用前 使用后 差异来源
跨代算子迁移 每换一代芯片需重写底层指令编排逻辑 同一份PTO代码编译到不同平台,库自动完成映射 PTO虚拟ISA消除了平台特定指令的直接暴露
算子开发调试周期 必须在NPU硬件或模拟器上迭代验证 CPU Simulator支持在普通服务器上功能验证 CPU软仿真降低了硬件依赖门槛
通信与计算重叠 通信和计算串行执行,算子完成后调用HCCL 通信指令嵌入kernel内部,可构建通算融合流水线 通信扩展指令允许在算子内直接发起数据传输
流水线同步控制 需要直接操作硬件同步原语,平台差异大 set/wait flag抽象统一,开发体验一致 标准事件模型屏蔽了底层同步机制的差异

pto-isa仓库是昇腾CANN生态中面向tile级算子编程的关键基础设施。它用虚拟指令集的方式统一了不同昇腾芯片代际之间的tile操作语义,让算子代码具备可移植性的同时保留了性能优化调整的空间。从计算类指令到通信类指令,从Auto模式到Manual模式,从CPU Simulator到CostModel,这套工具链覆盖了算子开发从功能验证到性能极限的完整路径。集合通信、系统调度、微指令等能力的持续扩展,让PTO ISA正在成为昇腾平台上算子开发的通用中间表达层。


仓库地址:https://atomgit.com/cann/pto-isa

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐