Tensor program optimization method and apparatus

Abstract

The present invention relates to the technical field of deep learning. Provided are a tensor program optimization method and apparatus. The method comprises: acquiring a computation graph corresponding to a tensor program to be optimized; using a performance model to obtain a computation performance limitation operator and an access and storage performance limitation operator in the computation graph; calling a hardware operator library to compute the computation performance limitation operator, so as to generate a first code; using a graph intermediate representation to optimize the access and storage performance limitation operator, so as to generate a second code, wherein an edge in the graph intermediate representation represents a data slice on a specified memory hierarchy, and a node in the graph intermediate representation represents a hardware instruction or a group of hardware instructions and is used for executing a specified operation; and merging the first code with the second code, so as to generate an optimized tensor program. The tensor program optimization method and apparatus provided in the present invention can optimize memory performance, and the optimization effect thereof is more efficient than that of the prior art.

Description

Tensor program optimization method and device Technical Field

[0001] The present invention relates to the field of deep learning technology, and in particular to a tensor program optimization method and device based on explicit data movement and instruction-level graph intermediate representation. Background Art

[0002] Currently, memory performance is increasingly becoming a bottleneck for deep neural network applications. Numerous studies have attempted to propose optimization solutions to this memory performance issue. Tensor compilers are a common solution. This type of work improves program execution efficiency by representing tensor programs in an intermediate representation and then optimizing based on this intermediate representation. Depending on the level of abstraction of the intermediate representation, related work falls into two main categories: one focuses on code organization within operators, known as operator-specific tensor compilers; the other focuses on data reuse and operator merging between operators, known as computational graph transformation-based optimization.

[0003] Existing tensor compilers for operators, such as TVM and Ansor, use the Halide abstraction to describe programs as a computation + scheduling model. They search for suitable combinations to find the appropriate layout and execution method for a computation operator. However, this model has the following drawbacks:

[0004] 1. Existing tensor compilers lack explicit descriptions of data movement operations, making it impossible to directly optimize memory operations. For example, in TVM, the intermediate representation is loop-based, requiring the program to undergo equivalent transformations before being mapped to different memory hierarchies. This makes it difficult to optimize memory performance specifically.

[0005] 2. Existing tensor compilers have a limited search space, making it difficult to cover diverse memory operations. For example, in Ansor, the program must first merge and then search. This search order prevents the optimization process from searching all possible combinations of memory operation parameters, and thus may not find the optimal operation combination, resulting in inefficient memory usage.

[0006] Another type of optimization focuses on operator merging, such as TensorRT and DNN Fusion. This explores opportunities for merging operators within the program at the computational layer. By analyzing the data dependencies between operators, multiple operators are merged into one, reducing overall memory usage. However, this approach still has the following drawbacks:

[0007] 1. Computational graph optimization is limited by backend operators. Whether based on an operator library (such as TensorRT) or code generation (such as DNN Fusion), this type of work requires a merging rule. This means the optimization process needs to know in advance which operator combinations can be replaced with which new operator combinations. Therefore, the optimization process is limited by the computing modes supported by the backend, making it difficult to extend to new models. Furthermore, porting to other hardware platforms is difficult, making this type of work difficult to scale.

[0008] 2. On the other hand, the data dependency analysis in this type of work is too coarse-grained; the program can only analyze the reuse relationships between tensors. Data reuse in complex memory hierarchies is typically more granular, requiring even finer-grained data dependency analysis. Therefore, this type of work struggles to optimize program performance by leveraging high-level memory hierarchies.

[0009] Summary of the Invention

[0010] In view of this, the present invention provides a tensor program optimization method and apparatus to solve at least one of the above-mentioned problems.

[0011] In order to achieve the above object, the present invention adopts the following scheme:

[0012] According to a first aspect of the present invention, a tensor program optimization method is provided, the method comprising: obtaining a computational graph corresponding to a tensor program to be optimized; obtaining computationally performance-constrained operators and memory access performance-constrained operators in the computational graph using a performance model; calling a hardware operator library to calculate the computationally performance-constrained operators to generate a first code; optimizing the memory access performance-constrained operators using a graph intermediate representation to generate a second code, wherein the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing specified operations; and merging the first code and the second code to generate an optimized tensor program.

[0013] As an embodiment of the present invention, the above method uses the performance model to obtain the computing performance-limited operators and memory access performance-limited operators in the computing graph, including: using the performance model to predict the computing time and memory access time of each operator in the computing graph; and determining whether each operator belongs to a computing performance-limited operator or a memory access performance-limited operator based on the computing time and memory access time.

[0014] As an embodiment of the present invention, the specified operation in the above method includes: a data movement operation, a calculation operation, a synchronization operation or a dimension transformation operation.

[0015] As an embodiment of the present invention, the above method uses graph intermediate representation to optimize the memory access performance-limited operator, and generating the second code includes: grouping all memory access performance-limited operators into several subgraphs without external dependencies; for each operator in each subgraph, obtaining all graph intermediate representations that can implement the operator; combining the graph intermediate representations of all operators in each subgraph to obtain the graph intermediate representation combination corresponding to the subgraph; rewriting the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination; using a performance model to select several alternative graph intermediate representation combinations that meet the requirements from the optimized graph intermediate representation combination; and generating the second code based on several alternative graph intermediate representation combinations of each subgraph.

[0016] As an embodiment of the present invention, the preset rules in the above method include a first rule, a second rule and a third rule, wherein: the first rule is: when multiple data movement operations access the same data slice, a synchronization operation of the corresponding scope is added according to the scope of the movement operation; the second rule is: replace the write-synchronization-read operation with a synchronization operation, and replace the synchronization-read-read operation with a synchronization-read operation; the third rule is: force operations that do not affect memory performance to be exchanged with operations that are interchangeable with their predecessors.

