HyperAIHyperAI

Command Palette

Search for a command to run...

毕昇编译器核心 AscendNPU IR 开源,华为架构师海丽娟详解技术进展,支持 Triton 等算子生态

Featured Image

8 月 1 日,由 HyperAI 主办的 Meet AI Compiler 技术沙龙第 9 期如约而至。来自智源研究院、 TileRT 团队、腾讯、华为、智元创新等多家头部企业与科研机构的分享嘉宾,围绕 AI 编译器在语言表达、算子计算、推理执行、场景落地等多维度的协同演进展开深度研讨,聚焦技术痛点、实践方案和行业趋势,为 AI 编译技术的迭代路径与产业价值的深度挖掘贡献多元思路与实践经验。

现场,华为 AscendNPU IR 架构师海丽娟以「AscendNPU IR:编译底座开源开放,支持多语言接入昇腾」为主题进行分享。她系统介绍了昇腾毕昇编译器组件 AscendNPU IR 的整体技术架构与设计理念,并围绕 AscendNPU IR 面向昇腾 950 扩展的新特性、技术迭代优势、生态社区建设等内容展开细致解读,拆解开源底座的核心能力和技术优势,展示其对多类前端语言的适配能力,向开发者呈现出 AscendNPU IR 易用友好、开放共建的开发体验。

海丽娟老师出席 Meet AI Compiler 第 9 期

HyperAI 在不违原意的前提下,对演讲分享进行了整理汇总。

关注微信公众号「HyperAI 超神经」,后台回复关键字「0801 AI 编译器」,即可获取确认授权的讲师演讲 PPT 。

为什么 Triton 编程高效易用?

Triton 作为一门编程语言,在具备高效性的同时还兼顾友好的编程体验,那这份开发优势究竟从何而来?

简单来说,Triton 使用了类 Python 语法,提供 Block Level Tile 级编程,向开发者屏蔽抽象核内硬件内存、指令、流水等底层细节操作,从而降低算子开发门槛,提升开发效率。

Triton 所提供的 Tile 编程模型,相比如 C++ 等异构编程语言来说,可以让开发者更专注于对算子做数据切分,聚焦 Tensor  op 抽象、逻辑表达。基于 Tile 级编程抽象,除核间数据调度由用户管理之外,核内访存合并、 Shared-Memory 管理都可由编译器自动完成。

更重要的是,Triton 方言和 Triton 官方 GPU 编译器,都是基于 MLIR 构建的多级 IR 抽象,面对不同硬件深度优化。

AscendNPU IR 持续进行架构优化

AscendNPU IR 是基于 MLIR 构建的昇腾硬件 Tile 级抽象,服务于底层编译优化。下图可见其架构,AscendNPU IR 向下对接 LLVM IR,最后编译到昇腾硬件二进制。上层 Triton 等三方语言和编译器都可以接入 AscendNPU IR,从而实现昇腾的适配。

AscendNPU IR 关键技术特征,其一是面向昇腾硬件做了自下而上的抽象,抽象核内不同资源,如内存、数据搬运单元等;其二是提供抽象 Tile 级 OP,可对底层硬件实现跨架构统一支持,覆盖从 A2/A3 到 昇腾 950 SIMD 、 SIMT 最新硬件架构。基于 MLIR 开放性和可扩展性,AscendNPU IR 支持更多上层语言接入。

下图是 AscendNPU IR 编译架构设计,从昇腾 A2/A3,再到今年最新上市的 昇腾 950 硬件。左侧是去年发布的第一版 AscendNPU IR 架构,主要由两层构成。上层是硬件无关层 —— HFusion,主要实现上层多维度融合优化。它的定位基于社区 Linalg 和扩展 OP,向上对接各类标准化方言,完成数据类型和 OP 规范化预处理,以及 Auto Schedule 等融合优化。

硬件相关层 HIVM 主要做昇腾硬件的高层抽象,依次完成核映射编译、片上内存映射和处理单元映射。

首先是核映射编译,包括 Cube 核、 Vector 核,高层 Tile 表达通过核映射,Cube 的归 Cube 核,Vector 的归 Vector 核。它们之间的数据通信、数据同步以及中间 Workspace 内存管理,都由编译器自动补齐,同时完成 CV 流水并行高阶优化工作。

接下来是片上内存映射。 Cube 核和 Vector 核都有各自的片上内存。 Tensor 要在昇腾硬件上工作,需要把那些逻辑 Tensor 数据推导到对应的物理内存上,还有 Cube 特有的分型矩阵格式推导,片上内存分配,都由编译器自动完成。

