A neural network heterogeneous partition rewriting method and system based on graph intermediate representation

Abstract

This invention relates to the field of artificial intelligence compiler technology, and discloses a method and system for rewriting heterogeneous partitioning of neural networks based on graph intermediate representation. The method includes: obtaining the intermediate representation of the high-level computational graph of the neural network, the hardware configuration of the heterogeneous hardware platform, and the description of the heterogeneous backend devices; extracting subgraph nodes and tensor edges of core operators from the intermediate representation of the high-level computational graph; constructing a cost model based on the hardware configuration; solving the heterogeneous backend device allocation result for each subgraph node; rewriting the allocation result back to the intermediate representation of the high-level computational graph in the form of attributes; generating a wrapper function marked with heterogeneous backend devices; and rewriting the computational graph structure inside the target forward function. This invention realizes hardware-aware graph-level automatic partitioning, tightly coupling the partitioning result with the high-level IR, and connecting to backend code generation through a unified wrapper function interface, solving the problems of existing technologies that rely on static rules, lack quantitative cost models, and have loose coupling between partitioning results and backend calls.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence compiler technology, specifically to a method and system for rewriting heterogeneous partitions of neural networks based on graph intermediate representation. Background Technology

[0002] With the widespread application of deep learning technology in fields such as image recognition, speech processing, and natural language processing, the number of layers and the scale of parameters in neural network models are constantly increasing. Their computation process requires a large number of floating-point operations such as matrix multiplication and convolution, resulting in significant demands on computing power and bandwidth. Especially in resource-constrained scenarios such as edge devices, embedded terminals, and IoT devices, how to complete large-scale neural network inference under power-constrained conditions has become one of the key issues of concern in industry and academia.

[0003] To meet these computing power demands, chip architectures have evolved from the traditional "single general-purpose processor" to a heterogeneous architecture of "general-purpose processor + dedicated accelerator." Combinations such as "CPU+GPU," "CPU+NPU," and "CPU+DSP," widely used in mobile and data center scenarios, achieve a trade-off between overall performance and power consumption by offloading computationally intensive matrix and convolution operations to dedicated accelerators, while retaining control logic and lightweight operators on the general-purpose CPU. In the open-source instruction set architecture RISC-V field, due to its open instruction set and strong scalability, an increasing number of system-on-chips (SoCs) integrate vector extension units, matrix multiply-accumulate arrays, convolution accelerators, and self-developed neural network IPs on the basis of RISC-V processors, delegating a large number of core operators to on-chip dedicated accelerators or vector units for execution.

[0004] In terms of the software stack, to reuse upper-layer models and perform unified optimization across different hardware platforms, deep learning compilation frameworks centered on graph intermediate representations (Graph IR) and operator-level intermediate representations have gradually emerged. These frameworks sit between the training framework and the underlying hardware drivers, mapping the same model to multiple hardware backends through a series of transformations, including operator fusion, memory planning, scheduling optimization, and code generation. In recent years, next-generation high-level computational graph intermediate representations (such as TVM Relax IR) have been widely used to describe the forward computation and parameter layout of neural networks, providing a unified entry point for subsequent graph-level optimization and heterogeneous scheduling.

[0005] The prior art related to this invention is mainly divided into two categories: Existing technology 1: Graph-level heterogeneous distribution / external code generation mechanism based on a general deep learning compilation framework. This approach converts the neural network model into a graph intermediate representation, matches subgraphs using predefined operator patterns, adds backend labels, and calls an external code generation module. The remaining operators are still executed by the CPU backend. Subgraph partitioning and allocation decisions in this type of approach typically rely on static rules such as "whether the external backend supports a certain operator" or "whether it matches a predefined pattern," lacking quantitative cost assessments based on specific hardware parameters (such as CPU floating-point performance, accelerator computing unit size, on-chip storage capacity, DMA bandwidth, etc.). This makes it impossible to accurately estimate the computation time of the subgraph on different backends, cross-device data transfer overhead, and storage resource constraints. Furthermore, this approach has limited support for next-generation high-level IR (such as Relax), and the partitioning results are loosely coupled with the backend calling interface.

[0006] Existing technology two: Manual partitioning and explicit invocation of acceleration libraries. This approach involves engineers manually replacing core operators with acceleration library functions provided by hardware vendors in the business code based on experience, or encapsulating accelerator calls through a custom operator mechanism and embedding them into the model graph. This type of solution heavily relies on manual experience, requiring repeated modifications to the business code for different models and hardware platforms, resulting in high development and maintenance costs. Furthermore, manual partitioning only focuses on local operator acceleration, lacking a comprehensive optimization of computational costs, communication overhead, and storage resources from a global graph perspective, easily leading to situations where "local deployment results in limited or even decreased overall performance improvement." Simultaneously, the partitioning logic is implicit in the business code, making it difficult to reuse and migrate, and failing to form a unified IR layer interface.

[0007] In summary, existing technologies have the following technical shortcomings when dealing with heterogeneous hardware platforms consisting of general-purpose processors and dedicated accelerators: First, there is a lack of graph-level automatic partitioning mechanisms tailored to specific hardware parameters. Existing solutions rely on static rules or predefined patterns, which cannot quantitatively evaluate the execution cost and cross-device communication cost of operator subgraphs on different backends based on specific hardware parameters such as CPU performance, accelerator computing power, on-chip storage capacity, and DMA bandwidth. Consequently, it is difficult to provide a near-globally optimal graph-level partitioning scheme while meeting resource constraints.

[0008] Second, there is a lack of a unified "subgraph + tensor edge" abstraction applicable to cost models. Existing technologies mostly use "operator sets" as the unit of subgraphs, lacking a unified and systematic data structure description of the shape parameters, data types, layout of the subgraph itself, and tensor dependencies between subgraphs. This directly affects the accurate estimation of computational load, memory usage, and communication overhead by cost models, limiting the accuracy of graph-level performance analysis and partition optimization.

[0009] Third, there is a lack of a reusable accelerator call interface that is tightly coupled with the partitioning results. In existing technologies, graph-level partitioning results are separated from the code that actually calls the accelerator: the former only adds an "external execution" mark to the subgraph, while the latter requires manually writing the call logic in the business code or underlying implementation. This loose coupling makes it difficult for backend code generation and heterogeneous runtime to directly identify and uniformly process the subgraphs that should be sent to the accelerator from the IR, and it is also not conducive to reusing partitioning decisions and call logic across different models and hardware platforms. Summary of the Invention