[0017] As an embodiment of the present invention, in the above method, the graph intermediate representation combination is rewritten based on preset rules to obtain an optimized graph intermediate representation combination, which includes: optimizing the graph intermediate representation combination in the order of the first rule, the third rule, and the second rule until all rules can no longer optimize the graph intermediate representation combination, thereby obtaining an optimized graph intermediate representation combination.

[0018] As an embodiment of the present invention, the above method for generating a second code based on several alternative graph intermediate representation combinations of each subgraph includes: generating a kernel program for each graph intermediate representation in the alternative graph intermediate representation combination; selecting a topological order for each graph intermediate representation that minimizes the total amount of memory slices activated at the same time; generating operation instructions in the kernel program according to the topological order for the nodes in each graph intermediate representation; mapping each graph intermediate representation to specific hardware, and generating operation instructions on the specific hardware based on the operation instructions, thereby forming the code of the alternative graph intermediate representation combination; performing performance testing on the code of the alternative graph intermediate representation combination of each subgraph, and selecting the optimal code for each subgraph as the second code.

[0019] According to a second aspect of the present invention, a tensor program optimization device is provided, the device comprising: a computational graph acquisition unit for acquiring a computational graph corresponding to a tensor program to be optimized; an operator classification unit for obtaining computational performance-limited operators and memory access performance-limited operators in the computational graph using a performance model; a computational performance optimization unit for calling a hardware operator library to calculate the computational performance-limited operators and generate a first code; a memory access performance optimization unit for optimizing the memory access performance-limited operators using a graph intermediate representation and generating a second code, wherein the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing specified operations; and a code merging unit for merging the first code and the second code to generate an optimized tensor program.

[0020] As an embodiment of the present invention, the operator classification unit in the above-mentioned device includes: a performance calculation module, which is used to use a performance model to predict the calculation time and memory access time of each operator in the calculation graph; an operator classification module, which is used to determine whether each operator belongs to a computing performance-limited operator or a memory access performance-limited operator based on the calculation time and memory access time.

[0021] As an embodiment of the present invention, the specified operation in the above apparatus includes: a data movement operation, a calculation operation, a synchronization operation or a dimension transformation operation.

[0022] As an embodiment of the present invention, the memory access performance optimization unit in the above-mentioned device includes: a subgraph acquisition module, which is used to group all memory access performance-limited operators into several subgraphs without external dependencies; a graph intermediate representation acquisition module, which is used to obtain all graph intermediate representations that can implement each operator in each subgraph; a combination module, which is used to combine the graph intermediate representations of all operators in each subgraph to obtain the graph intermediate representation combination corresponding to the subgraph; a graph rewriting module, which is used to rewrite the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination; an alternative combination selection module, which is used to use a performance model to select several alternative graph intermediate representation combinations that meet the requirements from the optimized graph intermediate representation combination; a second code generation module, which is used to generate a second code based on several alternative graph intermediate representation combinations of each subgraph.

[0023] As an embodiment of the present invention, the preset rules in the above-mentioned device include a first rule, a second rule and a third rule, wherein: the first rule is: when multiple data movement operations access the same data slice, a synchronization operation of the corresponding scope is added according to the scope of the movement operation; the second rule is: replace the write-synchronization-read operation with a synchronization operation, and replace the synchronization-read-read operation with a synchronization-read operation; the third rule is: force operations that do not affect memory performance to be exchanged with operations that are interchangeable with their predecessors.

[0024] As an embodiment of the present invention, the graph rewriting module in the above-mentioned device is specifically used to: optimize the graph intermediate representation combination in the order of the first rule, the third rule, and the second rule until all rules can no longer continue to optimize the graph intermediate representation combination, thereby obtaining an optimized graph intermediate representation combination.

[0025] As an embodiment of the present invention, the second code generation module in the above-mentioned device includes: a kernel program generation submodule, which is used to generate a kernel program for each graph intermediate representation in the alternative graph intermediate representation combination; a topological order selection submodule, which is used to select a topological order for each graph intermediate representation that minimizes the total amount of memory slices activated at the same time; an operation instruction generation submodule, which is used to generate operation instructions in the kernel program for the nodes in each graph intermediate representation according to the topological order; a mapping submodule, which is used to map each graph intermediate representation to specific hardware, and generate operation instructions on the specific hardware based on the operation instructions, thereby forming the code of the alternative graph intermediate representation combination; a performance testing submodule, which is used to perform performance testing on the code of the alternative graph intermediate representation combination of each subgraph, and select the optimal code for each subgraph as the second code.

[0026] According to a third aspect of the present invention, an electronic device is provided, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the steps of the above method when executing the computer program.

[0027] According to a fourth aspect of the present invention, there is provided a computer-readable storage medium having a computer program stored thereon, wherein the computer program implements the steps of the above method when executed by a processor.

[0028] As can be seen from the above technical solutions, the tensor program optimization method and device provided by the present invention can optimize memory performance, and its optimization effect is more efficient than that of the existing technology. In addition, since the edges in the graph intermediate representation of the present application represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation, this graph intermediate representation based on explicit data movement and instruction level can also bring the following beneficial effects:

[0029] First, data movement operations are explicitly represented, so optimization of the intermediate graph representation can directly affect the program's memory operation performance, making it more conducive to subsequent search and optimization.

[0030] Second, using a graph-based intermediate representation, we can clearly express the dependencies between data pieces. This representation makes it easy to analyze the scope of data and determine at which level of the memory hierarchy the data should be stored, thereby improving program performance.