最底层是处理单元映射。流水同步保证时序和并行效率。向量化和张量化使能高性能硬件指令完成最终编译优化工作。

图中右侧可以看到,从去年 A2/A3 到今年 昇腾 950 架构发生的一些主要变化。上层硬件无关层 HFusion 基本保持稳定,主要扩展来自于硬件相关层 HIVM ,从 Membased SIMD 扩展到支持 Regbased SIMD 和 SIMT 。 CV 融合编译过去通过 Global Memory 交互 Cube 核和 Vector 核数据,昇腾 950 实现紧耦合方式做更高效的数据交互。

总体来说新架构依然保持稳定两层分层,但随着架构变化,Tile 编译已经进行较大幅度的升级。

先来看下 Vector Regbased SIMD 架构演进。 MemBased SIMD 时 Vector 计算数据基于片上 Unified Buffer 进行向量化计算,指令基于片上内存直接完成。昇腾 950 上,数据基于高速寄存器进行计算,多了一层寄存器层级。对于后端尤其 Vector 优化而言,编译工作有如下变化:

一,面向 Regbased 的向量化。 Tensor 操作通过向量化切分映射到固定向量位宽的寄存器指令。从片上内存把数据加载到寄存器,基于寄存器完成计算再存储回片上内存,这样 Load-Compute-Store 形成一个循环(loop)。

二,寄存器粒度融合。 OP 计算每次进出片上内存效率不理想。基于寄存器做深度融合,让数据常驻寄存器。把多个 Triton op 融合到一个循环 loop 从而减少 Load-Store,这就涉及如何利用不同搜索算法,以及 Cost Model 去做更好的融合策略。

三,AVE 方言抽象。  昇腾亲和的 Vector 方言扩展,屏蔽底层硬件掩码差异,完成掩码分析与推导。硬件 Vector Load 有丰富的随路优化能力。使能随路操作可以把多条指令合并成一条。这些都是针对 Regbase SIMD 方言上做的优化设计。

昇腾 950 硬件新增支持 SIMT 单元,如下图绿色部分所示。它可以使 Vector 核有三种不同的计算模式,纯 SIMD 的方式,纯 SIMT 的方式,还有 SIMT 和 SIMD 混合交替执行的方式。 SIMT 可以支持离散访存场景加速,而稠密计算可考虑通过 SIMD 完成更激进的向量优化。

编译流程在进入 Vector SIMD 或 SIMT 编译前先做融合分析优化。识别哪些适合 SIMT ,提取出来分别做 SIMT 和 SIMD 编译,最后再合到一起做 Vector 整体计算,过程中会对 SIMT 和 SIMD 进行建模分析。

跟 Attention 类算子密切相关的是 CV 核的升级变化。 A2/A3 上 Cube 和 Vector 没有片上通路,只能通过 Global Memory 片外显存来进行交互。而昇腾 950 ,如下图红色所示,Cube 结果在 L0C,可以直接拷贝到 Vector 片上内存。反过来,Vector 计算结果也可以直接拷贝到 Cube 片上内存接着执行 Cube 计算。总而言之,CV 紧耦合快速数据交换对于 Attention 性能优化有非常大的帮助。

CV 算子优化关键技术之一是 CV 流水并行。昇腾自动化 CV Pipeline 可以更好地隐藏时延,充分利用 Cube 核和 Vector 核算力资源。

另一个关键优化是 AutoSubTiling 。昇腾 Cube 和 Vector 核配比 1:2,两个 Vector 核对应一个 Cube 核。编译器完成 Vector 操作 1:2 切块,从而让两个 Vector 核分别同时计算其中一半数据,更快完成计算。

最新版本编译器功能泛化工作

针对刚发布的 AscendNPU IR 最新版本,逐一介绍其中重要的功能泛化工作。

第一:CV 数据交互。 Cube 和 Vector 分核涉及的核间数据交互由编译器自动补齐。 Triton 算子 Cube 和 Vector 计算可能分布在控制流不同分支,给编译分析工作带来一定复杂性。不能直接通过 Pattern Match 进行简单的指令插入。如上图例子中,两个分支分别做 Cube 和 Vector 计算。 CV 数据交互需要支持复杂控制流推导以保证功能完备性。

新版本增强了 InsertCVLoadStore 重要 Pass 分析流程,通过全局的跨复杂控制流的推导补齐 CV 数据交互。首先插入确定性锚点,矩阵乘 op 输入在 L1,输出在 L0C,向量 op 输入输出都在 UB 。之后基于起始锚点插入强制类型转换操作,进行跨控制流的向上向下的推导和传播。传播结束,所有操作数 Tensor 都明确所在的内存层级。

