Hardware-aware onnx graph optimization method and hardware-aware graph optimization and compilation engine

Abstract

The present application relates to the field of artificial intelligence, and provides a hardware-aware ONNX graph optimization method and a hardware-aware graph optimization and compilation engine, which breaks the barrier between hardware-independent optimization and hardware-dependent compilation, optimizes and compiles an original ONNX model calculation graph based on the hardware characteristics of an AI chip through a unified and closed-loop optimization framework, realizes deep collaborative optimization from the ONNX calculation graph to the AI chip code, and thus maximizes the performance of the AI chip.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and more specifically, to a hardware-aware ONNX graph optimization method and a hardware-aware graph optimization and compilation engine. Background Technology

[0002] With the rapid development of artificial intelligence technology, deep learning models are widely used in fields such as computer vision, natural language processing, and autonomous driving. To facilitate the migration of deep learning models between different frameworks (such as PyTorch and TensorFlow) and different hardware platforms, the Open Neural Network Exchange Format (ONNX) has emerged. ONNX defines a neutral computation graph representation, enabling trained models to be deployed independently of native frameworks.

[0003] However, ONNX models still face significant challenges when deployed to specific AI chips (such as ASICs, FPGAs, and NPUs). AI chips typically have unique computational units (such as tensor cores), memory hierarchies (such as global caches, local caches, and registers), and instruction sets. The general-purpose ONNX runtime cannot fully utilize these hardware characteristics, resulting in high latency, low throughput, and poor energy efficiency in model inference.

[0004] Existing deployment schemes for the ONNX model mainly include a two-stage approach: "ONNX graph optimization based on general rules + independent backend compiler".

[0005] ① General ONNX Graph Optimization Stage: Utilize tools such as ONNX Runtime and ONNX Optimizer to perform graph optimization at a hardware-independent level. Common optimizations include: constant folding, redundant node elimination (e.g., continuous reshape merging), operator fusion (e.g., Conv-BN-ReLU fusion), and common subexpression elimination.