[0010] To address the aforementioned shortcomings in existing technologies, this invention provides a method and system for rewriting heterogeneous partitions of neural networks based on graph intermediate representation. This method solves the problems of existing technologies, such as reliance on static rules, lack of quantitative cost models, and loose coupling between partitioning results and backend calls. It achieves automatic graph-level partitioning at the high-level computational graph intermediate representation level by combining specific hardware parameters, and tightly binds the partitioning results to the intermediate neural network heterogeneous partition rewriting in a unified interface.

[0011] To achieve the above-mentioned objectives, the technical solution adopted by this invention is as follows: A neural network heterogeneous partitioning rewriting method based on graph intermediate representation includes the following steps: S1. Obtain the intermediate representation of the high-level computation graph of the neural network, the hardware configuration of the heterogeneous hardware platform and the description of the heterogeneous back-end devices, and extract the subgraph nodes of the core operators and the tensor edges between the subgraph nodes from the intermediate representation of the high-level computation graph. S2. Based on the hardware configuration of the heterogeneous hardware platform, construct a cost model to quantitatively evaluate the execution cost of subgraph nodes on each heterogeneous backend device, and solve the heterogeneous backend device allocation result for each subgraph node. S3. Rewrite the heterogeneous backend device allocation results back to the intermediate representation of the high-level computation graph in the form of attributes, generate an encapsulated function with heterogeneous backend device tags, and rewrite the computation graph structure inside the target forward function of the intermediate representation of the high-level computation graph to generate a rewritten intermediate representation of the high-level computation graph that can be executed by heterogeneous backend devices.

[0012] Furthermore, step S1 specifically includes: S11. Obtain the intermediate representation of the high-level computation graph of the neural network to be compiled, and the intermediate representation of the high-level computation graph is an abstract intermediate representation layer containing data flow graph structure and type information. S12. Obtain the hardware configuration and heterogeneous back-end device description of the heterogeneous hardware platform, and the hardware configuration and heterogeneous back-end device description include the computing power parameters, storage resource parameters and communication parameters of the interconnection links between the heterogeneous back-end devices. S13. Traverse the target forward function of the intermediate representation of the high-level computation graph, perform pattern matching on the definition expression of the intermediate variable based on the predefined core operator pattern, identify the graph nodes that satisfy the core operator and mark them as subgraph nodes. S14. Assign a unique identifier to each subgraph node, extract the shape parameters, data type, and memory layout information of its input and output tensors, calculate the number of floating-point operations, and form the subgraph node information structure. S15. Based on the data dependencies between subgraph nodes, when the output tensor of one subgraph node is directly referenced as the input tensor of another subgraph node, a tensor edge is constructed between the two subgraph nodes. S16. Record the source node identifier, target node identifier, and data size of each tensor edge to form a tensor edge information structure; S17. Based on the subgraph node information structure and tensor edge information structure, the target forward function is abstracted into a unified graph structure representation composed of the subgraph node set and the tensor edge set.

[0013] Furthermore, step S2 specifically includes: S21. Based on the hardware configuration of the heterogeneous hardware platform, construct a cost model, which includes a computation cost sub-model, a communication cost sub-model, and a storage constraint sub-model. S22. For each combination of a subgraph node and each heterogeneous backend device, based on the number of floating-point operations of the subgraph node, the effective floating-point performance of the heterogeneous backend device, and the hardware efficiency correction coefficient, a computational cost sub-model is used to generate an estimated computational execution time of the subgraph node on the heterogeneous backend device. S23. For each tensor edge and the two subgraph nodes connected to it, different combinations of heterogeneous back-end devices may be allocated. Based on the data volume of the tensor edge, the link bandwidth between the source heterogeneous back-end device and the target heterogeneous back-end device and the startup delay, a communication cost sub-model is used to generate an estimated value of data transmission time across heterogeneous back-end devices. S24. Based on the on-chip storage capacity, memory alignment requirements, and memory usage of the active tensor set during the execution of subgraph nodes of heterogeneous back-end devices, a storage constraint sub-model is adopted to generate storage feasibility constraints for each heterogeneous back-end device at any execution time. S25. Input the estimated values ​​of the unified graph structure representation and cost model into the partition solver, and establish an integer programming problem with the optimization objective of minimizing the estimated computation execution time and the estimated data transfer time, under the constraint of storage feasibility. S26. Solve the integer programming problem to generate a heterogeneous back-end device allocation result that corresponds one-to-one between each subgraph node and the target heterogeneous back-end device, and the heterogeneous back-end device allocation result is represented as a subgraph node-heterogeneous back-end device allocation mapping table.

[0014] Furthermore, step S3 specifically includes: S31. Serialize the allocation results of heterogeneous back-end devices and write them into the intermediate representation of the high-level computation graph in the form of function attributes and module attributes; and the function attributes are used to record the complete mapping relationship between the subgraph nodes contained in the target forward function and the heterogeneous back-end device identifiers, and the module attributes are used to record the operator type, shape parameters and upstream and downstream dependencies of each subgraph node. S32. Traverse the target forward function. For subgraph nodes in the heterogeneous backend device allocation result that are allocated to non-default heterogeneous backend devices, execute the encapsulation function generation operation: S33. Create an independent wrapper function for the current subgraph node, and the parameter list and return type of the wrapper function are completely consistent with the original operator corresponding to the subgraph node. S34. Preserve the computational semantics of the original operators within the function body of the encapsulated function; S35. Add a target heterogeneous backend device tag field to the function properties of the encapsulated function. This target heterogeneous backend device tag field is used to uniquely identify that the encapsulated function needs to be generated or scheduled at runtime by a specified heterogeneous backend device. S36. Check whether the memory layout of the input tensor of the current subgraph node is consistent with the hardware preference layout of the target heterogeneous backend device. If they are inconsistent, automatically insert a layout conversion operator at the boundary of the encapsulation function and execute step S37. Otherwise, directly execute step S37. S37. Inside the computation body of the target forward function, replace the original expression that calls the operator corresponding to the subgraph node with the expression that calls the encapsulated function. S38. After the traversal is completed, output the rewritten intermediate representation of the high-level graph; and the rewritten intermediate representation of the high-level graph contains the semantics of the original computation graph, complete heterogeneous partitioning decision information, and a unified format backend call interface, so as to use it as input for code generation and heterogeneous runtime calls.