如上面例子所展示,Tensor 原本不包含内存层级语义,插入一个强制类型转换 Op 之后,将其限定在 Vector 核 UB 内存;下方是 Cube 核 L1 内存。两者在推导过程中遇到冲突 —— 内存层级不一致,需要插入 Load-Store 或拷贝语句打通数据。不同代际处理方式略微不同。 A2/A3 Cube 和 Vector 需要插入  Global Memory Load-Store 打通数据,而昇腾 950 通过片上数据拷贝打通数据。

第二:MultiBuffer 针对复杂控制流的增强。 MultiBuffer 是做流水并行的重要功能前提。那么 Tensor 如何能变成 MultiBuffer 呢?

此前方案仅支持 for 循环实现。循环遍历时,MultiBuffer 复用 for 循环的迭代变量,根据这一迭代变量轮转 Buffer Slot 。这个方案限制是不能支持其他 While 和嵌套控制流。于是引入一个独立 Buffer 计数器。计数器跟着复杂控制流进行分支轮转,从而准确跟踪 Slot 变化。

除此之外,alloc 操作需要在 Cube 和 Vector 公共循环层可见,方便 Cube/Vector 两边同时对内存分配操作进行变换,比如分配 2 个/3 个 Buffer,保证 MultiBuffer 变换一致性。

第三:与昇腾算子优化密切相关的多核变换 —— AutoBlockify 。平时写 CUDA 或 Triton 算子,基本上会切很多逻辑 block,通过硬件调度执行。昇腾仅靠重复轮次调度开销较大。把不同逻辑核归并到一个 for 循环,交给软件循环调度更好优化性能。同时 for 循环中还可以应用 MultiBuffer 和流水并行优化等高阶优化技巧,使开箱性能更优。

第四:CV Pipeline 。如 FlashAttention 算子有四段任务分布在 Cube 和 Vector 两个核内分别执行。任务间串行执行,Cube 和 Vector 计算之间需要相互等待。编译流程中可以通过 MultiBuffer 使 Cube 和 Vector 尽量前提执行,达到流水并行的加速效果。

首先把 C/V 边界点找出来,再将代码段准确划分出 Cube Scope 和 Vector Scope,然后根据调度算法应用 MultiBuffer 实现流水并行变换,包括多种策略模式:Unroll ,Skew 及动态。对于不同计算负载,流水并行可以选择最优模式。

第五:CV 1:2 切分。两个 Vector 核和一个 Cube 核之间协作处理,对于原始 Tensor 数据需要尽可能准确切半,完成高效 Vector 计算。理想状态是从高维找到并行轴做全局对半切分,而候选轴可能历经 Transpose 、广播、归约等计算操作,增加 CV 1:2 切分的复杂度。

首先做全局维度分析,定位更适合做切分的那根轴。确定切分轴后,从末尾 Store 操作开始插入一个切分 OP 起始标记,再从这个尾部逐 OP 往根节点冒泡,最终实现完美切分。期间如遇到一些复杂场景不可切,退回到保守 1:1 未切分状态。

另外某些复杂场景下仅切分并行轴性能还不够,比如 FB8 低精度 Attention 算法 Vector 压力大,需要对 reduce 轴进行切分归约。由于两个 Vector 核不能直接交换数据,需要经过 Global Memory 将中间 Reduce 结果交换,进行二次规约,完成整体规约切分。此外动态 Shape 等场景也需要增强 CV 1:2 切分手段。

Triton-Ascend 、 AscendNPU IR 已开源共建

当前 AscendNPU IR 已开源到昇腾社区,大家可以扫图中二维码关注开源项目。在社区活动日历中已发布了 AscendNPU IR 社区 SIG 双周例会,开发者们可以提前关注每周研讨议题及会议纪要。

昇腾社区开放了开源实习项目和社区任务,为广大开发者提供了参与生态共建的方式,大家可以在这里看到所有最新的任务状态,并完成认领。这些社区任务难度不一,请大家根据自己的实际情况选择和参与。

开源实习和社区任务在 AscendNPU IR 和 Triton-Ascend 开源仓主页有快速链接,其中包含任务领取流程,每人同时只能领取一个任务。社区任务会及时刷新,同时开发过程也将全程保持互动开放。每个任务有对应的丰厚奖励,欢迎大家自由选取。

对于没有昇腾开发环境的参与者,社区提供了免费的昇腾算力平台 HiDevLab(https://hidevlab.huawei.com/home),开发者可自行注册并申请算力环境,申请成功后默认派发 100 小时免费时长。领取了社区任务的开发者们都可以申请算力来完成任务的开发与验证。

感谢大家关注!