[0006] ② Independent backend compilation stage: The optimized ONNX graph is passed to the AI ​​chip-specific compiler (such as TVM or the manufacturer's own compiler). This compiler is responsible for converting the ONNX operators into low-level code or instructions for the chip and performing hardware-related optimizations, such as operator splitting, memory allocation, and pipeline scheduling.

[0007] In the existing two-stage optimization scheme, hardware-independent graph optimization and hardware-related compilation optimization are separated, which makes it impossible to perform deep collaborative optimization from a global perspective. This limits the full potential of AI chip hardware and makes it difficult to achieve optimal end-to-end inference performance.

[0008] The main drawbacks of the existing solution are as follows:

[0009] (1) Problem 1: Graph optimization may destroy the opportunity for backend optimization.

[0010] ① The general graph optimizer is unaware of the specific instructions of the backend hardware. For example, an AI chip may have a dedicated compound instruction Fused_Conv_BN_Add_Activation.

[0011] ② If the general optimizer first merges Conv and BN into a new Conv node, when the backend compiler sees this merged node, it may not be able to match it with the subsequent Add and Activation nodes to that dedicated compound instruction, thus missing the opportunity to use more efficient instructions.

[0012] ③Performance degradation.

[0013] (2) Problem 2: Graph optimization lacks hardware awareness.

[0014] ① When performing operator replacement or decomposition, the general optimizer only considers the correctness and generality of the computation graph, without considering the computation and memory characteristics of the target hardware.

[0015] ② For example, for the sake of generality, a DepthwiseConv might be decomposed into multiple Slices and Convs. However, for chips with efficient DepthwiseConv hardware units, such decomposition can lead to a sharp performance drop. Conversely, for hardware that does not support a certain complex operator (such as EinSum), a general-purpose optimizer may not actively decompose it into more basic operators.

[0016] ③ It may produce suboptimal or even inefficient computational graphs.

[0017] (3) Problem 3: The optimization goal is too singular and difficult to balance.

[0018] ① In two-stage optimization, the optimization objectives of the preceding and following stages (such as latency vs. memory usage) may be inconsistent.

[0019] ② General graph optimization may introduce additional transpose operations to reduce computation, but this increases memory movement overhead. For AI chips with limited memory bandwidth, this optimization may be counterproductive. Due to the fragmented stages, this trade-off cannot be evaluated in the early stages of optimization.

[0020] ③ It cannot achieve optimal performance for specific hardware and deployment scenarios (such as high throughput vs. low latency). Summary of the Invention

[0021] This invention addresses the technical problems existing in the prior art by providing a hardware-aware ONNX graph optimization method and a hardware-aware graph optimization and compilation engine.

[0022] According to a first aspect of the present invention, a hardware-aware ONNX graph optimization method is provided, comprising:

[0023] A hardware-aware graph optimization and compilation engine is constructed, and a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer are configured in the hardware-aware graph optimization and compilation engine;

[0024] The unified optimizer receives the input ONNX model and the target AI chip identifier, and obtains the hardware information database of the target AI chip based on the target AI chip identifier.

[0025] Based on the hardware information database, pattern matching is performed on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instruction of the AI ​​chip, and the subgraph is replaced with the corresponding fusion operator;

[0026] Based on the performance cost model of each operator in the hardware information database, multiple candidate decomposition schemes are generated for complex operators in the ONNX model computation graph that are not natively supported by the target AI chip. The optimal decomposition scheme is selected from the multiple candidate decomposition schemes. Based on the optimal decomposition scheme, the complex operators that are not natively supported by the target AI chip are decomposed to obtain the optimized ONNX model computation graph.

[0027] Based on the just-in-time compiler and performance analyzer, multiple versions of kernel code are generated for operators or subgraphs with multiple implementations in the optimized ONNX model computation graph.

[0028] Micro-benchmark tests are run on each version of the kernel code to obtain the real performance data of each version of the kernel code. The real performance data of each version of the kernel code is fed back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator.

[0029] The ONNX model computation graph is iteratively optimized in multiple rounds, and the final optimized ONNX model computation graph is compiled into an executable file or instruction stream on the target AI chip.

[0030] Based on the above technical solution, the present invention can also be improved as follows.

[0031] Optionally, the hardware information database includes a list of native operators supported by the target AI chip, dedicated composite instructions, the performance cost model of each operator, and memory hierarchy parameters. The list of native operators includes all native operators supported by the target AI chip, as well as the data layout and data type supported by each native operator on the target AI chip hardware.

[0032] Optionally, the step of performing pattern matching on the ONNX model computation graph based on the hardware information database to identify the subgraph corresponding to the dedicated composite instruction of the AI ​​chip, and replacing the subgraph with the corresponding fusion operator, includes:

[0033] Based on the unified optimizer, the ONNX model computation graph is traversed to find a subgraph pattern in the ONNX model computation graph that matches the dedicated composite instructions supported by the target AI chip.

[0034] Replace the subgraph pattern with a single custom operator and mark the dedicated composite instruction corresponding to the custom operator.

[0035] Optionally, based on the performance-cost model of each operator in the hardware information database, multiple candidate decomposition schemes are generated for complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, including:

[0036] Traverse the ONNX model computation graph. When encountering a complex operator, query the hardware information database. When the complex operator is not natively supported by the target AI chip, decompose the complex operator. There are multiple candidate decomposition schemes when decomposing the complex operator.

[0037] The step of selecting the optimal decomposition scheme from multiple candidate decomposition schemes includes:

[0038] Based on the performance cost model of each operator, the execution performance cost of multiple operators after decomposition using each candidate decomposition scheme is calculated on the target AI chip.

[0039] Based on the execution performance cost of each candidate decomposition scheme, the candidate decomposition scheme with the lowest performance cost is selected as the optimal decomposition scheme.

[0040] Optionally, the calculation of the execution performance cost of multiple operators decomposed using each candidate decomposition scheme on the target AI chip based on the performance cost model of each operator includes:

[0041] Based on the performance cost model of each operator, the computational cost and memory cost of each candidate decomposition scheme on the target AI chip hardware are calculated.

[0042] Based on the computation cost and the memory cost, the total cost is calculated by weighted summation and used as the execution performance cost.

[0043] Based on the total cost of each candidate decomposition scheme, the candidate decomposition scheme with the minimum total cost is selected as the optimal candidate decomposition scheme.

[0044] Optionally, the total cost is calculated by weighted summation based on the computational cost and the memory cost, including:

[0045] Cost(op) = α * Comp(op) + β * Memory(op)

[0046] Where Cost(op) is the total cost, Comp(op) is the estimated computation cost, Memory(op) is the estimated memory access cost, and α and β are weighting coefficients calibrated for the target AI chip hardware. They are adjusted according to the hardware characteristics of the target AI chip. On computationally intensive chips, the value of α is larger; on chips with limited memory bandwidth, the value of β is larger.

[0047] Optionally, based on the data layout and data type supported by each native operator in the ONNX model computation graph included in the native operator list on the target AI chip hardware, the unified optimizer intelligently inserts or deletes format conversion nodes in the ONNX model computation graph to convert the data format of the native operators preceding its own position, so that the format of the converted native operators is supported by the target AI chip hardware.

[0048] Optionally, the step of running micro-benchmark tests for each version of the kernel code to obtain real performance data for each version of the kernel code includes:

[0049] In a lightweight hardware simulator or real device environment, micro-benchmark tests are run on each version of the kernel code to obtain real performance data for each candidate kernel, including the number of clock cycles.

[0050] Optionally, the multi-round iterative optimization of the ONNX model computation graph includes:

[0051] The ONNX model computation graph is iteratively optimized multiple times based on the unified optimizer until the maximum number of iterations is reached or the optimized ONNX model computation graph no longer changes.

[0052] According to a second aspect of the present invention, a hardware-aware graph optimization and compilation engine is provided for optimizing ONNX graphs, wherein the hardware-aware graph optimization and compilation engine is configured with a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer.

[0053] The unified optimizer receives the input ONNX model and the target AI chip identifier. Based on the target AI chip identifier, it obtains the hardware information library of the target AI chip. Based on the hardware information library, it performs pattern matching on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instructions of the AI ​​chip, and replaces the subgraph with the corresponding fusion operator. Based on the performance cost model of each operator in the hardware information library, it generates multiple candidate decomposition schemes for complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, selects the optimal decomposition scheme from the multiple candidate decomposition schemes, and decomposes the complex operators that are not natively supported by the target AI chip based on the optimal decomposition scheme to obtain the optimized ONNX model computation graph.

[0054] The just-in-time compiler and performance analyzer is used to generate multiple versions of kernel code for operators or subgraphs with multiple implementations in the optimized ONNX model computation graph; to run micro-benchmark tests on each version of kernel code to obtain the real performance data of each version of kernel code; and to feed the real performance data of each version of kernel code back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator.

[0055] The hardware-aware graph optimization and compilation engine is used to perform multiple rounds of iterative optimization on the ONNX model computation graph based on the unified optimizer and the just-in-time compilation and performance analyzer, and compile the final optimized ONNX model computation graph into an executable file or instruction stream on the target AI chip.

[0056] This invention provides a hardware-aware ONNX graph optimization method and a hardware-aware graph optimization and compilation engine, which breaks down the barriers between hardware-independent optimization and hardware-dependent compilation. Through a unified, closed-loop optimization framework, it optimizes and compiles the original ONNX model computation graph based on the hardware characteristics of the AI ​​chip, realizing deep collaborative optimization from the ONNX computation graph to the AI ​​chip code, thereby maximizing the performance of the AI ​​chip. Attached Figure Description

[0057] Figure 1 A flowchart of a hardware-aware ONNX graph optimization method provided in one embodiment of the present invention;

[0058] Figure 2 A schematic diagram of the overall working architecture of a hardware-aware ONNX graph optimization method provided in an embodiment of the present invention;

[0059] Figure 3 This is a schematic diagram illustrating the optimization process of the ONNX model calculation graph;

[0060] Figure 4This is a schematic diagram of the compilation process for the ONNX model calculation graph;

[0061] Figure 5 Example diagram of the specific implementation process for optimizing the compilation of ONNX model computation graph on a self-developed NPU;

[0062] Figure 6 This is a schematic diagram of the architecture of a hardware-aware graph optimization and compilation engine according to an embodiment of the present invention. Detailed Implementation

[0063] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention. In addition, the technical features of the various embodiments or individual embodiments provided by the present invention can be arbitrarily combined with each other to form feasible technical solutions. Such combinations are not constrained by the order of steps and / or structural composition patterns, but must be based on the ability of those skilled in the art to implement them. When the combination of technical solutions is contradictory or cannot be implemented, it should be considered that such a combination of technical solutions does not exist and is not within the scope of protection claimed by the present invention.

[0064] Figure 1 A flowchart of a hardware-aware ONNX graph optimization method provided in one embodiment of the present invention is shown below. Figure 1 As shown, the method includes the following steps:

[0065] Step 1: Construct a hardware-aware graph optimization and compilation engine, and configure a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer in the hardware-aware graph optimization and compilation engine.

[0066] Step 2: The unified optimizer receives the input ONNX model and the target AI chip identifier, and obtains the hardware information database of the target AI chip based on the target AI chip identifier.

[0067] Step 3: Based on the hardware information database, perform pattern matching on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instruction of the AI ​​chip, and replace the subgraph with the corresponding fusion operator;

[0068] Step 4: Based on the performance cost model of each operator in the hardware information database, generate multiple candidate decomposition schemes for complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, select the optimal decomposition scheme from the multiple candidate decomposition schemes, and decompose the complex operators that are not natively supported by the target AI chip based on the optimal decomposition scheme to obtain the optimized ONNX model computation graph.

[0069] Step 5: Based on the just-in-time compiler and performance analyzer, calculate operators or subgraphs with multiple implementation methods in the optimized ONNX model computation graph, and generate multiple versions of kernel code;

[0070] Step 6: Run micro-benchmark tests on each version of the kernel code to obtain the real performance data of each version of the kernel code. Feed the real performance data of each version of the kernel code back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator.

[0071] Step 7: Perform multiple rounds of iterative optimization on the ONNX model computation graph, and compile the final optimized ONNX model computation graph into an executable file or instruction stream on the target AI chip.

[0072] The core of this invention is the construction of a hardware-aware, unified graph optimization and compilation system. Its overall architecture and workflow are as follows: Figure 2 As shown.

[0073] Figure 2 The unidirectional flow in the code represents the system's main data flow and control flow, including:

[0074] ① Start with the original ONNX model input.

[0075] ② Enter the core hardware-aware graph optimization and compilation engine, and optimize and compile the original ONNX model computation graph based on the hardware characteristics of the target AI chip.

[0076] ③ Based on the optimized ONNX model calculation graph, a highly optimized executable file is finally produced.

[0077] The closed-loop flow within the engine embodies the core innovation of this invention: a feedback optimization loop, including:

[0078] ① Unified Optimizer: First, it queries the hardware information database and makes decisions and performs graph transformations based on the performance cost model and operator information therein.

[0079] ② The transformed computation graph is fed into the just-in-time compiler and performance analyzer, where real performance data is obtained through actual micro-benchmark tests.

[0080] ③ This real performance data is fed back to the hardware information database through a feedback loop, which is used to dynamically update and calibrate its internal performance cost model, enabling it to make more accurate decisions in future iterations.

[0081] This closed loop enables the system to learn and continuously optimize itself.

[0082] The following is about Figure 2 The overall architecture and workflow shown in the document are explained in detail.

[0083] (1) Initialize the hardware perception graph optimization and compilation engine and load the hardware information database of the target AI chip:

[0084] ① The hardware-aware graph optimization and compilation engine receives the original ONNX model and the identifier of the target AI chip.

[0085] ② Load the hardware information database of the target AI chip based on its identifier. This hardware information database includes: a list of native operators supported by the target AI chip (including all native operators supported by the target AI chip), dedicated compound instructions (such as Fused_Conv_BN_Relu), the performance cost model of each operator (which can be a model based on theoretical calculations / memory access, or empirical data obtained through profiling), memory hierarchy parameters, etc.

[0086] (2) The iterative optimization process of the unified optimizer, wherein... Figure 3 The workflow for optimizing the ONNX model computation graph using a unified optimizer is demonstrated, including:

[0087] ① The unified optimizer no longer distinguishes between "graph optimization" and "back-end optimization," but performs a unified graph transformation under the guidance of hardware information.

[0088] ②Key Optimization Sub-process 1: Hardware-Aware Pattern Matching and Fusion

[0089] A) The unified optimizer traverses the ONNX model computation graph, searching for subgraph patterns that match the dedicated composite instructions supported by the target AI chip. For example, it identifies the pattern Conv -> BN -> Add -> ReLU and queries the hardware information database to find the corresponding Fused_Conv_BN_Add_ReLU instruction.