[0031] Third, the instruction-level description granularity is more precise than the operator-level description granularity. Compute layers and operator layers can be combined to generate efficient tensor program code. This covers a larger search space, allowing for the discovery of more sophisticated implementations. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] In order to more clearly illustrate the technical solutions in the embodiments of the present application, the following briefly introduces the drawings required for describing the embodiments. Obviously, the drawings described below are only some embodiments of the present application. For those skilled in the art, other drawings can be obtained based on these drawings without inventive efforts. In the drawings:

[0033] FIG1 is a flow chart of a tensor program optimization method provided in an embodiment of the present application;

[0034] FIG2 is a schematic diagram of a calculation graph provided in an embodiment of the present application;

[0035] FIG3 is a flow chart of classifying performance constraints using a performance model according to an embodiment of the present application;

[0036] FIG4 is a schematic diagram of a GIR provided in an embodiment of the present application;

[0037] FIG5 is a flow chart of optimizing a memory access performance-limited operator using GIR according to an embodiment of the present application;

[0038] FIG6 is a code diagram of an optimization process for a performance-constrained operator provided in an embodiment of the present application;

[0039] FIG7 is a schematic diagram of a flow chart of generating a second code from an alternative GIR combination according to an embodiment of the present application;

[0040] FIG8 is a schematic diagram of a scope analysis method provided in an embodiment of the present application;

[0041] FIG9 is a schematic diagram of the generation process from a computation graph to a GIR according to an embodiment of the present application;

[0042] FIG10 is a schematic diagram of a GIR fragment provided in an embodiment of the present application;

[0043] FIG11 is a schematic diagram of the optimization process of the attention module of the system provided by an embodiment of the present application when GPT-2 is used to generate a model;

[0044] FIG12 is a schematic diagram comparing the optimization effects of the optimization method of the present application and the prior art provided in an embodiment of the present application;

[0045] FIG13 is a schematic diagram of the structure of a tensor program optimization device provided in an embodiment of the present application;

[0046] FIG14 is a schematic diagram of the structure of an operator classification unit provided in an embodiment of the present application;

[0047] FIG15 is a schematic structural diagram of a memory access performance optimization unit provided in an embodiment of the present application;

[0048] FIG16 is a schematic diagram of the structure of a second code generation module provided in an embodiment of the present application;

[0049] FIG17 is a schematic block diagram of the system structure of the electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0050] To make the purpose, technical solutions and advantages of the embodiments of the present invention more clear, the embodiments of the present invention are further described in detail below with reference to the accompanying drawings. Here, the exemplary embodiments of the present invention and their descriptions are used to explain the present invention, but are not intended to limit the present invention.

[0051] FIG1 is a flow chart of a tensor program optimization method provided in an embodiment of the present application. This embodiment is described from the perspective of a tensor program optimization system. The method includes the following steps:

[0052] Step S101: Obtain the computational graph corresponding to the tensor program to be optimized.

[0053] A computation graph is a data structure composed of a set of nodes and edges, representing computational processes. Nodes represent operators, and edges represent data flows. Figure 2 shows a schematic diagram of a computation graph, which includes four operators: Rectified Linear Unit (RELU), Concatenation (Concat), Matrix Transpose (Transpose), and Split.

[0054] Step S102: using the performance model to obtain computing performance-limited operators and memory access performance-limited operators in the computation graph.

[0055] The performance model of this embodiment can classify all operators in the input computation graph into those with limited computational performance and those with limited memory access performance, depending on the target hardware platform. The performance of operators with limited computational performance is primarily affected by the amount of computation and depends primarily on the computational performance of the hardware, such as most matrix multiplications and two-dimensional convolutions. The performance of operators with limited memory access performance is primarily affected by the amount of memory access and access performance, and depends primarily on the memory performance of the hardware, such as operations such as point-by-point calculations, transpositions, and reductions.

[0056] Preferably, as shown in FIG3 , this step may further include the following sub-steps:

[0057] Step S1021: using the performance model to predict the computation time and memory access time of each operator in the computation graph.

[0058] Step S1022: Determine whether each operator is a compute-performance-limited operator or a memory-access-performance-limited operator based on the computation time and memory access time. Specifically, for example, the program computation amount and memory access amount can be calculated based on the computation time and memory access time predicted by the performance model. The ratio of the computation amount to the memory access amount is compared with a preset threshold. If the ratio exceeds the preset threshold, the operator is considered to be a compute-performance-limited operator; otherwise, it is a memory-access-limited operator.

[0059] Step S103: calling a hardware operator library to calculate the operator with limited computing performance and generate a first code.

[0060] Step S104: Utilize a graph intermediate representation to optimize the memory access performance-limited operator and generate a second code, wherein the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation.

[0061] The solution of this embodiment mainly optimizes the memory performance of tensor programs. The basis of this optimization lies in the aforementioned graph intermediate representation, that is, the graph intermediate representation based on explicit data movement and instruction level. For the convenience of description below, this graph intermediate representation is replaced by GIR (Instruction-Level Graph Intermediate Representation with Explicit Data Movement Description).

[0062] As shown in Figure 4, this is a schematic diagram of a GIR provided in an embodiment of the present application, which includes a complete GIR and a simplified GIR, wherein the edges (i.e., the arrows with pointing directions) represent a data slice on a specific memory hierarchy. For example, D0 in Figure 4 represents a data slice of group p*q on DRAM. The nodes in the GIR represent one or a group of hardware instructions, which are used to perform specific specified operations, such as the MOVE operation and RELU operation in Figure 4. Preferably, the specified operations here include synchronization operations or dimension transformation operations in addition to the above-mentioned move operations and calculation operations. In addition, each GIR also includes a degree of parallelism n, which indicates that the operation described by the GIR is executed n times in parallel.