[0015] A system applying a neural network heterogeneous partitioning rewriting method based on graph intermediate representation includes: Hardware Configuration and Cost Model Module: Used to store the hardware configuration of heterogeneous hardware platforms and descriptions of heterogeneous backend devices, and provides cost estimation interfaces for computation execution time estimation, data transmission time estimation, and storage feasibility constraints required by the cost model; Subgraph extraction module: used to traverse the target forward function of the intermediate representation of the high-level computation graph, identify and mark subgraph nodes based on the core operator pattern, extract subgraph node information and tensor edge information, and construct a unified graph structure representation composed of the subgraph node set and the tensor edge set; Partition Solving Module: Based on the cost model and unified graph structure representation, this module solves integer programming problems with the optimization objective of minimizing computation execution time and data transmission time, and with backend storage capacity as a constraint. It generates an allocation mapping table that corresponds one-to-one between each subgraph node and heterogeneous backend devices. Attribute annotation module: used to serialize the allocation mapping table and write it into the intermediate representation of the high-level computation graph in the form of function attributes and module attributes; The Encapsulation Function Generation and Graph Rewriting module is used to create encapsulation functions for subgraph nodes assigned to non-default heterogeneous backend devices, add target heterogeneous backend device tags to the function properties, check and automatically insert layout transformation operators, replace the original operator calls in the target forward function with calls to the encapsulation function, and output the rewritten intermediate representation of the high-level graph.

[0016] The present invention has the following beneficial effects: This invention proposes a neural network heterogeneous partitioning rewriting method and system based on graph intermediate representation. It introduces a hardware parameter-aware cost model at the high-level computational graph intermediate representation layer, elevating subgraph partitioning and heterogeneous backend device allocation from empirical manual specification to automated, measurable quantitative decision-making. The partitioning results are bound to the intermediate representation using a unified data structure, achieving stable compilation flow and reusable encapsulated function rewriting across platforms. Simultaneously, while ensuring type consistency and compilation pipeline integrity, this method significantly reduces end-to-end inference latency on heterogeneous backend devices, effectively controls cross-heterogeneous backend device transmission overhead, and supports rapid adaptation of a single graph-level process to multiple heterogeneous hardware platforms, greatly reducing system integration and maintenance costs. Attached Figure Description

[0017] Figure 1 This is a flowchart illustrating a neural network heterogeneous partitioning rewriting method based on graph center representation proposed in this invention. Figure 2 This is a schematic diagram of the overall system architecture to which the method proposed in this invention applies in the embodiments; Figure 3 This is a schematic diagram illustrating the process of traversing the target forward function of the neural network in the intermediate representation of the high-level computation graph (e.g., Relax IR) in the embodiment. Figure 4 This is a schematic diagram illustrating the heterogeneous backend configuration and cost model modeling process in the embodiment. Figure 5 This is a schematic diagram of the process of generating the encapsulation function and rewriting the target forward function based on the allocation result of heterogeneous backend devices in the embodiment. Figure 6 This is a schematic diagram of the structure of a neural network heterogeneous partitioning rewriting system based on graph center representation proposed in this invention. Detailed Implementation

[0018] The specific embodiments of the present invention are described below to enable those skilled in the art to understand the present invention. However, it should be understood that the present invention is not limited to the scope of the specific embodiments. For those skilled in the art, various changes are obvious as long as they are within the spirit and scope of the present invention as defined and determined by the appended claims. All inventions utilizing the concept of the present invention are protected.

[0019] like Figure 1 As shown, a neural network heterogeneous partitioning rewriting method based on graph intermediate representation includes the following steps S1-S3: S1. Obtain the intermediate representation of the high-level computation graph of the neural network, the hardware configuration of the heterogeneous hardware platform and the description of the heterogeneous back-end devices, and extract the subgraph nodes of the core operators and the tensor edges between the subgraph nodes from the intermediate representation of the high-level computation graph.

[0020] In this embodiment, Figure 2 The overall system architecture to which the method proposed in this invention is applicable is illustrated, including a heterogeneous hardware platform and a deep learning compilation and execution software stack: (1) The heterogeneous hardware platform includes at least one general-purpose processor (e.g., RISC-V CPU, Arm CPU, x86 CPU, etc.) and one or more dedicated acceleration back-ends. These acceleration back-ends are collectively referred to as heterogeneous back-end devices, including but not limited to GPUs, NPUs, DSPs, accelerators (accelerators based on RISC-V vector extensions or self-developed matrix / convolution IPs), and reconfigurable computing units (FPGAs or other reconfigurable computing units). The general-purpose processor and the heterogeneous back-end devices are connected through on-chip interconnects, system buses, or high-speed interfaces (e.g., AXI, PCIe), and have certain bandwidth and latency characteristics, enabling the transfer and sharing of tensor data.

[0021] (2) The deep learning compilation and execution software stack includes a front-end, a compilation core, and heterogeneous back-end device code generation (back-end Codegen). The front-end receives models from training frameworks (such as PyTorch and TensorFlow) or directly constructs neural networks using model description languages. The compilation core uses a high-level computation graph intermediate representation (e.g., Relax IR) to represent the forward computation graph of the neural network (i.e., the computation graph structure inside the target forward function), and includes the graph-level partitioning module and intermediate representation (IR) rewriting module proposed in this invention. Heterogeneous back-end device code generation (back-end Codegen) generates corresponding kernel code and calling interfaces for different heterogeneous back-end devices, such as generating TIR, CUDA kernels, NPU instruction sequences, and calling external C libraries. Furthermore, the heterogeneous runtime can schedule the execution of the CPU and each heterogeneous back-end device based on the partitioning results, and manage the allocation and transmission of tensors between different heterogeneous back-end devices.