[0090] B) Decision: Directly replace the subgraph with a single custom operator Fused_Op and mark its corresponding dedicated compound instruction. This avoids premature fusion by the general optimizer, which could disrupt the matching pattern.

[0091] ③Key Optimization Sub-process 2: Operator Selection and Decomposition Based on Performance-Cost Model:

[0092] A) When a complex operator (such as EinSum) is encountered in the ONNX model computation graph, the unified optimizer queries the hardware information database. If the complex operator is not natively supported by the target AI chip, it selects the optimal decomposition scheme from multiple decomposition schemes based on the performance cost model of each operator stored in the hardware information database to decompose the complex operator, so that the decomposed multiple operators are all supported by the target AI chip.

[0093] A simplified cost model can be defined as follows:

[0094] Cost(op) = α* Comp(op) + β* Memory(op)

[0095] Wherein, Comp(op): estimated computational overhead (e.g., FLOPs), Memory(op): estimated memory access overhead (e.g., number of bytes read and written), α and β: weighting coefficients calibrated for the target hardware, which can be inferred from the performance of the profiling basic operators. For example, α is larger on compute-intensive chips, while β is larger on chips with limited memory bandwidth.

[0096] B) Example: EinSum can be decomposed into Transpose->MatMul or Reshape->BatchMatMul. The unified optimizer estimates the computational and memory costs of each decomposition scheme on the target hardware, calculates the total cost based on the computational and memory costs, selects the decomposition scheme with the lowest total cost as the optimal decomposition scheme, and decomposes the complex operator based on the optimal decomposition scheme.

