Bi Sheng's Compiler Core, AscendNPU IR, Has Been open-sourced. Huawei Architect Hai Lijuan Explains the Technical Progress and Highlights Its Support for Operator Ecosystems Such As Triton.

August 1,The 9th Meet AI Compiler technical salon, hosted by HyperAI, has arrived as scheduled.Guest speakers from leading companies and research institutions such as the Beijing Academy of Artificial Intelligence, the TileRT team, Tencent, Huawei, and Zhiyuan Innovation conducted in-depth discussions on the collaborative evolution of AI compilers across multiple dimensions, including language expression, operator computation, inference execution, and scenario implementation. They focused on technical pain points, practical solutions, and industry trends, contributing diverse ideas and practical experience to the iterative path of AI compiler technology and the in-depth exploration of its industrial value.

on site,Hai Lijuan, architect of Huawei AscendNPU IR, shared her insights on the topic of "AscendNPU IR: Open Source Compilation Base, Supporting Multi-Language Access to Ascend".She systematically introduced the overall technical architecture and design philosophy of the AscendNPU IR compiler component, and provided a detailed explanation of AscendNPU IR's new features for Ascend 950 extensions, its technological iteration advantages, and ecosystem community building. She dissected the core capabilities and technical advantages of the open-source platform, demonstrated its adaptability to multiple front-end languages, and presented developers with an easy-to-use, user-friendly, and open collaborative development experience with AscendNPU IR.

HyperAI has compiled and summarized the presentations without altering the original meaning.

Follow the WeChat official account "HyperAI" and reply with the keyword "0801 AI Compiler" to get the authorized lecturer's presentation PPT.

Why is Triton programming efficient and easy to use?

As a programming language, Triton is both efficient and user-friendly. So where does this development advantage come from?

In simple terms, Triton uses Python-like syntax and provides Block Level Tile programming, which shields developers from the low-level details of the abstract kernel's hardware memory, instructions, pipelining, and other operations, thereby lowering the barrier to operator development and improving development efficiency.

Triton's Tile programming model, compared to heterogeneous programming languages such as C++, offers advantages over traditional programming languages.This allows developers to focus more on data partitioning for operators, concentrating on Tensor op abstraction and logical expression.Based on the Tile-level programming abstraction, except for inter-core data scheduling which is managed by the user, intra-core memory access merging and shared-memory management can be completed automatically by the compiler.

More importantly, both the Triton dialect and the official Triton GPU compiler are based on a multi-level IR abstraction built on MLIR, allowing for in-depth optimization for different hardware.

AscendNPU IR continues to undergo architecture optimization.

AscendNPU IR is an Ascend hardware tile-level abstraction built on MLIR, serving underlying compilation optimizations. Its architecture is shown in the diagram below.AscendNPU IR interfaces downwards with LLVM IR, and is ultimately compiled into the Ascend hardware binary.Upper-layer third-party languages and compilers such as Triton can be integrated into AscendNPU IR, thereby achieving Ascend compatibility.

Key technical features of AscendNPU IRFirstly, it provides a bottom-up abstraction for Ascend hardware, abstracting different resources within the core, such as memory and data transport units. Secondly, it provides abstract tile-level operators, enabling unified cross-architecture support for underlying hardware, covering the latest hardware architectures from A2/A3 to Ascend 950 SIMD and SIMT. Based on the openness and scalability of MLIR, AscendNPU IR supports access from more upper-layer languages.

The diagram below shows the AscendNPU IR compilation architecture design, from Ascend A2/A3 to the latest Ascend 950 hardware released this year. The left side shows the first version of the AscendNPU IR architecture released last year, mainly composed of two layers. The upper layer is the hardware-independent layer—HFusion, which mainly implements multi-dimensional fusion optimization. Its role is based on the community Linalg and extended OPs, connecting upwards to various standardized dialects, completing data type and OP normalization preprocessing, and fusion optimizations such as Auto Schedule.

The hardware-related layer HIVM mainly performs high-level abstraction of Ascend hardware, completing core mapping compilation, on-chip memory mapping, and processing unit mapping in sequence.

The first step is kernel mapping compilation.This includes Cube kernels and Vector kernels. High-level tile representations are mapped through kernels, with Cube kernels assigned to Cube kernels and Vector kernels to Vector kernels. Data communication, data synchronization, and intermediate workspace memory management between them are all automatically completed by the compiler, which also performs high-order optimizations for CV pipelined parallelism.

Next is on-chip memory mapping. Both the Cube core and the Vector core have their own on-chip memory. For a Tensor to work on Ascend hardware, the logical Tensor data needs to be deduced into the corresponding physical memory, as well as the Cube-specific fractal matrix format derivation. On-chip memory allocation is all handled automatically by the compiler.