[0063] The above-mentioned GIR description can bring the following beneficial effects to this application: on the one hand, the data movement operation is explicitly represented, so the optimization of GIR can directly affect the memory operation performance of the program, which is more conducive to subsequent search and optimization. On the one hand, the use of graph-based intermediate representation can clearly represent the dependencies between data slices. Based on this representation, the scope of the data can be easily analyzed to determine at which level of the memory hierarchy the data should be stored, thereby improving program performance. On the other hand, the description granularity at the instruction level is more accurate than that at the operator level. The computational layer and the operator layer can be combined to generate efficient tensor program code. Covering a larger search space allows more sophisticated implementations to be searched.

[0064] The following further describes in detail how to optimize memory access performance-limited operators based on the GIR. FIG5 is a flow chart of optimizing memory access performance-limited operators using GIR according to an embodiment of the present application, which includes the following sub-steps:

[0065] Step S1041: Group all operators with limited memory access performance into several subgraphs without external dependencies.

[0066] In this embodiment, if some performance-constrained operators have only internal connections and no external dependencies, these operators can be grouped into the same subgraph. Subsequent operations are performed on the subgraphs. The code for each subgraph is then merged to obtain the optimized code for all performance-constrained operators in the computation graph, i.e., the second code.

[0067] Step S1042: For each operator in each subgraph, obtain all GIRs that can implement the operator.

[0068] In this embodiment, for each operator in the subgraph, all possible GIRs can be generated based on its computation mode. For example, the Transpose operator can generate different implementation modes and different parameter sizes. Therefore, each operator will correspond to several GIR codes.

[0069] Step S1043: Combine the GIRs of all operators in each subgraph to obtain the GIR combination corresponding to the subgraph.

[0070] Since each operator corresponds to several GIRs, there are many ways to combine GIRs between operators. For example, in a subgraph, there are two operators, operator 1 and operator 2, where operator 1 has 3 GIRs and operator 2 has 2 GIRs. Then there are 6 GIR combinations in the subgraph.

[0071] Step S1044: rewriting the GIR combination based on preset rules to obtain an optimized GIR combination.

[0072] The graph rewriting in this step optimizes the GIR combination obtained in the above step to minimize the amount of memory access in the code corresponding to this GIR combination, thereby improving memory performance. The graph rewriting in this embodiment is based on three preset rules. Based on these three preset rules, three types of transformations are allowed for the GIR. Here, we refer to these three preset rules as the first rule, the second rule, and the third rule, where:

[0073] The first rule states: When multiple data move operations access the same data slice, a synchronization operation is added based on the scope of the move operations. For example, if two data move operations access the same data slice on a DRAM block, a synchronization operation can be added based on the access pattern. On the GPU, this synchronization scope can be a thread, a warp, or the entire device. This operation ensures that the data scope is always reduced and the data storage hierarchy is always improved.

[0074] The second rule is: replace the write-sync-read operation with a sync operation, and replace the sync-read-read operation with a sync-read operation; such transformations can ensure that the total memory access volume is reduced.

[0075] The third rule states: operations that do not affect memory performance must be swapped with their predecessors. Examples of such operations include point-by-point operations, transposition operations, reduction operations, and broadcast operations. By forcing these operations to swap with their predecessors, any adjacent operations that could be eliminated are eliminated. This rule is irreversible because it defines the direction of the swap.

[0076] After defining the above three rules, this application designed a heuristic algorithm to sequentially apply the above three rules until the program can no longer be optimized. The code of the optimization process is shown in Figure 6. Specifically, the GIR combination is optimized in the order of the first rule, the third rule, and the second rule until all rules can no longer optimize the GIR combination. Finally, the optimized GIR combination is obtained. The complexity of this optimization process is O(N+M), where N is the number of nodes in the GIR and M is the number of edges in the GIR. This complexity allows the optimization process of the GIR combination to be executed quickly.

[0077] Step S1045: using the performance model to select several candidate GIR combinations that meet the requirements from the optimized GIR combinations.

[0078] Since the number of GIR combinations obtained after step S1044 may be relatively large, this application continues to apply the performance model to analyze the performance of the optimized GIR combinations on the basis of the above, and selects a predetermined number of GIR combinations with the best performance as alternative GIR combinations for subsequent code generation.

[0079] Step S1046: Generate a second code based on several candidate GIR combinations of each subgraph.

[0080] This step involves GIR code generation optimization, that is, it is necessary to select the best GIR combination from the above-mentioned alternative GIR combinations to generate the second code.

[0081] Preferably, as shown in FIG7 , this step may further include the following sub-steps:

[0082] Step S10461: Generate a kernel program for each graph intermediate representation in the candidate GIR combination.

[0083] Step S10462: For each GIR, a topological order is selected that minimizes the total number of memory slices activated at the same time. The topological order selection criterion here is to minimize the total number of memory slices activated at the same time, that is, to perform graph reordering after graph rewriting.

[0084] Step S10463: Each node in the GIR generates operation instructions in the kernel program according to the topological order.

[0085] Step S10464: Map each GIR to specific hardware, and generate an operation instruction on the specific hardware based on the operation instruction, thereby forming a code for the candidate GIR combination.

[0086] The operation instructions generated in step S10463 are general instructions, but tensor programs need to run on specific hardware platforms, so each GIR needs to be mapped to specific hardware. For example, on NVIDIA GPUs, parallelism is mapped to the warp granularity, and the block and device number of each warp is analyzed according to the warp-block-device model to generate GPU code. On the Cambrian MLU, considering the necessity of pipeline parallelism in kernel programs, each degree of parallelism is mapped to a pipeline stage. The scope of each data movement operation is analyzed according to the pipeline-core-device model. Figure 8 shows a schematic diagram of the scope analysis method provided by this application, which is performed according to the warp-block-device model. The figure illustrates the synchronization operations required for different memory access operation combinations on NVIDIA GPUs. Different grayscales represent different warps participating in the memory access operation, and different shapes represent different memory access modes. In the figure, the first row shows cases where synchronization is required within a warp, the second row shows cases where synchronization is required within a block, and the third row shows cases where synchronization is required within a device.