[0097] ④ Key optimization sub-process 3: Tensor data layout transformation optimization.

[0098] The hardware database stores the data layout and data type of each operator in the ONNX model computation graph that are supported by the target AI chip. The unified optimizer can intelligently insert format conversion nodes into the ONNX model computation graph to convert the tensor data format of the operators in the ONNX model computation graph into the data format supported by the target AI chip, so that the tensor layout is best suited to the hardware computing unit. The decision is based on a trade-off between computational overhead and memory handling overhead, which is guided by the hardware cost model.

[0099] (3) Just-in-time compilation and performance analyzer (feedback loop). Figure 4 The workflow for compiling the optimized ONNX model computation graph is demonstrated, including:

[0100] ① For operators or subgraphs in the optimized computation graph that have multiple implementations (e.g., convolution has different loop unrolling strategies and block strategies), the system will start the just-in-time compiler and performance analyzer.

[0101] ② The just-in-time compiler and performance analyzer generate multiple versions of kernel code for these "candidate points".

[0102] ③Then, in a lightweight hardware simulator or real machine environment, run micro-benchmark tests on each version of the kernel code to obtain real performance data (such as clock cycles) for each version of the kernel code.

[0103] ④ This real performance data is fed back to the hardware information database to update and calibrate the corresponding performance cost model, for example, by adjusting the weighting coefficients of α and β to make them more accurate.

[0104] 4) Code generation:

[0105] After multiple rounds of iterative optimization and final decision-making, the system compiles the optimized computation graph into the most efficient executable file or instruction stream on the target AI chip.

[0106] The hardware-aware ONNX graph optimization method provided by this invention is based on the construction of a unified and closed hardware-aware graph optimization and compilation engine, which loads the hardware information library of the target AI chip, a unified optimizer, and a just-in-time compiler and performance analyzer.