The lowest level is the processing unit mapping.Pipeline synchronization ensures timing and parallel efficiency. Vectorization and tensor quantization enable high-performance hardware instructions to complete the final compilation optimization.

The right side of the diagram shows some of the major changes in the Ascend 950 architecture from last year's A2/A3 to this year's. The upper hardware-independent layer HFusion remains largely stable, with the main expansion coming from the hardware-dependent layer HIVM, extending from Membased SIMD to support Regbased SIMD and SIMT. Previously, CV fusion compilation involved interaction between Cube cores and Vector cores via Global Memory; the Ascend 950 implements a more tightly coupled approach for more efficient data interaction.

Overall, the new architecture still maintains a stable two-layer structure, but with the changes in architecture, Tile compilation has undergone a significant upgrade.

Let's first look at the evolution of the Vector Regbased SIMD architecture. In MemBased SIMD, Vector computation is performed on-chip using a unified buffer, and instructions are executed directly from on-chip memory. On the Ascend 950, computation is performed using high-speed registers, adding an extra register layer. For backend optimization, especially for Vector, the compilation process changes as follows:

First, vectorization for Regbased systems. Tensor operations are mapped to fixed-width register instructions via vectorized partitioning. Data is loaded from on-chip memory into registers, computations are performed based on the registers, and then the data is stored back into on-chip memory, thus forming a Load-Compute-Store loop.

Second, register granularity fusion. Operations that repeatedly load and unload from on-chip memory are inefficient. Deep fusion based on registers allows data to reside in registers. Merging multiple Triton operations into a single loop reduces load storage. This involves utilizing different search algorithms and cost models to develop better fusion strategies.

Third, AVE dialect abstraction.  Ascend's user-friendly Vector dialect extension masks underlying hardware mask differences, enabling mask analysis and derivation. Hardware Vector Load offers rich in-path optimization capabilities. Enabling in-path operations allows multiple instructions to be combined into a single instruction. These are all optimizations specifically designed for the Regbase SIMD dialect.

The Ascend 950 hardware adds support for SIMT units, as shown in green in the image below. This allows the Vector core to have three different computation modes: pure SIMD, pure SIMT, and a hybrid alternating execution of SIMT and SIMD. SIMT can accelerate discrete memory access scenarios, while dense computation can utilize SIMD for more aggressive vector optimization.

Before proceeding with Vector SIMD or SIMT compilation, the compilation process performs a fusion analysis and optimization. It identifies which components are suitable for SIMT, extracts them for separate SIMT and SIMD compilations, and then combines them for overall Vector computation. During this process, both SIMT and SIMD are modeled and analyzed.

Closely related to Attention-type operators is the upgrade and change of CV kernels. On A2/A3 processors, the Cube and Vector lack on-chip connectivity and can only interact via off-chip Global Memory. However, on the Ascend 950, as shown in red in the diagram, the Cube result resides in the L0C and can be directly copied to the Vector's on-chip memory. Conversely, the Vector's computation result can also be directly copied to the Cube's on-chip memory for subsequent Cube computation. In summary, tight coupling and rapid data exchange in computer vision greatly benefit attention performance optimization.

One of the key technologies for optimizing CV operators is CV pipeline parallelism. Ascend Automation's CV Pipeline can better hide latency and fully utilize the computing power resources of Cube cores and Vector cores.

Another key optimization is AutoSubTiling.The Ascend Cube and Vector cores are configured in a 1:2 ratio, with two Vector cores corresponding to one Cube core. The compiler performs Vector operations in a 1:2 block split, allowing each Vector core to compute half of the data simultaneously, thus completing the computation faster.

The latest version of the compiler's feature generalization work

This article will introduce the important generalization work of the newly released AscendNPU IR version.

First: CV data interaction. The inter-core data interaction involved in Cube and Vector operations is automatically completed by the compiler. Triton operator Cube and Vector calculations may be distributed across different branches of the control flow, adding complexity to the compilation and analysis process. Simple instruction insertion via Pattern Match is not feasible. In the example above, the two branches perform Cube and Vector calculations respectively. CV data interaction needs to support complex control flow derivation to ensure functional completeness.

The new version enhances the analysis process of key passes in InsertCVLoadStore, completing the CV data interaction through global derivation across complex control flows. First, deterministic anchor points are inserted: matrix multiplication ops input at L1 and output at L0C, while vector ops input and output are both at UB. Then, based on the initial anchor points, type casting operations are inserted to perform upward and downward derivation and propagation across control flows. After propagation, the memory level of all operand tensors is clearly defined.