[0087] As can be seen, the scope analysis method described above abstracts different devices into constraints during the search process, eliminating the need for platform-specific optimization. Furthermore, for each computational operation on the GIR, the system directly generates an instruction or set of instructions for the specific device. Therefore, porting the system to a new hardware platform requires only the translation of different types of operations into instructions, resulting in a very low code footprint. For example, porting the system to an NVIDIA GPU requires no more than 200 lines of code.

[0088] Step S10465: Perform a performance test on the codes of the candidate GIR combinations of each subgraph, and select the optimal code of each subgraph as the second code.

[0089] After obtaining the code for the candidate GIR combinations, we perform performance tests on each code and select the optimal code for each subgraph as the second code. This code is then merged with the first code for the previously mentioned performance-constrained operators to generate the program for the end-to-end model. This program is implemented in the native language of the specific hardware and can be compiled and run directly.

[0090] Step S105: merging the first code and the second code to generate an optimized tensor program.

[0091] As can be seen from the above technical solutions, the tensor program optimization method provided by the present invention can optimize memory performance, and its optimization effect is more efficient than that of the existing technology. In addition, since the edges in the graph intermediate representation of the present application represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation, this graph intermediate representation based on explicit data movement and instruction level can also bring the following beneficial effects:

[0092] First, data movement operations are explicitly represented, so optimization of the intermediate graph representation can directly affect the program's memory operation performance, making it more conducive to subsequent search and optimization.

[0093] Second, using a graph-based intermediate representation, we can clearly express the dependencies between data pieces. This representation makes it easy to analyze the scope of data and determine at which level of the memory hierarchy the data should be stored, thereby improving program performance.

[0094] Third, the instruction-level description granularity is more precise than the operator-level description granularity. Compute layers and operator layers can be combined to generate efficient tensor program code. This covers a larger search space, allowing for the discovery of more sophisticated implementations.

[0095] The above steps are further described below through several specific embodiments.

[0096] As shown in Figure 9, it is a schematic diagram of the generation process from the computation graph to the GIR provided by an embodiment of the present application. As shown in Figure 9, for the activation function SiLU, it has two operators Sigmoid and Mul. For the operators Sigmoid and Mul, it will generate several alternative GIRs respectively, and then search each combination for them, and then merge them to generate several alternative GIR combinations. Figure 9 illustrates the selection of a GIR for the operator Sigmoid and the operator Mul respectively, and then the first rule is used to add two synchronization operations between the two data movement operations to obtain the merged GIR combination.

[0097] As shown in Figure 10, a schematic diagram of a GIR fragment provided by an embodiment of the present application shows the connection relationship between common operations in GIR. Due to the existence of the point-by-point operation of the blank circle, the original GIR cannot be further optimized. However, the third rule is applied to perform the exchange operation, and after the operation is moved to the front, a synchronization operation is added between DRAM Read B and DRAM Write B based on the first rule. Then DRAM Read B, DRAM Write B and the intermediate synchronization operation can be converted into a synchronization operation using rule two, and then DRAM Write A, DRAM Read A and the intermediate synchronization operation can continue to be converted into a synchronization operation according to rule two. Then only point-by-point operations and synchronization operations are left. After removing the synchronization operation, only an optimized point-by-point operation is left. This example proves that the GIR optimization process can effectively reduce the memory access volume of the entire program.

[0098] As shown in Figure 11, this is the optimization process of the attention module of the system provided by the embodiment of the present application when GPT-2 is used to generate a model. For the two matrix multiplication (MatMul) operations in the calculation graph, the performance model determines that they are memory performance-limited operations, and generates several alternative GIRs for these two operations. In the GIR combination with the best performance, the two matrix multiplication operations use different GIR implementations. The specific difference lies in the different dimensions of the bcast and reduce operations. After optimizing the GIR combination, the system can generate the entire attention operator into a kernel code. On NVIDIA GPUs, compared with TensorRT, TVM and other workloads, the performance improvement in this example can be up to 1.98 times.

[0099] FIG12 is a schematic diagram showing a comparison of the optimization effects of the present invention's optimization method and the prior art provided in an embodiment of the present invention. In this embodiment, three platforms, NVIDIA TESLA A100 GPU, AMD MI100 GPU, and Cambricon MLU-370, are used. The first row of FIG12 is an end-to-end performance comparison chart on the NVIDIA TESLA A100 GPU, the second row is an end-to-end performance comparison chart on the AMD MI100 GPU, and the third row is an end-to-end performance comparison chart on the Cambricon MLU-370.

[0100] This test program experiment used seven real-world DNN models: BERT, ViT, GPT2, SAR-DRN, EfficientNet, ShuffleNet, and RedNet-50. BERT, ViT, and GPT2 are based on the Transformer architecture and are designed for applications such as natural language processing and image recognition. SAR-DRN, EfficientNet, ShuffleNet, and RedNet-50 are based on the CNN architecture and are designed for various image tasks, including image classification and super-resolution.

[0101] The existing technologies compared with this application are A: PyTorch, B: TorchScript, C: TensorFlow, D: TF-XLA, E: TensorRT, F: TVM (Ansor), G: MagicMind, and this application is displayed with the label H in the figure.

[0102] As shown in Figure 12, in an end-to-end experiment on an NVIDIA A100 GPU, the present invention achieves a speedup of up to 1.98x compared to the current best-performing TensorRT. Furthermore, Figure 12 as a whole shows that the present invention's speedup is among the best across various platforms. Therefore, the present invention clearly outperforms existing technologies in memory performance optimization.