[0107] Among them, the construction and use of the hardware information database: abstracting hardware characteristics (operators, instructions, cost models) into a queryable database and using it as the core decision basis for graph optimization.

[0108] A unified, hardware-aware optimizer: integrates the traditional two-stage optimization into a single stage, directly considering hardware constraints and advantages during graph transformation.

[0109] ③ Feedback loop based on just-in-time compilation and micro-benchmarking: Use real performance data to dynamically calibrate optimization decisions, enabling the system to have self-learning and self-adaptive capabilities.

[0110] ④ Cost-based trade-off decision-making mechanism: Among multiple feasible graph transformation schemes, a unified cost model is used for quantitative comparison and selection.

[0111] It should be noted that the scenarios in which the present invention is used include:

[0112] ① AI chip manufacturers: As the core compilation and optimization component of their software stack (SDK), it highlights their hardware performance advantages.

[0113] ② Cloud service provider: Provides automated performance optimization for model services deployed on heterogeneous AI accelerator cards.

[0114] ③ Edge computing devices: On resource-constrained edge devices, extreme optimization is used to meet the stringent requirements of low latency and low power consumption.

[0115] For a detailed implementation process of optimizing the ONNX model computation graph on the self-developed NPU chip, please refer to [link to relevant documentation]. Figure 5 It implements a hardware-aware ONNX compiler, which specifies the target NPU architecture and loads the corresponding hardware information library during initialization, while creating a unified optimizer and a just-in-time compiler with performance analysis. The main compilation process consists of three core stages:

[0116] 1) First, initialize and load the NPU chip's hardware information library;

[0117] 2) Subsequently, the ONNX model computation graph is iteratively optimized using a unified optimizer, a process that fully considers the characteristics of the target hardware;

[0118] 3) Finally, executable code is generated based on the optimized computation graph.

[0119] The entire design implements a complete compilation chain from the hardware-independent ONNX model to hardware-optimized code.

[0120] The following section provides a detailed explanation of the implementation of the multiple core stages included in the main compilation process.

[0121] 1. Hardware information database implementation example:

[0122] A complete hardware information base was constructed to store and manage the architectural characteristics of the target NPU. During initialization, the hardware information base loads the basic operator set supported by the NPU, the dedicated compound instruction set, the performance cost model, and memory architecture information.

[0123] 1) The basic operator library defines in detail the data layout and data types supported by each operator;

[0124] 2) The composite instruction library describes the fusionable operation modes, corresponding hardware instructions, and execution constraints;

[0125] 3) The performance query interface can estimate the running cost based on the operator type and input shape.

[0126] This hardware information base provides complete infrastructure support for compiler hardware-aware optimization.

[0127] 2. Implementation of the performance-cost model

[0128] A dynamic performance-cost model was implemented to accurately estimate the execution cost of neural network operators on hardware. During initialization, the performance-cost model loads a basic performance profile and presets weight coefficients for computation and memory access. The cost estimation process comprehensively considers both computational and memory overhead costs, obtaining the total cost through a weighted sum of the weighted coefficients of these costs. These coefficients can be adjusted for specific implementation methods. The performance-cost model also possesses online learning capabilities, dynamically adjusting weight parameters based on actual performance feedback: when the prediction error exceeds a threshold, it automatically identifies computational or memory bottlenecks and adjusts the weight allocation accordingly, thereby continuously optimizing the accuracy of cost predictions.

[0129] 3. Unified optimizer implementation process:

[0130] The unified optimizer functionality has been implemented. The following is a description of the code's text fields:

[0131] The class, named UnifiedOptimizer, is a hardware-aware optimizer primarily used for optimizing the computation graph. During initialization, it receives hardware information and sets up four optimization channels: hardware-aware pattern matching, cost-based operator decomposition, layout transformation optimization, and memory allocation optimization.

[0132] The unified optimizer employs an iterative optimization strategy, performing a maximum of three rounds of optimization. Each round of optimization executes all optimization channels sequentially until the maximum number of iterations is reached or the computation graph no longer changes.

[0133] Key optimization features include:

[0134] 1) Hardware-aware pattern matching: Traverse the nodes in the ONNX model computation graph to find subgraph patterns that match hardware composite instructions. Once a matching pattern is found, the subgraph is merged into a dedicated composite instruction supported by the NPU to replace the original structure.

[0135] 2) Cost-based operator decomposition: For operators not natively supported by the hardware, multiple decomposition schemes are provided, and the optimal decomposition strategy is selected through a performance-cost model. Cost estimation helps identify the most cost-effective decomposition method.

[0136] In addition, the code includes helper methods for pattern matching and decomposition scheme selection, ensuring that the optimization process is both efficient and adaptable to specific hardware characteristics.

[0137] In summary, this unified optimizer aims to improve the execution efficiency of the computation graph on specific hardware through a combination of multiple iterations and various optimization techniques.

[0138] 4. Just-in-Time (JIT) Compilation and Performance Analysis Implementation:

[0139] A just-in-time compilation feature with performance analysis has been implemented. The following is a text description of the code:

