推出新型面向 AI 加速器的高性能编程框架——PyPTO

1.简介

PyPTO（发音：pai p-t-o）是 CANN 推出的一款面向 AI 加速器的高性能编程框架，旨在简化算子开发流程，同时保持高性能计算能力。该框架采用创新的 PTO（Parallel Tensor/Tile Operation）编程范式，以基于 Tile 的编程模型 为核心设计理念，通过多层次的中间表示（IR）系统，将用户通过 API 构建的 AI 模型应用从高层次的 Tensor 图逐步编译成硬件指令，最终生成可在目标平台上高效执行的可执行代码

2.PyPTO 产生的背景

高性能算子开发一直是 AI 加速器编程的核心挑战。算子开发人员不仅需要理解算子的数学计算属性，还必须考虑如何将其转换为对硬件友好的执行方式。这种双重知识要求使得算子开发成为一个高度专业化的领域。

尽管存在诸多挑战，行业在技术趋势上已经形成了一些重要共识：

（1）编程模型的重要性：编程模型决定学习曲线，学习曲线决定生态上限。一个友好的编程模型能够显著降低开发门槛，扩大开发者群体，从而推动生态发展。

（2）分层暴露复杂性：复杂性不可避免，但必须"分层暴露、按角色定价"。绝大多数用户应该只看到高层抽象，而底层复杂性应该只暴露给需要它的专家用户。

（3）指令式编译+运行时优化：这是快速演化时代的主流范式。支持 JIT 编译和动态图执行，通过 Runtime API 支持从 eager 执行调试到图执行，多种编译路径共存，适应不同场景需求。

（4）工具链和工程体验：这是决定"生态成色"的第二战场。提供完整的可视化工具和分析工具，一致的错误信息和诊断机制，高质量示例和文档，与主流 IDE 和构建系统的集成，这些都是生态成功的关键因素。

（5）基于 Tile 的编程模型成为趋势：无论是业界主流的 GPU 厂商还是开源社区，近期都推出了基于 Tile 的编程框架。这验证了 Tile 编程模型的方向正确性。然而，现有基于 Tile 的框架仍然相对底层，需要开发者具备丰富的底层开发经验和硬件系统知识。

（6）通用计算能力扩展：AI 专用只是入口而非终点，需要扩展对通用并行计算的支持，包括数值计算、科学计算等，以及 3D 高斯溅射等渲染应用，为硬件找到足够多的"落地场景"和"主战场"。通用计算能力是生态的"放大器"，能够吸引更广泛的开发者群体和应用场景。

面对上述挑战和局限性，PyPTO 应运而生。

3.PyPTO 的设计目的与创新点

本节讨论的是 PyPTO 的设计初衷、试图解决的核心问题，以及与现有方案的区别PyPTO 的设计初衷如下：

（1）降低开发门槛：让算法开发者能够直接实现高性能算子，无需深入了解硬件细节。通过 Tensor 级别抽象，算法开发者可以专注于算法逻辑的表达，而将硬件优化交给框架处理。

（2）提升开发效率：通过 Tensor 级别抽象，快速获得性能良好的实现。这种设计显著缩短了从算法设计到可运行代码的开发周期，支持快速迭代和验证。

（3）保留优化空间：通过多层级 IR 和 Pass 系统，支持从快速可用到极致性能的平滑过渡。开发者可以根据实际需求选择合适的优化深度，既保证了开发效率，又保留了追求极致性能的可能性。

（4）支持复杂场景：强调复杂融合算子，甚至于整个模型网络的开发能力。框架支持分布式执行、自动微分等高级特性，使其能够适应从单算子优化到大规模模型训练的多样化场景。

（5）作为补充而非替换：PyPTO 作为已有编程框架的补充而非替换，为开发者提供了另一种选择，特别适合需要深度性能优化和细粒度控制的场景。

PyPTO 致力于解决以下核心问题：

（1）消除算法与算子开发的鸿沟：通过 Tensor 级别抽象，算法开发者可以直接实现高性能算子，无需等待专门的算子开发团队。这消除了传统分工模式带来的沟通成本和迭代延迟。

（2）平衡编程简单性和控制力：通过分层抽象设计，对不同开发者暴露不同抽象层次。算法开发者可以使用高层次的 Tensor API，而性能优化专家和系统开发者可以根据需要深入到 Tile、Block 甚至虚拟指令集层次。

（3）支持快速迭代和深度优化：提供"快速可用 → 灵活调优 → 深度优化"的平滑过渡。大多数开发者可以快速获得可用的实现，而性能敏感的开发者可以通过工具链进行深度优化，追求极致性能。

（4）提供完整的工具链支持：可视化、调试、性能分析一体化支持，帮助开发者快速定位问题、理解性能瓶颈，并实现深度定制。

PyPTO 与现有方案相比，具有以下显著区别：

（1）多层级 IR 系统：支持从 Tensor 到硬件指令的渐进式降低，每个层次都可以进行针对性优化。这种设计使得每个层次都可以进行针对性的优化，同时支持不同抽象层次的编程接口。

（2）基于 Tile 的编程模型：所有计算基于 Tile 进行，充分利用硬件并行性和内存层次结构。这种设计在保持易用性的同时实现了高性能。

（3）更高层次的抽象：PyPTO 提供更高层次的 Tensor 抽象，算法开发者无需深入了解底层硬件细节即可开发高性能 kernel。这显著降低了开发门槛，使得更多开发者能够参与高性能算子开发。

（4）MPMD 执行调度：避免全局同步，灵活调度，更好的资源利用。这种设计适配多核架构，能够更好地利用硬件资源，提升整体利用率。

（5）Man-In-The-Loop 工作流：人工参与的优化流程，在开箱即用与性能极致之间提供平滑过渡。通过 IDE 集成、可视化、性能分析、调试支持等工具，帮助开发者理解编译过程、识别性能瓶颈，并进行深度定制。

（6）分层抽象设计：对不同开发者暴露不同抽象层次，平衡编程简单性和控制力。算法开发者使用 Tensor 层次，快速实现功能；性能专家使用 Tile 层次，进行深度优化；系统开发者使用 Block 层次，实现框架集成。

PyPTO 通过多层级 IR 系统、基于 Tile 的编程模型、更高层次的 Tensor 抽象、MPMD 执行调度、Man-In-The-Loop 工作流和分层抽象设计，致力于消除算法与算子开发的鸿沟，平衡编程简单性和控制力，为不同层次的开发者提供相应的抽象层次和工具支持。

PyPTO 的分层抽象

PyPTO 通过多层次抽象设计，为不同的用户群体提供不同层次的控制粒度。基于分层抽象的能力，框架既能够为大多数开发者提供易用的高层抽象，又能够为专家用户提供细粒度的控制能力。

非常欢迎开发者朋友们在公众号评论区提出你们关心的问题，我们会将这些问题收集起来，在专场直播中为大家进行解答。

关注昇腾CANN公众号，及时获取直播时间与观看链接，相信通过这次直播，大家会对CANN PyPTO有更深入的了解，为昇腾芯片的应用开发带来更多的创新和突破！