[0103] As shown in Figure 13, a structural diagram of a tensor program optimization device provided in an embodiment of the present application, the device includes: a computational graph acquisition unit 100, an operator classification unit 200, a computational performance optimization unit 300, a memory access performance optimization unit 400 and a code merging unit 500, wherein the operator classification unit 200 is respectively connected to the computational graph acquisition unit 100, the computational performance optimization unit 300 and the memory access performance optimization unit 400, and the code merging unit 500 is respectively connected to the computational performance optimization unit 300 and the memory access performance optimization unit 400.

[0104] The computation graph acquisition unit 100 is used to obtain the computation graph corresponding to the tensor program to be optimized.

[0105] The operator classification unit 200 is used to obtain the computing performance-limited operators and the memory access performance-limited operators in the computation graph using the performance model.

[0106] The computing performance optimization unit 300 is used to call the hardware operator library to calculate the computing performance-limited operator and generate a first code.

[0107] The memory access performance optimization unit 400 is used to optimize the memory access performance-limited operator using a graph intermediate representation to generate a second code, where the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation.

[0108] The code merging unit 500 is configured to merge the first code and the second code to generate an optimized tensor program.

[0109] Preferably, as shown in FIG14 , the operator classification unit 200 in this embodiment may include: a performance calculation module 210 and an operator classification module 220 , which are interconnected.

[0110] The performance calculation module 210 is used to predict the computation time and memory access time of each operator in the computation graph using a performance model.

[0111] The operator classification module 220 is configured to determine whether each operator is a computing performance-limited operator or a memory access performance-limited operator based on the computing time and the memory access time.

[0112] Preferably, the above-mentioned specified operation includes: a data movement operation, a calculation operation, a synchronization operation or a dimension transformation operation.

[0113] Preferably, as shown in FIG15 , the memory access performance optimization unit 400 includes: a subgraph acquisition module 410, a graph intermediate representation acquisition module 420, a combination module 430, a graph rewriting module 440, an alternative combination selection module 450 and a second code generation module 460, which are connected in sequence.

[0114] The subgraph acquisition module 410 is used to group all operators with limited memory access performance into several subgraphs without external dependencies.

[0115] The graph intermediate representation acquisition module 420 is used to acquire, for each operator in each subgraph, all graph intermediate representations that can realize the operator.

[0116] The combining module 430 is used to combine the graph intermediate representations of all operators in each subgraph to obtain a graph intermediate representation combination corresponding to the subgraph.

[0117] The graph rewriting module 440 is configured to rewrite the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination.

[0118] The candidate combination selection module 450 is used to select several candidate graph intermediate representation combinations that meet the requirements from the optimized graph intermediate representation combinations by using the performance model.

[0119] The second code generation module 460 is configured to generate a second code based on a combination of several candidate graph intermediate representations of each subgraph.

[0120] Preferably, the above-mentioned preset rules include a first rule, a second rule and a third rule, wherein:

[0121] The first rule is: when multiple data movement operations access the same data slice, a synchronization operation corresponding to the scope of the movement operation is added;

[0122] The second rule is: replace the write-sync-read operation with a sync operation, and replace the sync-read-read operation with a sync-read operation;

[0123] The third rule is: operations that do not affect memory performance are forced to be exchanged with operations that are exchangeable with their predecessors.

[0124] Preferably, the above-mentioned graph rewriting module 440 is specifically used to optimize the graph intermediate representation combination in the order of the first rule, the third rule, and the second rule until all rules can no longer optimize the graph intermediate representation combination, thereby obtaining an optimized graph intermediate representation combination.

[0125] Preferably, as shown in FIG16 , the second code generation module 460 may further include: a kernel program generation submodule 461 , a topology order selection submodule 462 , an operation instruction generation submodule 463 , a mapping submodule 464 and a performance test submodule 465 , which are connected in sequence.

[0126] The kernel program generation submodule 461 is used to generate a kernel program for each graph intermediate representation in the candidate graph intermediate representation combination.

[0127] The topological order selection submodule 462 is used to select, for each graph intermediate representation, a topological order that minimizes the total amount of memory slices activated at the same time.

[0128] The operation instruction generation submodule 463 is used to generate operation instructions in the kernel program from the nodes in each graph intermediate representation according to the topological order.

[0129] The mapping submodule 464 is used to map each graph intermediate representation to specific hardware, and generate operation instructions on the specific hardware based on the operation instructions, thereby forming the code of the candidate graph intermediate representation combination.

[0130] The performance testing submodule 465 is used to perform a performance test on the codes of the candidate graph intermediate representation combinations of each subgraph, and select the optimal code of each subgraph as the second code.

[0131] The detailed description of each of the above units can be found in the corresponding description in the aforementioned method embodiment, and will not be repeated here.

[0132] As can be seen from the above technical solutions, the tensor program optimization device provided by the present invention can optimize memory performance, and its optimization effect is more efficient than that of the existing technology. In addition, since the edges in the graph intermediate representation of the present application represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation, this graph intermediate representation based on explicit data movement and instruction level can also bring the following beneficial effects:

[0133] First, data movement operations are explicitly represented, so optimization of the intermediate graph representation can directly affect the program's memory operation performance, making it more conducive to subsequent search and optimization.

[0134] Second, using a graph-based intermediate representation, we can clearly express the dependencies between data pieces. This representation makes it easy to analyze the scope of data and determine at which level of the memory hierarchy the data should be stored, thereby improving program performance.