[0140] The class, named JITCompilerWithProfiling, is a just-in-time compiler that integrates performance profiling capabilities. During initialization, it automatically sets up the performance profiling environment, preparing it for subsequent benchmarking tests.

[0141] The core function of this compiler is to perform performance analysis on candidate implementations. It takes a computation graph and a set of candidate operations as input and evaluates the performance of each candidate implementation through systematic micro-benchmarking. The analysis process executes performance tests and collects key metrics for each implementation scheme under each operation identifier.

[0142] The performance testing phase involves running kernel code on an NPU simulator or real hardware to capture multi-dimensional performance data, including:

[0143] 1) Number of execution cycles;

[0144] 2) Memory bandwidth utilization;

[0145] 3) Power consumption;

[0146] 4) Cache efficiency;

[0147] These detailed performance metrics are fed back to the hardware database, updating and calibrating the performance-cost model within it. This provides data support for subsequent optimization decisions, ensuring that the compiler can select the optimal implementation based on actual runtime characteristics. The entire design embodies a data-driven optimization philosophy, guiding the compiler optimization process through empirical analysis rather than theoretical speculation.

[0148] To achieve a similar "hardware-aware optimization" goal, the following alternatives may exist, but each has its limitations:

[0149] 1) Graph optimization scheme based on reinforcement learning:

[0150] ① Description: The graph optimization process is modeled as a Markov decision process, where the agent (optimizer) learns through trial and error to select graph transformation actions (such as fusion and decomposition) and uses the final inference performance as a reward.

[0151] ② Comparison: While this method can also achieve deep optimization, it requires enormous computational resources and time for training, and the training results are difficult to interpret and guarantee. The present invention, based on a clear cost model and rules, is more interpretable and practical for engineering applications.

[0152] 2) Template-based code generation solution:

[0153] ① Description: Manually optimized template code is pre-written for all possible operator fusion patterns and subgraph patterns. The optimizer is only responsible for matching patterns and calling the corresponding templates.

[0154] ② Comparison: The performance limit of this method depends on the quality of the hand-crafted template, has poor scalability (adding new operators or patterns requires rewriting the template), and cannot handle unseen graph structures. The solution of this invention, through a cost model and just-in-time compilation, possesses stronger generalization ability and flexibility.

[0155] 3) Scheme to extend the ONNX operator set:

[0156] ① Description: Define all AI chip-specific instructions as standard ONNX operators, and then use them in the general optimization phase.

[0157] ② Comparison: This requires modifying the ONNX standard, making it difficult to unify the ecosystem, and causing the ONNX model to be tied to specific hardware, losing its cross-platform compatibility. The solution in this invention, while maintaining the ONNX standard unchanged, performs internal mapping during the compilation stage, making it more feasible.

[0158] See Figure 6 This invention provides a hardware-aware graph optimization and compilation engine for optimizing ONNX graphs. The hardware-aware graph optimization and compilation engine is configured with a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer.

[0159] The unified optimizer receives the input ONNX model and the target AI chip identifier. Based on the target AI chip identifier, it obtains the hardware information library of the target AI chip. Based on the hardware information library, it performs pattern matching on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instructions of the AI ​​chip, and replaces the subgraph with the corresponding fusion operator. Based on the performance cost model of each operator in the hardware information library, it generates multiple candidate decomposition schemes for complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, selects the optimal decomposition scheme from the multiple candidate decomposition schemes, and decomposes the complex operators that are not natively supported by the target AI chip based on the optimal decomposition scheme to obtain the optimized ONNX model computation graph.

[0160] The just-in-time compiler and performance analyzer is used to generate multiple versions of kernel code for operators or subgraphs with multiple implementations in the optimized ONNX model computation graph; to run micro-benchmark tests on each version of kernel code to obtain the real performance data of each version of kernel code; and to feed the real performance data of each version of kernel code back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator.

[0161] The hardware-aware graph optimization and compilation engine is used to perform multiple rounds of iterative optimization on the ONNX model computation graph based on the unified optimizer and the just-in-time compilation and performance analyzer, and compile the final optimized ONNX model computation graph into an executable file or instruction stream on the target AI chip.

[0162] It is understood that the hardware-aware graph optimization and compilation engine optimization method for ONNX model computation graph provided by this invention corresponds to the hardware-aware ONNX graph optimization method provided in the foregoing embodiments. Therefore, the optimization method for ONNX model computation graph using the hardware-aware graph optimization and compilation engine can be found in the relevant technical features of the foregoing hardware-aware ONNX graph optimization method, which will not be repeated here.

[0163] The hardware-aware ONNX graph optimization method and hardware-aware graph optimization and compilation engine provided in this invention have the following advantages:

[0164] ① This invention uses a hardware information database and a unified optimizer to ensure that graph optimization decisions are always consistent with the characteristics of the target hardware.

[0165] ② Since the unified optimizer directly targets hardware composite instructions, it can make the most of the dedicated computing units provided by the chip, avoid missing the opportunity for a better fusion mode due to premature fusion, and achieve higher computing efficiency.

[0166] ③ All operator decomposition and replacement decisions are driven by a performance-cost model specific to the hardware, ensuring that the decomposed subgraph is best suited for computation on that hardware. This avoids performance regression caused by general optimization and improves the robustness and effectiveness of the optimization.