As shown in the example above, a Tensor originally does not contain memory hierarchy semantics. After inserting a type cast Op, it is confined to the Vector core's UB memory; below it is the Cube core's L1 memory. A conflict arises during the derivation process—the memory hierarchy is inconsistent, requiring the insertion of a Load-Store or copy statement to establish data connectivity. The handling methods differ slightly between generations. A2/A3 Cube and Vector require inserting a Global Memory Load-Store to establish data connectivity, while Ascend 950 establishes data connectivity through on-chip data copying.

Second: MultiBuffer enhances complex control flow. MultiBuffer is a crucial prerequisite for pipelining and parallelism. So how can a Tensor be transformed into a MultiBuffer?

The previous solution only supported for loops. During loop iteration, MultiBuffer reused the iteration variable of the for loop, rotating the Buffer Slot according to this iteration variable. This solution was limited by its inability to support other while loops and nested control flows. Therefore, an independent Buffer counter was introduced. The counter follows the complex control flow through branching loops, thus accurately tracking Slot changes.

In addition, the alloc operation needs to be visible in the common loop layer of Cube and Vector, so that both Cube and Vector can perform memory allocation operations simultaneously, such as allocating 2 or 3 Buffers, to ensure consistency of MultiBuffer transformation.

Third: AutoBlockify, a multi-core transformation closely related to Ascend operator optimization.When writing CUDA or Triton operators, we typically divide the code into many logic blocks, which are then executed via hardware scheduling. Ascend relies heavily on repetitive round-robin scheduling. Merging different logic cores into a single for loop and letting software loop scheduling optimize performance. Furthermore, advanced optimization techniques such as MultiBuffer and pipelining can be applied within the for loop to further improve out-of-the-box performance.

Fourth: CV Pipeline.For example, the FlashAttention operator has four tasks distributed across two cores, Cube and Vector, executed separately. These tasks are executed sequentially, and computations in Cube and Vector need to wait for each other. During the compilation process, MultiBuffer can be used to ensure that Cube and Vector are executed as early as possible, achieving a pipelined parallel acceleration effect.

First, identify the C/V boundary points, then accurately divide the code segment into Cube Scope and Vector Scope. Next, apply MultiBuffer to implement pipelined parallel transformations based on the scheduling algorithm, including various strategy modes: Unroll, Skew, and dynamic. For different computational loads, the optimal pipelined parallel mode can be selected.

Fifth: CV 1:2 split.The two Vector kernels and one Cube kernel collaborate to process the raw Tensor data in half as accurately as possible to achieve efficient Vector computation. Ideally, parallel axes would be found in the high dimension to perform a global halving. However, candidate axes may undergo computational operations such as Transpose, Broadcast, and Reduction, increasing the complexity of the CV 1:2 split.

First, perform a global dimensional analysis to identify the most suitable axis for splitting. After determining the splitting axis, insert a splitting operation start marker starting from the end of the Store operation, and then bubble up from this end to the root node, ultimately achieving a perfect split. If some complex scenarios are encountered that cannot be split, revert to the conservative 1:1 unsplit state.

Furthermore, in some complex scenarios, simply splitting the parallel axis is insufficient for performance. For example, the FB8 low-precision attention algorithm experiences high vector load and requires splitting and reducing the reduce axis. Since the two vector kernels cannot directly exchange data, intermediate reduce results need to be exchanged through Global Memory for secondary reduction, completing the overall reduction and splitting. In addition, scenarios such as dynamic shapes also require enhanced CV 1:2 splitting methods.

Triton-Ascend and AscendNPU IR are open source and collaboratively developed.

AscendNPU IR is now open source on the Ascend Community. You can scan the QR code in the image to follow the open source project. The AscendNPU IR Community SIG bi-weekly meeting has been published in the community activity calendar. Developers can pay attention to the weekly discussion topics and meeting minutes in advance.

The Ascend Community has launched open-source internship projects and community tasks, providing developers with a way to participate in ecosystem building. You can see the latest task status here and claim your tasks. These community tasks vary in difficulty; please choose and participate according to your own abilities.

Open source internships and community tasks are available via quick links on the AscendNPU IR and Triton-Ascend open source repository homepages, which include the task claiming process. Each person can only claim one task at a time. Community tasks will be refreshed regularly, and the development process will remain interactive and open throughout. Each task offers generous rewards; everyone is welcome to choose freely.

For participants without an Ascend development environment, the community provides a free Ascend computing power platform, HiDevLab (https://hidevlab.huawei.com/home). Developers can register and apply for computing power, and upon successful application, they will be allocated 100 hours of free computing time by default. Developers who have accepted community tasks can apply for computing power to complete the development and verification of the tasks.

Thank you for your attention!