[0135] Third, the instruction-level description granularity is more precise than the operator-level description granularity. Compute layers and operator layers can be combined to generate efficient tensor program code. This covers a larger search space, allowing for the discovery of more sophisticated implementations.

[0136] An embodiment of the present invention further provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the above method is implemented when the processor executes the program.

[0137] An embodiment of the present invention further provides a computer-readable storage medium storing a computer program for executing the above method.

[0138] As shown in FIG17 , the electronic device 600 may further include: a communication module 110, an input unit 120, an audio processor 130, a display 160, and a power supply 170. It is worth noting that the electronic device 600 does not necessarily include all of the components shown in FIG17 ; in addition, the electronic device 600 may also include components not shown in FIG17 , for which reference may be made to the prior art.

[0139] As shown in FIG. 17 , the central processing unit 100 is sometimes also referred to as a controller or an operation control unit, and may include a microprocessor or other processor device and / or logic device. The central processing unit 100 receives inputs and controls the operations of various components of the electronic device 600 .

[0140] Memory 140 may be, for example, one or more of a cache, flash memory, hard drive, removable media, volatile memory, non-volatile memory, or other suitable devices. It may store the aforementioned failure-related information and may also store programs that execute the relevant information. The CPU 100 may execute the programs stored in memory 140 to implement information storage or processing.

[0141] The input unit 120 provides input to the CPU 100. The input unit 120 may be, for example, a keypad or touch input device. The power supply 170 is used to provide power to the electronic device 600. The display 160 is used to display objects such as images and text. The display may be, for example, an LCD display, but is not limited thereto.

[0142] The memory 140 may be a solid-state memory, such as a read-only memory (ROM), a random access memory (RAM), or a SIM card. Alternatively, it may be a memory that retains information even when power is off, can be selectively erased, and is provided with more data. Examples of such memory are sometimes referred to as EPROMs. The memory 140 may also be some other type of device. The memory 140 includes a buffer memory 141 (sometimes referred to as a buffer). The memory 140 may include an application / function storage unit 142 for storing application programs and function programs or processes for executing the operations of the electronic device 600 via the central processing unit 100.

[0143] The memory 140 may also include a data storage unit 143 for storing data, such as contacts, digital data, pictures, sounds, and / or any other data used by the electronic device. The driver storage unit 144 of the memory 140 may include various driver programs for communication functions of the electronic device and / or for executing other functions of the electronic device (such as messaging applications, address book applications, etc.).

[0144] The communication module 110 is a transmitter / receiver 110 that transmits and receives signals via an antenna 111. The communication module (transmitter / receiver) 110 is coupled to the central processor 100 to provide input signals and receive output signals, which may be the same as in a conventional mobile communication terminal.

[0145] Based on different communication technologies, multiple communication modules 110 may be provided in the same electronic device, such as a cellular network module, a Bluetooth module, and / or a wireless local area network module. The communication module (transmitter / receiver) 110 is also coupled to a speaker 131 and a microphone 132 via an audio processor 130 to provide audio output via the speaker 131 and receive audio input from the microphone 132, thereby implementing common telecommunication functions. The audio processor 130 may include any suitable buffer, decoder, amplifier, etc. Furthermore, the audio processor 130 is also coupled to the central processing unit 100, enabling local recording via the microphone 132 and playback of stored audio via the speaker 131.

[0146] It will be understood by those skilled in the art that embodiments of the present invention may be provided as methods, systems, or computer program products. Thus, the present invention may take the form of an entirely hardware embodiment, an entirely software embodiment, or an embodiment combining software and hardware. Furthermore, the present invention may take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to magnetic disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0147] The present invention is described with reference to the flowcharts and / or block diagrams of the methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or box in the flowchart and / or block diagram, as well as the combination of processes and / or boxes in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device produce a device for implementing the functions specified in one or more processes in the flowchart and / or one or more boxes in the block diagram.

[0148] These computer program instructions may also be stored in a computer-readable memory that can direct a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce a product including an instruction device that implements the functions specified in one or more processes in the flowchart and / or one or more boxes in the block diagram.

[0149] These computer program instructions can also be loaded onto a computer or other programmable data processing device so that a series of operating steps are executed on the computer or other programmable device to produce a computer-implemented process, so that the instructions executed on the computer or other programmable device provide steps for implementing the functions specified in one or more processes in the flowchart and / or one or more boxes in the block diagram.

[0150] Specific embodiments are used in the present invention to illustrate the principles and implementation methods of the present invention. The description of the above embodiments is only used to help understand the method of the present invention and its core ideas. At the same time, for those skilled in the art, according to the ideas of the present invention, there may be changes in the specific implementation methods and application scopes. In summary, the contents of this specification should not be understood as limiting the present invention.

Claims

1. A tensor program optimization method, characterized in that: The method comprises: Get the computational graph corresponding to the tensor program to be optimized; Obtaining computing performance-limited operators and memory access performance-limited operators in the computation graph using a performance model; Calling a hardware operator library to calculate the operator with limited computing performance, and generating a first code; Optimizing the memory access performance-constrained operator using a graph intermediate representation to generate a second code, wherein the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation; The first code and the second code are combined to generate an optimized tensor program.

2. The tensor program optimization method according to claim 1, characterized in that: The obtaining of the computing performance-limited operators and the memory access performance-limited operators in the computing graph by using the performance model comprises: Using the performance model to predict the computation time and memory access time of each operator in the computation graph; Based on the calculation time and the memory access time, it is determined whether each operator is a calculation performance-limited operator or a memory access performance-limited operator.

3. The tensor program optimization method according to claim 1, characterized in that: The specified operation includes: a data movement operation, a calculation operation, a synchronization operation or a dimension transformation operation.