[0167] ④ A unified performance cost model can comprehensively consider multiple dimensions such as computing, memory, and bandwidth, enabling the unified optimizer to make intelligent trade-offs between different objectives (such as latency and memory) and find the most suitable optimization solution for a specific deployment scenario, thus achieving better end-to-end performance.

[0168] ⑤ By calibrating the cost model through real-time performance feedback, the system can adapt to performance fluctuations caused by different chip batches and temperature variations, and continuously evolves with use, resulting in increasingly precise optimization. Result: The system possesses long-term optimization capabilities and adaptability.

[0169] Therefore, the present invention can systematically and automatically generate deployment code that performs significantly better than existing two-stage solutions on specific AI chips.

[0170] It should be noted that the descriptions of each embodiment in the above embodiments have different focuses. For parts that are not described in detail in a certain embodiment, please refer to the relevant descriptions in other embodiments.

[0171] Those skilled in the art will understand that embodiments of the present invention can be provided as methods, systems, or computer program products. Therefore, the present invention can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Furthermore, the present invention can take the form of a computer program product embodied on one or more computer-usable storage media (including, but not limited to, disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0172] This invention is described with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the invention. It will be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, special-purpose computer, embedded computer, or other programmable data processing apparatus to produce a machine, such that the instructions, which execute via the processor of the computer or other programmable data processing apparatus, generate instructions for implementing the flowchart illustrations. Figure 1 One or more processes and / or boxes Figure 1 A device that provides the functions specified in one or more boxes.

[0173] These computer program instructions may also be stored in a computer-readable storage medium that can direct a computer or other programmable data processing device to function in a particular manner, such that the instructions stored in the computer-readable storage medium produce an article of manufacture including instruction means, which are implemented in a process Figure 1 One or more processes and / or boxes Figure 1 The function specified in one or more boxes.

[0174] These computer program instructions may also be loaded onto a computer or other programmable data processing equipment to cause a series of operational steps to be performed on the computer or other programmable equipment to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable equipment for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0175] Although preferred embodiments of the invention have been described, those skilled in the art, upon learning the basic inventive concept, can make other changes and modifications to these embodiments. Therefore, the appended claims are intended to be interpreted as including both the preferred embodiments and all changes and modifications falling within the scope of the invention.

[0176] Obviously, those skilled in the art can make various modifications and variations to this invention without departing from its spirit and scope. Therefore, if these modifications and variations fall within the scope of the claims of this invention and their equivalents, this invention also intends to include these modifications and variations.

Claims

1. A hardware-aware ONNX graph optimization method, characterized in that, include: A hardware-aware graph optimization and compilation engine is constructed, and a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer are configured in the hardware-aware graph optimization and compilation engine; The unified optimizer receives the input ONNX model and the target AI chip identifier, and obtains the hardware information database of the target AI chip based on the target AI chip identifier. Based on the hardware information database, pattern matching is performed on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instruction of the target AI chip, and the subgraph is replaced with the corresponding fusion operator; For complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, multiple candidate decomposition schemes are generated. Based on the performance cost model of each operator in the hardware information library, the optimal decomposition scheme is selected from the multiple candidate decomposition schemes. Based on the optimal decomposition scheme, the complex operators that are not natively supported by the target AI chip are decomposed to obtain the optimized ONNX model computation graph. Based on the just-in-time compiler and performance analyzer, multiple versions of kernel code are generated for operators or subgraphs with multiple implementations in the optimized ONNX model computation graph. Micro-benchmark tests are run on each version of the kernel code to obtain the real performance data of each version of the kernel code. The real performance data of each version of the kernel code is fed back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator. The ONNX model computation graph is iteratively optimized in multiple rounds, and the final optimized ONNX model computation graph is compiled into an executable file or instruction stream on the target AI chip. For complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, multiple candidate decomposition schemes are generated. Based on the performance-cost model of each operator in the hardware information database, the optimal decomposition scheme is selected from the multiple candidate decomposition schemes, including: Traverse the ONNX model computation graph. When encountering a complex operator, query the hardware information database. When the complex operator is not natively supported by the target AI chip, decompose the complex operator. There are multiple candidate decomposition schemes when decomposing the complex operator. The step of selecting the optimal decomposition scheme from multiple candidate decomposition schemes includes: Based on the performance cost model of each operator, the execution performance cost of multiple operators after decomposition using each candidate decomposition scheme is calculated on the target AI chip. Based on the execution performance cost of each candidate decomposition scheme, the candidate decomposition scheme with the lowest performance cost is selected as the optimal decomposition scheme.

2. The hardware-aware ONNX graph optimization method according to claim 1, characterized in that, The hardware information database includes a list of native operators supported by the target AI chip, dedicated compound instructions, the performance cost model of each operator, and memory hierarchy parameters. The list of native operators includes all native operators supported by the target AI chip, as well as the data layout and data type supported by each native operator on the target AI chip hardware.