[0022] Based on the above-mentioned overall system architecture, the core work of the method proposed in this invention occurs at the layer of intermediate representation of the high-level computation graph, graph-level partitioning, and intermediate representation rewriting, which corresponds to the compilation core. By completing subgraph extraction, cost model analysis, partitioning decision, partitioning result attribute annotation, and wrapper construction on the intermediate representation of the high-level computation graph, a unified interface is provided for heterogeneous backend device code generation (backend Codegen) and runtime.

[0023] In summary, the specific operations for obtaining the intermediate representation of the high-level computation graph of the neural network, the hardware configuration of the heterogeneous hardware platform, and the description of the heterogeneous backend devices, and extracting the subgraph nodes of the core operators and the tensor edges between the subgraph nodes from the intermediate representation of the high-level computation graph are as follows: Obtain the intermediate representation of the high-level computation graph of the neural network to be compiled, and the intermediate representation of the high-level computation graph is an abstraction of the intermediate representation layer containing data flow graph structure and type information.

[0024] Obtain the hardware configuration and heterogeneous backend device description of the heterogeneous hardware platform, and the hardware configuration and heterogeneous backend device description include the computing power parameters, storage resource parameters and communication parameters of the interconnection links between the heterogeneous backend devices.

[0025] Traverse the target forward function of the intermediate representation of the high-level computation graph, perform pattern matching on the definition expression of the intermediate variable based on the predefined core operator pattern, identify the graph nodes that satisfy the core operator and mark them as subgraph nodes.

[0026] In this step, Figure 3This invention demonstrates the process of traversing the target forward function of a neural network in a high-level computational graph intermediate representation (e.g., Relax IR), thereby automatically constructing a graph structure of "subgraph + tensor edge" suitable for the cost model.

[0027] The acquisition and traversal of the target forward function are as follows: select the target forward function from the intermediate representation rewriting module, traverse its computation graph structure (e.g., data flow blocks or basic blocks); and obtain the definition expression and usage relationship of each intermediate variable by parsing the data dependency relationship of the nodes in the computation graph.

[0028] The core operator identification and subgraph node construction process involves taking the definition expression and usage relationship of the intermediate variables of the target forward function in the intermediate representation of the high-level computation graph as input. First, a set of core operator patterns are predefined. Each core operator pattern is used to characterize the type or combination of operators that need to be preferentially mapped to specific heterogeneous backend devices, such as: matrix multiplication operators: relax.matmul, fully connected layers, etc.; convolution operators: relax.nn.conv2d and its depthwise and pointwise variants; other computationally intensive operators: such as parts of certain large-scale normalization or attention modules that can be abstracted into standard operators.

[0029] Then, the definition expressions of intermediate variables and their usage relationships are traversed, the operators corresponding to the intermediate variables are pattern matched, the intermediate variables that satisfy the core operator pattern are identified, and the graph nodes corresponding to the intermediate variables are marked as subgraph nodes.

[0030] Furthermore, during the aforementioned traversal process, when encountering a calling expression that satisfies the core operator pattern described above, it is necessary to assign a unique identifier (node_id) to each subgraph node and extract its input and output tensor shape parameters (static dimension information can be obtained through struct_info or type deduction), data type, and memory layout information. The number of floating-point operations is calculated (recording the operator type: matrix multiplication, standard convolution, depthwise convolution, pointwise convolution, etc., and calculating the matrix dimension, such as M, K, N, or the convolution dimension, such as N, C_in, C_out, H, W, K_h, K_w, stride_h, stride_w, groups), forming the subgraph node information structure, thereby storing the above information uniformly.

[0031] Therefore, through the aforementioned core operator identification and subgraph node construction process, a subgraph composed of core operators can be obtained on the intermediate representation of the high-level computation graph, providing a foundation for subsequent heterogeneous partitioning and subgraph rewriting. Furthermore, after obtaining each subgraph node, the subgraph nodes can be connected according to the data dependencies between them to obtain tensor edges, thereby constructing the corresponding subgraph structure. The process of tensor edge extraction and the corresponding subgraph structure is as follows: Based on the data dependencies between subgraph nodes, when the output tensor of one subgraph node is directly referenced as the input tensor of another subgraph node, a tensor edge is constructed between the two subgraph nodes. For each tensor edge, the source node identifier, the target node identifier, and the data size of the tensor edge are recorded, forming a tensor edge information structure. Based on the subgraph node information structure and the tensor edge information structure, the target forward function is abstracted into a unified graph structure representation composed of the subgraph node set and the tensor edge set. This graph structure serves as the direct input for subsequent cost models and partitioned solutions.

[0032] In an optional embodiment of the present invention, by establishing a mapping from Var to node_id, the subgraph node from which each intermediate tensor is generated is identified, where Var is an intermediate variable in the intermediate representation of the high-level computation graph, used to characterize the output tensor generated by the operator; when the defining node of a Var is different from the node it uses, the source node and target node of the tensor edge are determined accordingly; when the output of a core operator is used as the input of another core operator, it can be determined that there is a tensor edge from the source node to the target node; by using the shape and data type of the output tensor, the data size, element type, layout (such as NCHW, NHWC, or row principal / column principal of a matrix), etc., corresponding to the edge are calculated and encapsulated in a TensorEdgeInfo structure. Finally, the present invention abstracts the entire target forward function as: a set of subgraph nodes. and tensor edge set , means as follows: in, , , They are the 0th, 1st, and 1st respectively. Subgraph nodes , , Article 0, Article 1, and Article 2 are respectively. Each edge contains a source node, a target node, and tensor features.

[0033] S2. Based on the hardware configuration of the heterogeneous hardware platform, construct a cost model to quantitatively evaluate the execution cost of subgraph nodes on each heterogeneous backend device, and solve the heterogeneous backend device allocation result for each subgraph node.

[0034] In this embodiment, the step is to construct a cost model and use the cost model to establish a partitioning problem for each heterogeneous backend device. By solving the partitioning problem, the heterogeneous backend device allocation result for each subgraph node is generated, which is the graph partitioning result.