4. The tensor program optimization method according to claim 1, characterized in that: The step of optimizing the memory access performance-limited operator by using the graph intermediate representation to generate the second code comprises: Group all memory-limited operators into several subgraphs without external dependencies; For each operator in each subgraph, obtain all graph intermediate representations that can implement the operator; Combine the graph intermediate representations of all operators in each subgraph to obtain the graph intermediate representation combination corresponding to the subgraph; Rewriting the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination; Selecting a plurality of candidate graph intermediate representation combinations that meet the requirements from the optimized graph intermediate representation combinations using a performance model; A second code is generated based on a combination of several candidate graph intermediate representations of each subgraph.

5. The tensor program optimization method according to claim 4, characterized in that: The preset rules include a first rule, a second rule and a third rule, wherein: The first rule is: when multiple data movement operations access the same data slice, a synchronization operation corresponding to the scope is added according to the scope of the movement operation; The second rule is: replace the write-sync-read operation with a sync operation, and replace the sync-read-read operation with a sync-read operation; The third rule is: operations that do not affect memory performance are forced to be exchanged with operations that are exchangeable with their predecessors.

6. The tensor program optimization method according to claim 5, characterized in that: The step of rewriting the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination includes: optimizing the graph intermediate representation combination in the order of the first rule, the third rule, and the second rule until all rules can no longer optimize the graph intermediate representation combination, thereby obtaining an optimized graph intermediate representation combination.

7. The tensor program optimization method according to claim 4, characterized in that: The generating of the second code based on the combination of several candidate graph intermediate representations of each subgraph comprises: Generate a kernel program for each graph intermediate representation in the candidate graph intermediate representation combination; For each graph intermediate representation, select the topological order that minimizes the total amount of memory slices activated at the same time; The nodes in the intermediate representation of each graph generate operation instructions in the kernel program according to the topological order; Mapping each graph intermediate representation to specific hardware, and generating operation instructions on the specific hardware based on the operation instructions, thereby forming a code for the candidate graph intermediate representation combination; The codes of the candidate graph intermediate representation combinations of each subgraph are performance tested, and the optimal code of each subgraph is selected as the second code.

8. A tensor program optimization device, characterized in that: The device comprises: A computational graph acquisition unit, used to acquire a computational graph corresponding to the tensor program to be optimized; An operator classification unit, used for obtaining computing performance-limited operators and memory access performance-limited operators in the computation graph by using a performance model; A computing performance optimization unit, used for calling a hardware operator library to calculate the computing performance-limited operator and generate a first code; A memory access performance optimization unit, configured to optimize the memory access performance-constrained operator by using a graph intermediate representation to generate a second code, wherein the edges in the graph intermediate representation represent data slices on a specified memory hierarchy, and the nodes in the graph intermediate representation represent one or a group of hardware instructions for performing a specified operation; A code merging unit is used to merge the first code and the second code to generate an optimized tensor program.

9. The tensor program optimization device according to claim 8, characterized in that: The operator classification unit comprises: A performance calculation module, used to predict the calculation time and memory access time of each operator in the calculation graph by using a performance model; The operator classification module is used to determine whether each operator is a computing performance-limited operator or a memory access performance-limited operator based on the computing time and the memory access time.

10. The tensor program optimization device according to claim 8, characterized in that: The specified operation includes: a data movement operation, a calculation operation, a synchronization operation or a dimension transformation operation.

11. The tensor program optimization device according to claim 8, characterized in that: The memory access performance optimization unit comprises: A subgraph acquisition module is used to group all memory-access performance-constrained operators into several subgraphs without external dependencies; A graph intermediate representation acquisition module is used to acquire all graph intermediate representations that can realize each operator in each subgraph; A combination module is used to combine the graph intermediate representations of all operators in each subgraph to obtain the graph intermediate representation combination corresponding to the subgraph; A graph rewriting module, used to rewrite the graph intermediate representation combination based on preset rules to obtain an optimized graph intermediate representation combination; An alternative combination selection module, used for selecting a plurality of alternative graph intermediate representation combinations that meet the requirements from the optimized graph intermediate representation combinations by using a performance model; The second code generation module is used to generate a second code based on a combination of several candidate graph intermediate representations of each subgraph.

12. The tensor program optimization device according to claim 11, characterized in that: The preset rules include a first rule, a second rule and a third rule, wherein: The first rule is: when multiple data movement operations access the same data slice, a synchronization operation corresponding to the scope is added according to the scope of the movement operation; The second rule is: replace the write-sync-read operation with a sync operation, and replace the sync-read-read operation with a sync-read operation; The third rule is: operations that do not affect memory performance are forced to be exchanged with operations that are exchangeable with their predecessors.

13. The tensor program optimization device according to claim 12, characterized in that: The graph rewriting module is specifically used to optimize the graph intermediate representation combination in the order of the first rule, the third rule, and the second rule until all rules can no longer optimize the graph intermediate representation combination, thereby obtaining an optimized graph intermediate representation combination.

14. The tensor program optimization device according to claim 11, characterized in that: The second code generation module comprises: A kernel program generation submodule, used for generating a kernel program for each graph intermediate representation in the candidate graph intermediate representation combination; The topological order selection submodule is used to select a topological order for each graph intermediate representation that minimizes the total amount of memory slices activated at the same time; An operation instruction generation submodule, used for generating operation instructions in a kernel program from the nodes in each intermediate representation of the graph according to the topological order; A mapping submodule, used for mapping each graph intermediate representation to specific hardware, and generating operation instructions on the specific hardware based on the operation instructions, so as to form a code for the candidate graph intermediate representation combination; The performance testing submodule is used to perform performance testing on the codes of the candidate graph intermediate representation combination of each subgraph, and select the optimal code of each subgraph as the second code.

15. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 7 are implemented.

16. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images