3. The hardware-aware ONNX graph optimization method according to claim 2, characterized in that, The step of performing pattern matching on the ONNX model computation graph based on the hardware information database to identify the subgraph corresponding to the dedicated composite instruction of the AI ​​chip, and replacing the subgraph with the corresponding fusion operator, includes: Based on the unified optimizer, the ONNX model computation graph is traversed to find a subgraph pattern in the ONNX model computation graph that matches the dedicated composite instructions supported by the target AI chip. Replace the subgraph pattern with a single custom operator and mark the dedicated composite instruction corresponding to the custom operator.

4. The hardware-aware ONNX graph optimization method according to claim 1, characterized in that, The performance cost model based on each operator calculates the execution performance cost of multiple operators decomposed using each candidate decomposition scheme on the target AI chip, including: Based on the performance cost model of each operator, the computational cost and memory cost of each candidate decomposition scheme on the target AI chip hardware are calculated. Based on the computation cost and the memory cost, the total cost is calculated by weighted summation and used as the execution performance cost. Based on the total cost of each candidate decomposition scheme, the candidate decomposition scheme with the minimum total cost is selected as the optimal candidate decomposition scheme.

5. The hardware-aware ONNX graph optimization method according to claim 4, characterized in that, Based on the computational cost and the memory cost, the total cost is calculated by weighted summation, including: Cost(op) = α × Comp(op) + β ×Memory(op) Where Cost(op) is the total cost, Comp(op) is the estimated computation cost, Memory(op) is the estimated memory access cost, and α and β are weighting coefficients calibrated for the target AI chip hardware. They are adjusted according to the hardware characteristics of the target AI chip. On computationally intensive chips, the value of α is larger; on chips with limited memory bandwidth, the value of β is larger.

6. The hardware-aware ONNX graph optimization method according to claim 2, characterized in that, Based on the data layout and data type supported by each native operator in the ONNX model computation graph included in the native operator list on the target AI chip hardware, the unified optimizer intelligently inserts or deletes format conversion nodes in the ONNX model computation graph to convert the data format of the native operators preceding its own position, so that the format of the converted native operators is supported by the target AI chip hardware.

7. The hardware-aware ONNX graph optimization method according to claim 1, characterized in that, The process of running micro-benchmark tests on each version of the kernel code to obtain real performance data for each version includes: In a lightweight hardware simulator or real device environment, micro-benchmark tests are run on each version of the kernel code to obtain real performance data for each candidate kernel, including the number of clock cycles.

8. The hardware-aware ONNX graph optimization method according to claim 1, characterized in that, The multi-round iterative optimization of the ONNX model computation graph includes: The ONNX model computation graph is iteratively optimized multiple times based on the unified optimizer until the maximum number of iterations is reached or the optimized ONNX model computation graph no longer changes.

9. A hardware-aware graph optimization and compilation platform for optimizing ONNX graphs, characterized in that, The hardware-aware graph optimization and compilation platform integrates a hardware-aware graph optimization and compilation engine, which is configured with a hardware information base, a unified optimizer, and a just-in-time compiler and performance analyzer. The unified optimizer receives the input ONNX model and the target AI chip identifier. Based on the target AI chip identifier, it obtains the hardware information library of the target AI chip. Based on the hardware information library, it performs pattern matching on the ONNX model computation graph to identify the subgraph corresponding to the dedicated composite instructions of the AI ​​chip, and replaces the subgraph with the corresponding fusion operator. For complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, it generates multiple candidate decomposition schemes. Based on the performance cost model of each operator in the hardware information library, it selects the optimal decomposition scheme from the multiple candidate decomposition schemes. Based on the optimal decomposition scheme, it decomposes the complex operators that are not natively supported by the target AI chip to obtain the optimized ONNX model computation graph. The just-in-time compiler and performance analyzer is used to generate multiple versions of kernel code for operators or subgraphs with multiple implementations in the optimized ONNX model computation graph; to run micro-benchmark tests on each version of kernel code to obtain the real performance data of each version of kernel code; and to feed the real performance data of each version of kernel code back to the hardware information database for updating and calibrating the performance cost model of the corresponding operator. The hardware-aware graph optimization and compilation engine is used to perform multiple rounds of iterative optimization on the ONNX model computation graph based on the unified optimizer and the just-in-time compilation and performance analyzer, and compile the final optimized ONNX model computation graph into an executable file or instruction stream on the target AI chip. For complex operators in the ONNX model computation graph that are not natively supported by the target AI chip, multiple candidate decomposition schemes are generated. Based on the performance-cost model of each operator in the hardware information database, the optimal decomposition scheme is selected from the multiple candidate decomposition schemes, including: Traverse the ONNX model computation graph. When encountering a complex operator, query the hardware information database. When the complex operator is not natively supported by the target AI chip, decompose the complex operator. There are multiple candidate decomposition schemes when decomposing the complex operator. The step of selecting the optimal decomposition scheme from multiple candidate decomposition schemes includes: Based on the performance cost model of each operator, the execution performance cost of multiple operators after decomposition using each candidate decomposition scheme is calculated on the target AI chip. Based on the execution performance cost of each candidate decomposition scheme, the candidate decomposition scheme with the lowest performance cost is selected as the optimal decomposition scheme.

Images