[0035] in, Figure 4 This paper demonstrates the process of heterogeneous backend configuration and cost modeling, which supports multiple heterogeneous backend devices. A hardware configuration and cost model parameter table is built during the compilation phase and used as input for subsequent partitioning methods. The heterogeneous backend configuration modeling is as follows: For each available heterogeneous backend device, a heterogeneous backend device description structure (i.e., backend description structure) is defined. This structure not only includes basic floating-point operation count (FLOPs) information but also introduces finer-grained micro-architectural parameters, including: `backend_id`: a backend device identifier that uniquely identifies a specific device or device category; `backend_type`: the backend device type, such as CPU, GPU, NPU, or a self-developed ACC; computational capability parameters: peak FLOPS, SIMD width, and systolic array dimensions of dedicated accelerators; storage and caching parameters: on-chip available storage capacity, cache size, supported alignment methods, and supported data types; communication parameters: bandwidth with other devices, single-boot latency, and supported DMA modes; and actual efficiency parameters: empirical coefficients used to correct the difference between theoretical values ​​and actual performance, such as `acc_efficiency` and `cpu_efficiency`.

[0036] In one optional embodiment of the present invention, a hardware configuration object may also be defined, which contains a backend list, namely a heterogeneous backend device description structure backends[], where each element is a heterogeneous backend device description structure; at the same time, an "interconnection matrix" may also be defined to describe the data transmission bandwidth and latency between any two heterogeneous backend devices, which is used to estimate the cross-device tensor transfer overhead.

[0037] The cost model is modeled as follows: Based on the hardware configuration and description of heterogeneous backend devices on a heterogeneous hardware platform, a cost model is constructed, including a computational cost sub-model, a communication cost sub-model, and a storage constraint sub-model, thereby establishing a cost estimation function for each subgraph node of the core operator and each tensor edge: (1) Computational Cost Sub-model: For each combination of a subgraph node and each heterogeneous backend device, based on the number of floating-point operations of the subgraph node, the effective floating-point performance of the heterogeneous backend device, and the hardware efficiency correction coefficient, a computational cost sub-model is used to generate an estimated computational execution time of the subgraph node on the heterogeneous backend device, expressed as: in, Subgraph nodes In heterogeneous backend devices The estimated execution time of the calculation is as follows. Subgraph nodes The number of floating-point operations corresponds to the number of operations performed by the operator. This number of floating-point operations is determined by a preset calculation formula based on the operator type and the shape parameters of its input and output tensors (such as batch size, number of channels, height, width, convolution kernel size, etc.). (The calculation methods vary greatly between different operators. The number of floating-point operations performed by an operator is a relatively basic concept in the field, and those skilled in the art can easily calculate or design it.) For heterogeneous backend devices Effective floating-point arithmetic performance, Subgraph nodes In heterogeneous backend devices Above, hardware efficiency correction coefficients for specific operator types (such as Depthwise convolution vs. ordinary convolution).

[0038] (2) Storage Constraint Sub-model: Based on the on-chip storage capacity of heterogeneous back-end devices, memory alignment requirements, and memory usage of the active tensor set during the execution of subgraph nodes, a storage constraint sub-model is adopted to generate storage feasibility constraints for each heterogeneous back-end device at any execution time, expressed as: in, For at any time An intermediate tensor index set that is still in a "live" state is a set of intermediate tensors whose corresponding production operators have been completed, but at least one consumption operator has not yet been completed. A production operator is a subgraph node that generates an intermediate tensor as output, and a consumption operator is a subgraph node that takes an intermediate tensor as input. When an intermediate tensor corresponds to multiple consumption operators, the intermediate tensor is still in a live state as long as at least one consumption operator has not been completed. To perform the round-up operation, simulate hardware memory alignment. Expenses; intermediate tensor Storage footprint; Memory alignment granularity (in bytes). As a buffer scaling factor, for example, for the back end of a pipeline that supports Ping-Pong buffers, a 2x output buffer needs to be reserved; The size of the temporary workspace required for operator execution; For heterogeneous backend devices The upper limit of storage capacity.

[0039] (3) Communication Cost Sub-model: For each tensor edge and the two subgraph nodes it connects, which may be assigned to different heterogeneous back-end device combinations, based on the data size of the tensor edge, the link bandwidth between the source heterogeneous back-end device and the target heterogeneous back-end device, and the startup delay, a communication cost sub-model is used to generate an estimated value of the data transmission time across heterogeneous back-end devices, expressed as: in, Subgraph nodes With subgraph nodes Distributed on heterogeneous back-end devices heterogeneous back-end devices hour( The estimated data transmission time; The concurrency coefficient characterizes whether the hardware supports pipeline overlap between computation and data transfer. ; Source equipment (heterogeneous back-end equipment) From the target device (heterogeneous backend device) Direct memory access (DMA) startup latency; The amount of data in the tensor to be transmitted; Source equipment (heterogeneous back-end equipment) ) and target device (heterogeneous back-end device) The physical link bandwidth between ( ).

[0040] The partitioning problem is modeled as follows: After obtaining the subgraph nodes, tensor edges, hardware configuration, and cost model, graph-level partitioning is performed to determine the heterogeneous backend devices corresponding to each subgraph node. Therefore, graph-level partitioning is modeled as a partitioning problem (integer programming problem), represented as: Decision variables: for each subgraph node and each heterogeneous backend device Define a place variable ,Right now: And meet the requirements ; Objective function: The overall objective considering computation and communication time. for: in, This is an operation to retrieve the minimum value; Subgraph nodes and heterogeneous back-end devices Placement variables; Subgraph nodes and heterogeneous back-end devices Placement variables; Subgraph nodes With subgraph nodes Distributed on heterogeneous back-end devices heterogeneous back-end devices hour( The estimated data transmission time; Constraints: For each heterogeneous backend device Its on-chip storage usage at any given time does not exceed For cases where pipelined or batch processing needs to be considered, memory constraints and parallelism limits on time slices are further added.

[0041] The graph-level partitioning result can then be obtained by solving this optimization problem using the partitioning solution method.

[0042] The selection of partitioning solution methods may include: a greedy strategy: minimizing the overall objective function. With the objective of iterating the subgraph node migration benefits, the partitioning results are gradually optimized while satisfying storage constraints. Global solution: The Integer Linear Programming (ILP) solver is used to directly solve the above objective function and constraints globally. Combination strategy: An approximate solution is first obtained through a greedy method, and then ILP or other search methods are used for local refinement.

[0043] In an optional embodiment of the present invention, the above-mentioned graph-level partitioning optimization problem can be solved using a heuristic partitioning algorithm based on a combination of greedy expansion and local search, as follows: First, using the aforementioned cost model, we estimate parameters such as the computation time of each subgraph node in the forward graph on different heterogeneous back-end devices, the communication cost between adjacent nodes across heterogeneous back-end devices, the on-chip storage (e.g., BRAM) usage when executing on an accelerator, and the pipeline overlap benefits. During the initialization phase, all subgraph nodes are uniformly assigned to CPU devices for execution, serving as the initial partitioning scheme that satisfies the constraints. Then, the first phase of greedy expansion begins: in each iteration, from the nodes currently on the CPU and executable on the accelerator, the impact of migrating each node to the accelerator on "total execution time = computation cost + communication cost" is estimated, while simultaneously checking whether the on-chip storage capacity constraint of the accelerator is still satisfied after migration. If migrating a node can bring the maximum positive benefit without violating hardware constraints, then the node is migrated to the accelerator and the current partition is updated. When no feasible node with a positive benefit exists, the greedy phase ends.

[0044] Building upon the greedy partitioning approach, a second stage of local search optimization can be performed: In the current scheme, a pair of candidate subgraph nodes are selected from the CPU side and the accelerator side respectively, and their deployment devices are swapped to re-evaluate the total execution time under resource constraints. If the objective function is improved after the swap, the swap is accepted and the current scheme is updated. This process is repeated until no pair of node swaps can further reduce the total execution time. Through this two-stage combination, this embodiment can quickly obtain near-optimal graph-level partitioning results while ensuring hardware capabilities and on-chip storage constraints.

[0045] In an optional embodiment of the present invention, the method proposed in this invention can also support multiple heterogeneous backend devices: the above modeling method allows the backend set to include multiple non-CPU devices, such as NPU, GPU, and self-developed matrix accelerators coexisting. In the cost model... and Each pair (subgraph node, heterogeneous backend device) is defined, so this partitioning method can naturally select the most suitable combination of backends. For example, large matrix multiplications can be assigned to matrix accelerators; convolutional subgraphs can be assigned to the NPU; and certain operations requiring high-bandwidth memory access can be reserved for the CPU or GPU. After the partitioning is solved, each subgraph node can be obtained. The allocation results of heterogeneous backend devices, for example ,in It serves as a unique identifier (node_id) for subgraph nodes. This is the identifier for the heterogeneous backend device corresponding to subgraph node i. This represents the allocation mapping table of "subgraph node ID → backend ID" output by the partition solver.

[0046] S3. Rewrite the heterogeneous backend device allocation results back to the intermediate representation of the high-level computation graph in the form of attributes, generate an encapsulated function with heterogeneous backend device tags, and rewrite the computation graph structure inside the target forward function of the intermediate representation of the high-level computation graph to generate a rewritten intermediate representation of the high-level computation graph that can be executed by heterogeneous backend devices.

[0047] In this embodiment, this step involves writing the partitioning result, i.e., the heterogeneous backend device allocation result, back to the function attributes and module attributes represented in the middle of the high-level computation graph. Simultaneously, it involves generating an encapsulation function based on the heterogeneous backend device allocation result and rewriting the target forward function. The operation is as follows: The allocation results of heterogeneous backend devices are written back to the middle representation of the high-level computation graph in the form of attributes: Function-level partitioning attribute: On the target forward function, record the complete mapping table from "subgraph node ID to heterogeneous backend device ID" and the list of backend devices actually involved in the function.

[0048] Node-level details: Create a metadata index in the module properties to record the operator type, detailed shape parameters, upstream and downstream connection relationships, and the final assigned backend device identifier for each subgraph node.

[0049] Module-level summary information: Records statistical information such as the total number of subgraph nodes, the estimated computing load distribution of each heterogeneous backend device, and the resource constraint satisfaction status.

[0050] By using the above attribute annotations, this invention ensures that the partitioning results are closely attached to the intermediate representation of the high-level computation graph. Subsequent encapsulation function rewriting, heterogeneous backend device code generation (backend codegen), and runtime can directly read this information from the intermediate representation without relying on external configuration files or business code.

[0051] The process of generating the encapsulation function and rewriting the target forward function based on the allocation results of heterogeneous backend devices is as follows: Figure 5 As shown: Wrapper generation and memory layout consistency: for each core subgraph node assigned to a non-default backend Perform the following operations: (1) Check the memory layout of the input tensor and the target heterogeneous backend device. Check if the hardware layout preferences (e.g., NPU prefers NHWC while CPU prefers NCHW) are consistent. If not, automatically insert layout transformation operators at the boundaries of the wrapper function.

[0052] (2) Create a new wrapper function to ensure that its formal parameters and return types are consistent with the original operator.

[0053] (3) Preserve the original operator semantics in the function body and add explicit "backend code generation mark" and scheduling prompt information to the function properties.

[0054] Rewriting the target forward function: Traversing the subgraph nodes in the target forward function: (1) For nodes assigned to heterogeneous backends: replace the original operator call expression directly with a call to the newly generated Wrapper function; (2) For the default backend node: keep the original expression unchanged.

[0055] After rewriting, the backend compiler only needs to identify the Wrapper function with specific tags to complete the code generation for the specific hardware.

[0056] Multi-heterogeneous backend device wrapper function support: When there are multiple different heterogeneous backend devices, different wrapper functions are allowed to use different code generation codegen tags (e.g., "Codegen":"my_riscv_acc", "Codegen":"npu_v1", "Codegen":"cuda_gpu", etc.), thereby supporting: certain subgraph nodes to be assigned to the self-developed matrix accelerator; certain convolutional subgraphs to be assigned to the NPU; and certain large-scale matrix operations to be assigned to the GPU.

[0057] Since all wrapper functions follow the same intermediate representation abstraction, heterogeneous backend devices can implement different codegen modules or runtime calling strategies for different code generation codegen names, maintaining a unified overall architecture.

[0058] A system applying a neural network heterogeneous partitioning rewriting method based on graph intermediate representation includes: Hardware Configuration and Cost Model Module: This module stores the hardware configuration of heterogeneous hardware platforms and descriptions of heterogeneous backend devices, and provides cost estimation interfaces for computation execution time estimation, data transfer time estimation, and storage feasibility constraints required by the cost model.

[0059] Subgraph extraction module: used to traverse the target forward function of the intermediate representation of the high-level computation graph, identify and mark subgraph nodes based on the core operator pattern, extract subgraph node information and tensor edge information, and construct a unified graph structure representation composed of the subgraph node set and the tensor edge set.

[0060] Partition Solving Module: Based on the cost model and unified graph structure representation, this module solves integer programming problems with the optimization objective of minimizing computation execution time and data transmission time, and with backend storage capacity as a constraint. It generates an allocation mapping table that corresponds one-to-one between each subgraph node and heterogeneous backend devices.

[0061] Attribute annotation module: Used to serialize the allocation mapping table and write it into the intermediate representation of the high-level computation graph in the form of function attributes and module attributes.

[0062] The Encapsulation Function Generation and Graph Rewriting module is used to create encapsulation functions for subgraph nodes assigned to non-default heterogeneous backend devices, add target heterogeneous backend device tags to the function properties, check and automatically insert layout transformation operators, replace the original operator calls in the target forward function with calls to the encapsulation function, and output the rewritten intermediate representation of the high-level graph.

[0063] In summary, the proposed method and system for heterogeneous partitioning and rewriting neural networks based on graph intermediate representation automatically extracts subgraphs and tensor edge information from high-level computational graph intermediate representations (e.g., TVM Relax), constructs a cost model in conjunction with specific hardware configurations, and performs graph-level partitioning and encapsulation function rewriting for multiple heterogeneous backend devices. This transforms the task partitioning of neural networks between general-purpose CPUs and multiple heterogeneous accelerated backends from "empirical manual specification" to "measurable and reproducible automatic decision-making." While ensuring the stability of the compilation process and type consistency, it effectively improves end-to-end inference performance and reduces system integration complexity.

[0064] Meanwhile, taking the deployment of MobileNetV1 (input 224×224, batch=1, INT8 quantization) on a RISC-V SoC (dual-core CPU + on-chip KPU, acceleration unit has high inference performance for lightweight convolutional networks) based on Canaan Technology K230 as an example, without modifying the network structure: in the baseline scheme where all operators are executed by the CPU backend and only general compilation optimization is used, the single-frame inference latency is about 25-30 ms; by adopting the graph-level partitioning and encapsulation function rewriting of the present invention, the computationally intensive subgraphs such as backbone convolution and fully connected layers are automatically partitioned to the KPU or matrix kernel based on vector expansion, while the remaining lightweight operators are retained to be executed on the CPU, the end-to-end single-frame inference latency can be reduced to about 11-14 ms, and the overall throughput is improved by about 2-3 times. At the same time, the overhead of tensor transfer and layout transformation across CPU / KPU is controlled within 15% of the total latency, reserving more computing power and power consumption space for parallel operation of other tasks such as image processing and video encoding and decoding on the same SoC. The above data was obtained under the following test conditions: K230 SoC (CPU+KPU), MobileNetV1 input 224×224, batch=1, INT8 quantization; the baseline was that all operators were executed in the CPU backend and only general compilation optimizations were enabled; the solution of this invention is to partition the computationally intensive subgraphs to the KPU or vector / matrix kernels according to graph-level partitioning, while the remaining operators are retained in the CPU; the inference latency is the average of multiple repeated runs (including / excluding warm-up can be explained), and the timing range is end-to-end single-frame inference.

[0065] Furthermore, since the partitioning results are bound to the intermediate representation of the high-level computation graph with a unified data structure and function attributes, when porting to different hardware platforms (such as RISC-V SoCs with only vector extensions, edge AI chips with integrated KPU / NPU, etc.), only the hardware configuration and the code generation implementation of the corresponding heterogeneous backend devices need to be adjusted to reuse the same graph-level partitioning and rewriting process. There is no need to repeatedly modify the operator calling logic in the upper-level business code, which helps to reduce engineering integration and maintenance costs, and reserves clear interfaces for future expansion of more heterogeneous backends.

[0066] Specific embodiments have been used to illustrate the principles and implementation methods of this invention. The descriptions of the embodiments above are only for the purpose of helping to understand the method and core ideas of this invention. At the same time, for those skilled in the art, there will be changes in the specific implementation methods and application scope based on the ideas of this invention. Therefore, the content of this specification should not be construed as a limitation of this invention.

[0067] Those skilled in the art will recognize that the embodiments described herein are intended to help the reader understand the principles of the invention, and should be understood that the scope of protection of the invention is not limited to such specific statements and embodiments. Those skilled in the art can make various other specific modifications and combinations based on the technical teachings disclosed in this invention without departing from the spirit of the invention, and these modifications and combinations are still within the scope of protection of this invention.

Claims

1. A neural network heterogeneous partitioning rewriting method based on graph intermediate representation, characterized in that, Includes the following steps: S1. Obtain the intermediate representation of the high-level computation graph of the neural network, the hardware configuration of the heterogeneous hardware platform and the description of the heterogeneous back-end devices, and extract the subgraph nodes of the core operators and the tensor edges between the subgraph nodes from the intermediate representation of the high-level computation graph. S2. Based on the hardware configuration of the heterogeneous hardware platform, construct a cost model to quantitatively evaluate the execution cost of subgraph nodes on each heterogeneous backend device, and solve the heterogeneous backend device allocation result for each subgraph node. S3. Rewrite the heterogeneous backend device allocation results back to the intermediate representation of the high-level computation graph in the form of attributes, generate an encapsulated function with heterogeneous backend device tags, and rewrite the computation graph structure inside the target forward function of the intermediate representation of the high-level computation graph to generate a rewritten intermediate representation of the high-level computation graph that can be executed by heterogeneous backend devices.

2. The neural network heterogeneous partitioning rewriting method based on graph intermediate representation according to claim 1, characterized in that, Step S1 specifically includes: S11. Obtain the intermediate representation of the high-level computation graph of the neural network to be compiled, and the intermediate representation of the high-level computation graph is an abstract intermediate representation layer containing data flow graph structure and type information. S12. Obtain the hardware configuration and heterogeneous back-end device description of the heterogeneous hardware platform, and the hardware configuration and heterogeneous back-end device description include the computing power parameters, storage resource parameters and communication parameters of the interconnection links between the heterogeneous back-end devices. S13. Traverse the target forward function of the intermediate representation of the high-level computation graph, perform pattern matching on the definition expression of the intermediate variable based on the predefined core operator pattern, identify the graph nodes that satisfy the core operator and mark them as subgraph nodes. S14. Assign a unique identifier to each subgraph node, extract the shape parameters, data type, and memory layout information of its input and output tensors, calculate the number of floating-point operations, and form the subgraph node information structure. S15. Based on the data dependencies between subgraph nodes, when the output tensor of one subgraph node is directly referenced as the input tensor of another subgraph node, a tensor edge is constructed between the two subgraph nodes. S16. Record the source node identifier, target node identifier, and data size of each tensor edge to form a tensor edge information structure; S17. Based on the subgraph node information structure and tensor edge information structure, the target forward function is abstracted into a unified graph structure representation composed of the subgraph node set and the tensor edge set.

3. The neural network heterogeneous partitioning rewriting method based on graph intermediate representation according to claim 2, characterized in that, Step S2 specifically includes: S21. Based on the hardware configuration of the heterogeneous hardware platform, construct a cost model, which includes a computation cost sub-model, a communication cost sub-model, and a storage constraint sub-model. S22. For each combination of a subgraph node and each heterogeneous backend device, based on the number of floating-point operations of the subgraph node, the effective floating-point performance of the heterogeneous backend device, and the hardware efficiency correction coefficient, a computational cost sub-model is used to generate an estimated computational execution time of the subgraph node on the heterogeneous backend device. S23. For each tensor edge and the two subgraph nodes connected to it, different combinations of heterogeneous back-end devices may be allocated. Based on the data volume of the tensor edge, the link bandwidth between the source heterogeneous back-end device and the target heterogeneous back-end device and the startup delay, a communication cost sub-model is used to generate an estimated value of data transmission time across heterogeneous back-end devices. S24. Based on the on-chip storage capacity, memory alignment requirements, and memory usage of the active tensor set during the execution of subgraph nodes of heterogeneous back-end devices, a storage constraint sub-model is adopted to generate storage feasibility constraints for each heterogeneous back-end device at any execution time. S25. Input the estimated values ​​of the unified graph structure representation and cost model into the partition solver, and establish an integer programming problem with the optimization objective of minimizing the estimated computation execution time and the estimated data transfer time, under the constraint of storage feasibility. S26. Solve the integer programming problem to generate a heterogeneous back-end device allocation result that corresponds one-to-one between each subgraph node and the target heterogeneous back-end device, and the heterogeneous back-end device allocation result is represented as a subgraph node-heterogeneous back-end device allocation mapping table.

4. The neural network heterogeneous partitioning rewriting method based on graph intermediate representation according to claim 3, characterized in that, Step S3 specifically includes: S31. Serialize the allocation results of heterogeneous back-end devices and write them into the intermediate representation of the high-level computation graph in the form of function attributes and module attributes; and the function attributes are used to record the complete mapping relationship between the subgraph nodes contained in the target forward function and the heterogeneous back-end device identifiers, and the module attributes are used to record the operator type, shape parameters and upstream and downstream dependencies of each subgraph node. S32. Traverse the target forward function. For subgraph nodes in the heterogeneous backend device allocation result that are allocated to non-default heterogeneous backend devices, execute the encapsulation function generation operation: S33. Create an independent wrapper function for the current subgraph node, and the parameter list and return type of the wrapper function are completely consistent with the original operator corresponding to the subgraph node. S34. Preserve the computational semantics of the original operators within the function body of the encapsulated function; S35. Add a target heterogeneous backend device tag field to the function properties of the encapsulated function. This target heterogeneous backend device tag field is used to uniquely identify that the encapsulated function needs to be generated or scheduled at runtime by a specified heterogeneous backend device. S36. Check whether the memory layout of the input tensor of the current subgraph node is consistent with the hardware preference layout of the target heterogeneous backend device. If they are inconsistent, automatically insert a layout conversion operator at the boundary of the encapsulation function and execute step S37. Otherwise, directly execute step S37. S37. Inside the computation body of the target forward function, replace the original expression that calls the operator corresponding to the subgraph node with the expression that calls the encapsulated function. S38. After the traversal is completed, output the rewritten intermediate representation of the high-level graph; and the rewritten intermediate representation of the high-level graph contains the semantics of the original computation graph, complete heterogeneous partitioning decision information, and a unified format backend call interface, so as to use it as input for code generation and heterogeneous runtime calls.

5. A neural network heterogeneous partitioning rewriting system based on graph intermediate representation, characterized in that, The method applied to any one of claims 1-4 includes: Hardware Configuration and Cost Model Module: Used to store the hardware configuration of heterogeneous hardware platforms and descriptions of heterogeneous backend devices, and provides cost estimation interfaces for computation execution time estimation, data transmission time estimation, and storage feasibility constraints required by the cost model; Subgraph extraction module: used to traverse the target forward function of the intermediate representation of the high-level computation graph, identify and mark subgraph nodes based on the core operator pattern, extract subgraph node information and tensor edge information, and construct a unified graph structure representation composed of the subgraph node set and the tensor edge set; Partition Solving Module: Based on the cost model and unified graph structure representation, this module solves integer programming problems with the optimization objective of minimizing computation execution time and data transmission time, and with backend storage capacity as a constraint. It generates an allocation mapping table that corresponds one-to-one between each subgraph node and heterogeneous backend devices. Attribute annotation module: used to serialize the allocation mapping table and write it into the intermediate representation of the high-level computation graph in the form of function attributes and module attributes; The Encapsulation Function Generation and Graph Rewriting module is used to create encapsulation functions for subgraph nodes assigned to non-default heterogeneous backend devices, add target heterogeneous backend device tags to the function properties, check and automatically insert layout transformation operators, replace the original operator calls in the target forward function with calls to the encapsulation function, and output the rewritten intermediate representation of the high-level graph.

Images