Attention mechanism calculation optimization method and device, storage medium and program product

Abstract

The embodiment of the invention provides an attention mechanism calculation optimization method and device, a storage medium and a program product, and relates to the technical field of artificial intelligence, and the method comprises the steps: segmenting a calculation graph of attention mechanism calculation, and obtaining a plurality of calculation sub-graphs; and then compiling each calculation sub-graph into a corresponding kernel function, and executing the kernel function corresponding to each calculation sub-graph to obtain a target calculation result of attention mechanism calculation. As the number of the obtained calculation sub-graphs is smaller than the number of the operators contained in the calculation graph, compared with the mode that each operator in the calculation graph is compiled into one kernel function, the mode that each calculation sub-graph in the calculation graph is compiled into one kernel function reduces the number of the kernel functions obtained through compiling, and the calculation efficiency is improved. Therefore, in the attention mechanism calculation process, the time overhead of kernel function emission is reduced, and then the performance of attention mechanism calculation is improved.

Description

Technical Field

[0001] The embodiments of the present application relate to the field of artificial intelligence technology, and in particular, to an optimization method, device, storage medium and program product for attention mechanism calculation. Background Art

[0002] In the field of deep learning, attention mechanism calculation is widely present in various deep learning models, especially in the era of large models, so its execution performance becomes particularly important.

[0003] In the related technology, in the attention mechanism calculation process, according to the execution order of multiple operators in the calculation graph corresponding to the attention mechanism module (AttentionBlock), multiple operators are compiled into corresponding kernel functions and executed in turn, where each operator corresponds to a kernel function.

[0004] Before each operator is executed, the corresponding kernel function needs to be launched. Since the computation graph contains many operators and the compiled kernel functions are also many, the time overhead of launching the kernel function is relatively large during the attention mechanism calculation process. Summary of the invention

[0005] The embodiments of the present application provide an optimization method, device, storage medium and program product for attention mechanism calculation, which are used to reduce the time overhead of kernel function emission during attention mechanism calculation, thereby improving the efficiency of attention mechanism calculation.

[0006] On the one hand, an embodiment of the present application provides an optimization method for attention mechanism calculation, including:

[0007] Splitting the computation graph of the attention mechanism calculation to obtain a plurality of computation subgraphs, wherein each computation subgraph includes at least one operator, and the number of the plurality of computation subgraphs is less than the number of operators included in the computation graph;

[0008] Compile each computational subgraph into a corresponding kernel function, and execute the kernel function corresponding to each computational subgraph to obtain the target computational result of the attention mechanism calculation.

[0009] On the one hand, an embodiment of the present application provides an optimization device for attention mechanism calculation, including:

[0010] A segmentation module, used to segment the computation graph of the attention mechanism calculation to obtain multiple computation subgraphs, wherein each computation subgraph includes at least one operator, and the number of the multiple computation subgraphs is less than the number of operators included in the computation graph;

[0011] The execution module is used to compile each computational subgraph into a corresponding kernel function, and execute the kernel function corresponding to each of the computational subgraphs to obtain the target computational result of the attention mechanism calculation.

[0012] Optionally, the segmentation module is specifically used for:

[0013] The computation graph of the attention mechanism calculation is split to obtain multiple computation subgraphs and the execution order of the multiple computation subgraphs;

[0014] The execution module is specifically used for:

[0015] According to the execution order, each computational subgraph is compiled into a corresponding kernel function, and the kernel function corresponding to each computational subgraph is executed to obtain the target computational result of the attention mechanism calculation.

[0016] Optionally, the segmentation module is further used for:

[0017] Each computational subgraph is compiled into a corresponding kernel function, and the kernel function corresponding to each computational subgraph is executed. Before obtaining the target computational result calculated by the attention mechanism, a subgraph optimization operation is performed on the computational subgraphs among the multiple computational subgraphs that meet the preset optimization conditions.

[0018] Optionally, the segmentation module is specifically used for:

[0019] If the computation subgraph includes at least one operator combination for layout adjustment, the computation subgraph satisfies the preset optimization condition, and the following subgraph optimization operation is performed on the computation subgraph:

[0020] For at least one operator combination for layout adjustment included in the computation subgraph, the following operations are performed respectively: multiple operators included in one operator combination are replaced with a fused operator having the same function as that of the one operator combination.

[0021] Optionally, the one operator combination includes: a deformation operator and a reordering operator, wherein the deformation operator is executed before the reordering operator, or the deformation operator is executed after the reordering operator.

[0022] Optionally, the segmentation module is specifically used for:

[0023] If the computation subgraph includes a tensor core operator and a vector core operator, the computation subgraph satisfies the preset optimization conditions, and the following subgraph optimization operations are performed on the computation subgraph:

[0024] Thread local registers are allocated for data exchange between the tensor core operator and the vector core operator, wherein the tensor core operator refers to an operator executed by a tensor computing unit, and the vector core operator refers to an operator executed by a vector computing unit.

[0025] Optionally, the tensor core operator and the vector core operator are executed in parallel.

[0026] On the one hand, an embodiment of the present application provides a computer device, including a memory, an artificial intelligence chip, and a computer program stored in the memory and executable on the artificial intelligence chip, wherein the artificial intelligence chip implements the steps of the above-mentioned optimization method for attention mechanism calculation when executing the computer program.

[0027] On the one hand, an embodiment of the present application provides a computer-readable storage medium, which stores a computer program that can be executed by a computer device. When the computer program runs on the computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0028] On the one hand, an embodiment of the present application provides a computer program product, which includes a computer program stored on a computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0029] In an embodiment of the present application, the calculation graph of the attention mechanism calculation is segmented to obtain multiple calculation subgraphs. Each calculation subgraph is then compiled into a corresponding kernel function, and the kernel function corresponding to each calculation subgraph is executed to obtain the target calculation result of the attention mechanism calculation. Since the number of multiple calculation subgraphs obtained is less than the number of operators contained in the calculation graph, compared to compiling each operator in the calculation graph into a kernel function, the method of compiling each calculation subgraph in the calculation graph into a kernel function reduces the number of compiled kernel functions, thereby reducing the time overhead of kernel function emission during the attention mechanism calculation process, thereby improving the performance of the attention mechanism calculation. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.

[0031] Figure 1 A schematic diagram of the structure of an artificial intelligence chip provided in an embodiment of the present application;

[0032] Figure 2 A flowchart of an optimization method for attention mechanism calculation provided in an embodiment of the present application;

[0033] Figure 3 A computational graph for a forward attention mechanism calculation provided in an embodiment of the present application;

[0034] Figure 4 A schematic diagram of a segmentation result of a computation graph provided in an embodiment of the present application;

[0035] Figure 5 A computational graph for a reverse attention mechanism calculation provided in an embodiment of the present application;

[0036] Figure 6 A schematic diagram of a segmentation result of a computation graph provided in an embodiment of the present application;

[0037] Figure 7 A schematic diagram of a flow chart of a subgraph optimization method provided in an embodiment of the present application;

[0038] Figure 8 A schematic diagram of a flow chart of a subgraph optimization method provided in an embodiment of the present application;

[0039] Fig. 9 A schematic diagram of a flow chart of a subgraph optimization method provided in an embodiment of the present application;

[0040] Fig.10 A schematic diagram of the structure of an optimization device for attention mechanism calculation provided in an embodiment of the present application;

[0041] Fig.11 A schematic diagram of the structure of a computer device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0042] In order to make the purpose, technical solution and beneficial effects of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0043] refer to Figure 1 , which is a structural diagram of an artificial intelligence chip applicable to an embodiment of the present application, the artificial intelligence chip 100 at least includes: a video memory 101, a plurality of execution units 102, wherein each execution unit 102 includes: an on-chip cache 103, a register 104, a tensor calculation unit 105 and a vector calculation unit 106, wherein the tensor calculation unit 105 and the vector calculation unit 106 are heterogeneous cores.

[0044] The video memory 101 may be a high bandwidth memory (HBM) or other types of memory. The on-chip cache 103 is a temporary memory, which has a smaller capacity than the video memory 101 but a faster data exchange speed than the video memory 101 .

[0045] The on-chip cache 103 includes a cache shared by the tensor computing unit 105 and the vector computing unit 106 , and also includes a cache exclusively used by the tensor computing unit 105 or the vector computing unit 106 .

[0046] The register 104 is a register shared by the tensor calculation unit 105 and the vector calculation unit 106. The tensor calculation unit 105 and the vector calculation unit 106 can exchange data through the register 104. The register 104 can be a thread-local register (TLR).

[0047] Compared to the on-chip cache 103 , the register 104 has a smaller capacity than the on-chip cache 103 , but a faster data exchange speed than the on-chip cache 103 .

[0048] In addition to the above-mentioned structure, the artificial intelligence chip 100 in the present application may also include other structures, which are not specifically limited in the present application.

[0049] The artificial intelligence chip 100 can be: a central processing unit (CPU), a graphics processing unit (GPU), a general-purpose graphics processing unit (GPGPU), a domain-specific architecture (DSA), etc.

[0050] In the related technology, in the attention mechanism calculation process, according to the execution order of multiple operators in the calculation graph corresponding to the attention mechanism module, multiple operators are compiled into corresponding kernel functions and executed in turn, wherein each operator corresponds to a kernel function.

[0051] Before each operator is executed, the corresponding kernel function needs to be emitted. Since the computation graph contains many operators and the compiled kernel functions are also many, the kernel function emission time overhead is relatively large during the attention mechanism calculation process.

[0052] In view of this, this application is based on Figure 1The architecture diagram of the artificial intelligence chip shown in the figure provides an optimization method for attention mechanism calculation. Attention mechanism calculation can be applied to various scenarios. For example, image processing scenarios, speech processing scenarios, text processing scenarios, etc. In different application scenarios, the physical meaning of the input data of attention mechanism calculation can be different.

[0053] For example, in speech processing scenarios, the input data for attention mechanism calculations can be speech data used in tasks such as speech enhancement, speech recognition, and speech synthesis. For example, in image processing scenarios, the input data for attention mechanism calculations can be image data used in tasks such as image preprocessing, image segmentation, and object detection. For example, in text processing scenarios, the input data for attention mechanism calculations can be text data used in tasks such as text generation and text recognition.

[0054] The following is a detailed introduction to the optimization method of attention mechanism calculation process, such as Figure 2 As shown, the process of the method is executed by an artificial intelligence chip and includes the following steps:

[0055] Step 201, split the computational graph of the attention mechanism calculation to obtain multiple computational subgraphs.

[0056] Specifically, the calculation graph of the attention mechanism calculation includes multiple operators; according to the positions of the multiple operators in the calculation graph, the multiple operators are executed in sequence to obtain the calculation results of the attention mechanism calculation. When the type of attention mechanism calculation is different, the corresponding calculation graph is also different.

[0057] Multiple templates (patterns) are pre-set. The templates can be some complex fusion operators. For example, the fusion operator of the deformation operator and the reordering operator, the fusion operator of the matrix multiplication operator and the bit-by-bit calculation operator, etc. It can also be a single operator, such as the normalization operator, the reduction operator, etc.; it can also be some special functions, such as the loss function.

[0058] The computational graph of the attention mechanism calculation is segmented by template matching to obtain multiple computational subgraphs and the execution order of multiple computational subgraphs.

[0059] Specifically, starting from the start position of the computation graph, the operators in the computation graph are traversed in sequence according to the execution order of the operators in the computation graph. During the traversal process, the traversed operators are matched with the preset templates, and a computation subgraph is generated according to the matching results.

[0060] In a specific implementation, when traversing content that does not match any template, each operator in the content can be regarded as a computational subgraph; or the operators in the content can be divided into multiple operator sets, the number of operators in each operator set being less than a preset threshold (i.e., the upper limit of the number of operators in a computational subgraph), and then each operator set can be regarded as a computational subgraph.

[0061] When traversing multiple contents that match the same preset template and are executed in parallel, it is determined whether the total number of operators contained in the multiple contents is less than the preset threshold. If so, the multiple contents are treated as a calculation subgraph; otherwise, each of the multiple contents is treated as a calculation subgraph.

[0062] It is also possible to use each time a content matching a preset template is traversed as a calculation subgraph.

[0063] Each computation subgraph includes at least one operator, and the number of computation subgraphs obtained after segmentation is less than the number of operators contained in the computation graph; at the same time, the generation order of multiple computation subgraphs is used as the execution order of multiple computation subgraphs. In a specific implementation, each time a computation subgraph is generated, a corresponding subgraph number can be assigned to the computation subgraph, and the execution order of the computation subgraph is represented by the subgraph number.

[0064] In some embodiments, the computation graph of the attention mechanism calculation is: a computation graph of the forward attention mechanism calculation, or a computation graph of the reverse attention mechanism calculation. The structure of the computation graph of the forward attention mechanism calculation is different from the computation graph of the reverse attention mechanism calculation.

[0065] For example, see Figure 3 , which is a calculation graph for a forward attention mechanism calculation provided in an embodiment of the present application. The operators included in the calculation graph are: a matrix multiplication (Matmul) operator, a deformation (Reshape) operator, a reordering (Permute) operator, a forward flash attention mechanism (flash_attention_fwd) operator, and a linear (linear) operator.

[0066] The process of performing the forward attention mechanism calculation according to this calculation graph is:

[0067] Input the input tensor (hidden_states) and query weight (q_weight) to matrix multiplication operator 1 for calculation, and obtain the first matrix multiplication result (q_matmul_output). Input the first matrix multiplication result to deformation operator 1 to obtain the first deformation result (q_reshape_out); input the first deformation result to reordering operator 1 to obtain the query representation.

[0068] Input the input tensor (hidden_states) and key weight (k_weight) to matrix multiplication operator 2 for calculation, and obtain the second matrix multiplication result (k_matmul_output). Input the second matrix multiplication result to deformation operator 2 to obtain the second deformation result (k_reshape_out); input the second deformation result to reordering operator 2 to obtain the key representation.

[0069] Input the input tensor (hidden_states) and value weight (v_weight) to matrix multiplication operator 3 for calculation, and obtain the third matrix multiplication result (v_matmul_output). Input the third matrix multiplication result to deformation operator 3 to obtain the third deformation result (v_reshape_out); input the third deformation result to reordering operator 3 to obtain value representation.

[0070] Among them, the input tensor can be the input data of the model, or the output result of the previous module in the forward propagation calculation graph of the model; the query representation (query) can be a scalar, vector, multi-dimensional array, etc.; similarly, the key representation (key) can also be a scalar, vector, multi-dimensional array, etc.; the value representation (value) can also be a scalar, vector, multi-dimensional array, etc.

[0071] The query representation, key representation and value representation are input into the forward flash attention mechanism operator for calculation to obtain a preliminary calculation result (attention_output). The preliminary calculation result is input into the reordering operator 4 to obtain the first reordering result (attention_outputs_permute). The first reordering result is input into the deformation operator 4 to obtain the fourth deformation result (qkv).

[0072] Input the fourth deformation result, output weight (o_weight) and bias vector (o_bias) into the linear operator to obtain the target calculation result (proj_output) calculated by the forward attention mechanism.

[0073] The output of the computation graph calculated by the forward attention mechanism includes: target calculation result, query representation, key representation, value representation, preliminary calculation result, and fourth deformation result.

[0074] The upper limit value (ie, preset threshold) of the number of operators in a computation subgraph is set to 7; the preset template 1 is: a fusion operator consisting of a deformation operator and a reordering operator.

[0075] According to the preset template and Figure 3 The structural characteristics of the computational graph will Figure 3 The computational graph of is divided into five computational subgraphs (i.e., forward subgraphs), see Figure 4 , respectively: forward sub Figure 1 , forward Figure 2 , forward Figure 3 , forward Figure 4 , forward Figure 5 .

[0076] In the specific implementation, the traversal starts from the starting position of the calculation graph, and first traverses to matrix multiplication operator 1, matrix multiplication operator 2, and matrix multiplication operator 3. Since matrix multiplication operator 1 does not match template 1, matrix multiplication operator 2 does not match template 1, and matrix multiplication operator 3 does not match template 1; and the total number of matrix multiplication operator 1, matrix multiplication operator 2, and matrix multiplication operator 3 is less than the preset threshold, matrix multiplication operator 1, matrix multiplication operator 2, and matrix multiplication operator 3 are used as forward sub-operators. Figure 1 .

[0077] Then traverse to deformation operator 1, reordering operator 1, deformation operator 2, reordering operator 2, deformation operator 3, and reordering operator 3. Since deformation operator 1 and reordering operator 1 match template 1, deformation operator 2 and reordering operator 2 match template 1, and deformation operator 3 and reordering operator 3 match template 1; at the same time, the total number of deformation operator 1, reordering operator 1, deformation operator 2, reordering operator 2, deformation operator 3, and reordering operator 3 is less than the preset threshold, so deformation operator 1, reordering operator 1, deformation operator 2, reordering operator 2, deformation operator 3, and reordering operator 3 are used as forward sub-operators. Figure 2 .

[0078] Then traverse to the forward flash attention mechanism operator. If the forward flash attention mechanism operator does not match template 1, the forward flash attention mechanism operator is used as the forward sub Figure 3 .

[0079] Then traverse to deformation operator 4 and reordering operator 4. Since deformation operator 4 and reordering operator 4 match template 1, deformation operator 4 and reordering operator 4 are used as forward operators. Figure 4 .

[0080] Then we traverse to the linear operator. Since the linear operator does not match template 1, we use the linear operator as the forward child. Figure 5 .

[0081] Depend on Figure 3 and Figure 4 It can be seen that Figure 3 The computation graph of contains 13 operators, and Figure 3 After the computation graph shown in the figure is split, 5 computation subgraphs are obtained, that is, the number of computation subgraphs contained in the computation graph is less than the number of operators contained in the computation graph. The execution order of the 5 computation subgraphs is: forward subgraphs Figure 1 , forward Figure 2 , forward Figure 3 , forward Figure 4 , forward Figure 5 .

[0082] For example, see Figure 5 , which is a calculation graph for a reverse attention mechanism calculation provided in an embodiment of the present application. The operators included in the calculation graph are: matrix multiplication operator, deformation operator, reordering operator, reverse flash attention mechanism (flash_attention_bwd) operator, reduction (Reduce) operator, and addition (Add) operator.

[0083] According to the calculation graph, the calculation process of the reverse attention mechanism is as follows:

[0084] The output gradient (output_grad) and the output weight (o_weight) are input into the matrix multiplication operator 4 for calculation to obtain the reverse gradient 1 (i.e., the gradient qkv_grad of the fourth deformation result). The reverse gradient 1 is input into the reordering operator 5 to obtain the second reordering result (qkv_grad_permute). The second reordering result is input into the deformation operator 5 to obtain the reverse gradient 2 (i.e., the gradient attention_output_grad of the preliminary calculation result).

[0085] The output gradient (output_grad) and the fourth deformation result are input into the matrix multiplication operator 5 for calculation to obtain the reverse gradient 3 (ie, the gradient o_weight_grad of the output weight).

[0086] The output gradient (output_grad) is input into the reduction operator for calculation to obtain the reverse gradient 4 (i.e., the gradient of the bias vector o_bias_grad).

[0087] Input the reverse gradient 2, query representation, key representation, value representation, preliminary calculation results and forward intermediate results (inter_mediate) into the reverse flash attention mechanism operator for calculation to obtain the first query gradient (query_grad), the first key gradient (key_grad), and the first value gradient (value_grad).

[0088] The first query gradient is input into the deformation operator 6 to obtain the corresponding fifth deformation result (query_grad_reshape); the fifth deformation result is input into the reordering operator 6 to obtain the second query gradient. The second query gradient and the query weight are input into the matrix multiplication operator 6 for calculation to obtain the query input gradient (query_input_grad).

[0089] The first key gradient is input into the deformation operator 7 to obtain the sixth deformation result (key_grad_reshape); the sixth deformation result is input into the reordering operator 7 to obtain the second key gradient. The second key gradient and the key weight are input into the matrix multiplication operator 7 for calculation to obtain the key input gradient (ke_input_grad).

[0090] The first value gradient is input into the deformation operator 8 to obtain the seventh deformation result (value_grad_reshape); the seventh deformation result is input into the reordering operator 8 to obtain the second value gradient. The second value gradient and the value weight are input into the matrix multiplication operator 8 for calculation to obtain the value input gradient (value_input_grad).

[0091] Accumulate the query input gradient and the key input gradient input adder 1 to obtain the first accumulation result (input_grad_pre). Accumulate the first accumulation result and the key input gradient input adder 2 to obtain the gradient of the input tensor (hidden_states_grad).

[0092] The second query gradient and the input tensor are input into the matrix multiplication operator 9 for calculation to obtain the reverse gradient 5 (ie, the gradient q_weight_grad of the query weight).

[0093] The second key gradient and the input tensor are input into the matrix multiplication operator 10 for calculation to obtain the reverse gradient 6 (ie, the gradient key_states_grad of the key weight).

[0094] The second value gradient and the input tensor are input into the matrix multiplication operator 11 to calculate and obtain the reverse gradient 7 (ie, the gradient value_states_grad of the value weight).

[0095] The outputs of the computation graph computed by the inverse attention mechanism include: the gradient of the query weight, the gradient of the key weight, the gradient of the value weight, the gradient of the input tensor, the gradient of the output weight, and the gradient of the bias vector.

[0096] The upper limit of the number of operators in a computation subgraph (i.e., the preset threshold) is set to 7; the preset template 1 is: a fusion operator consisting of a deformation operator and a reordering operator; the preset template 2 is: a fusion operator consisting of a reordering operator and a deformation operator. The preset template 3 is: a fusion operator consisting of a matrix multiplication operator, a matrix multiplication operator, a matrix multiplication operator, an addition operator, and an addition operator.

[0097] According to the preset template and Figure 5 The structural characteristics of the computational graph will Figure 5 The computational graph of is divided into five computational subgraphs (i.e., reverse subgraphs), see Figure 6 , respectively: reverse Figure 1 , inverse sub Figure 2 , inverse sub Figure 3 , inverse sub Figure 4 , inverse sub Figure 5 , inverse sub Figure 6 .

[0098] In the specific implementation, the traversal starts from the starting position of the calculation graph, and first traverses to matrix multiplication operator 4, matrix multiplication operator 5 and reduction operator. Since matrix multiplication operator 4 does not match the above three templates, matrix multiplication operator 5 does not match the above three templates, and the reduction operator does not match the above three templates, and the total number of matrix multiplication operator 4, matrix multiplication operator 5 and reduction operator is less than the preset threshold, matrix multiplication operator 4, matrix multiplication operator 5 and reduction operator are used as reverse operators. Figure 1 .

[0099] Then traverse to reordering operator 5 and deformation operator 5. Since reordering operator 5 and deformation operator 5 match template 2, reordering operator 5 and deformation operator 5 are used as inverse operators. Figure 2 .

[0100] Then traverse to the reverse flash attention mechanism operator. The reverse flash attention mechanism operator does not match the above three templates, so the reverse flash attention mechanism operator is used as the reverse operator. Figure 3 .

[0101] Then traverse to deformation operator 6, reordering operator 6, deformation operator 7, reordering operator 7, deformation operator 8, and reordering operator 8. Since deformation operator 6 and reordering operator 6 match template 1, deformation operator 7 and reordering operator 7 match template 1, and deformation operator 8 and reordering operator 8 match template 1; at the same time, the total number of deformation operator 6, reordering operator 6, deformation operator 7, reordering operator 7, deformation operator 8, and reordering operator 8 is less than the preset threshold, therefore, deformation operator 6, reordering operator 6, deformation operator 7, reordering operator 7, deformation operator 8, and reordering operator 8 are used as reverse operators. Figure 4 .

[0102] Then traverse to matrix multiplication operator 6, matrix multiplication operator 7, matrix multiplication operator 8, addition operator 1 and addition operator 2. Since matrix multiplication operator 6, matrix multiplication operator 7, matrix multiplication operator 8, addition operator 1 and addition operator 2 match template 3, matrix multiplication operator 6, matrix multiplication operator 7, matrix multiplication operator 8, addition operator 1 and addition operator 2 are used as inverse sub- Figure 5 .

[0103] Then traverse to matrix multiplication operator 9, matrix multiplication operator 10, and matrix multiplication operator 11. Since matrix multiplication operator 9 does not match the above three templates, matrix multiplication operator 10 does not match the above three templates, and matrix multiplication operator 11 does not match the above three templates, and the total number of matrix multiplication operator 9, matrix multiplication operator 10, and matrix multiplication operator 11 is less than the preset threshold, matrix multiplication operator 9, matrix multiplication operator 10, and matrix multiplication operator 11 are used as reverse sub- Figure 6 .

[0104] Depend on Figure 5 and Figure 6 It can be seen that Figure 5 The computational graph shown contains 20 operators, and Figure 5 After the computation graph of is split, 6 computation subgraphs are obtained, that is, the number of computation subgraphs contained in the computation graph is less than the number of operators contained in the computation graph. The execution order of the 6 computation subgraphs is: reverse subgraphs Figure 1 , inverse sub Figure 2 , inverse sub Figure 3 , inverse sub Figure 4 , inverse sub Figure 5 , inverse sub Figure 6 .

[0105] Step 202, compile each computational subgraph into a kernel function and execute the kernel function to obtain the target computation result of the attention mechanism calculation.

[0106] Specifically, the central processing unit in the computer device compiles each computational subgraph into a kernel function, where the kernel function is a directly executable machine code. The central processing unit transmits each obtained kernel function to the artificial intelligence chip in the computer device. The artificial intelligence chip executes each received kernel function in turn to obtain the corresponding calculation results, where the calculation result of the last kernel function executed by the artificial intelligence chip is the target calculation result calculated by the attention mechanism.

[0107] In some embodiments, since multiple operators in a computational graph are executed in a specific order, after splitting the computational graph into multiple computational subgraphs, the multiple computational subgraphs also need to be executed in a certain order to ensure the accuracy of the results calculated by the attention mechanism.

[0108] Based on this, in this application, starting from the starting position of the calculation graph, each operator in the calculation graph is traversed in sequence according to the execution order of each operator in the calculation graph. During the traversal process, the traversed operator is matched with the preset template, and a calculation subgraph is generated according to the matching result. After the traversal is completed, the generation order of multiple calculation subgraphs is used as the execution order of multiple calculation subgraphs. In this way, when the calculation graph calculated by the attention mechanism is segmented, not only multiple calculation subgraphs are obtained, but also the execution order of multiple calculation subgraphs is obtained.

[0109] In the specific execution process, according to the execution order, the central processing unit in the computer device compiles each computational subgraph into a corresponding kernel function, transmits the kernel function to the artificial intelligence chip in the computer device, and the artificial intelligence chip executes the kernel function corresponding to each computational subgraph to obtain the target calculation result calculated by the attention mechanism. In addition, the process of obtaining the computational subgraph described in step 201 above is also executed by the central processing unit in the computer device.

[0110] In the specific implementation, for the computation subgraph that ranks first in the execution order, the CPU compiles the computation subgraph into a kernel function, and then sends the kernel function to the artificial intelligence chip for execution. The artificial intelligence chip loads the input data of the computation subgraph from the video memory, executes the kernel function based on the input data, obtains the calculation result of the computation subgraph, and saves the calculation result of the computation subgraph in the video memory.

[0111] For the computation subgraph that ranks second in the execution order, the CPU compiles the computation subgraph into a kernel function, and then sends the kernel function to the AI ​​chip for execution. The AI ​​chip loads the input data of the computation subgraph (including the calculation result of the computation subgraph that ranks first in the execution order) from the video memory, executes the kernel function based on the input data, obtains the calculation result of the computation subgraph, and saves the calculation result of the computation subgraph in the video memory.

[0112] And so on, until the calculation result of the last calculation subgraph in the execution order is obtained, and this calculation result is used as the target calculation result of the attention mechanism calculation.

[0113] In an embodiment of the present application, the calculation graph of the attention mechanism calculation is segmented to obtain multiple calculation subgraphs. Each calculation subgraph is then compiled into a corresponding kernel function, and the kernel function corresponding to each calculation subgraph is executed to obtain the target calculation result of the attention mechanism calculation. Since the number of multiple calculation subgraphs obtained is less than the number of operators contained in the calculation graph, compared to compiling each operator in the calculation graph into a kernel function, the method of compiling each calculation subgraph in the calculation graph into a kernel function reduces the number of compiled kernel functions, thereby reducing the time overhead of kernel function emission during the attention mechanism calculation process, thereby improving the performance of the attention mechanism calculation.

[0114] In some embodiments, each computational subgraph is compiled into a kernel function and the kernel function is executed. Before obtaining the target computational result of the attention mechanism calculation, a subgraph optimization operation is performed on the computational subgraphs among multiple computational subgraphs that meet preset optimization conditions.

[0115] Specifically, the preset optimization conditions can be set according to actual conditions. The following describes the specific process of subgraph optimization operation for several computational subgraphs that meet the preset optimization conditions.

[0116] In some embodiments, if the computation subgraph includes at least one operator combination for layout adjustment, the computation subgraph satisfies the preset optimization condition, and the following subgraph optimization operations are performed on the computation subgraph:

[0117] For at least one operator combination for layout adjustment included in the computation subgraph, the following operations are performed respectively: multiple operators included in an operator combination are replaced with a fused operator having the same function as the operator combination.

[0118] Specifically, the operator combination used for layout adjustment is generally an IO-intensive operator, that is, the computational complexity of this type of operator is very small, mainly involving data reading and writing operations. When multiple operators in the operator combination are executed sequentially, each operator needs to load input data from the video memory, and after performing layout adjustment on the loaded input data, the obtained adjustment result is written out to the video memory. The next operator loads the adjustment result of the previous operator from the video memory, and then performs layout adjustment on the adjustment result again, and writes the obtained adjustment result out to the video memory, and so on, until the adjustment result of the last operator (that is, the output result of the operator combination) is obtained. This results in frequent reading and writing of the video memory during the execution of this operator combination, thereby affecting the computing efficiency.

[0119] In view of this, the present application replaces multiple operators contained in an operator combination with a fusion operator with the same function as the operator combination. In other words, the output result of the operator combination can be obtained by executing a fusion operator. Specifically, the input data is loaded from the video memory, and the loaded input data is input into the fusion operator for layout adjustment to obtain the adjustment result (i.e., the output result of the operator combination), which can reduce the number of times the video memory is read and written, thereby improving the computing efficiency.

[0120] In some embodiments, an operator combination includes: a deformation operator and a reordering operator, and the deformation operator is executed before the reordering operator, or the deformation operator is executed after the reordering operator. This application does not specifically limit the execution order of the deformation operator and the reordering operator in the operator combination.

[0121] Specifically, the deformation operator and the reordering operator are replaced by a fusion operator, and the fusion operator can realize the functions of the deformation operator and the reordering operator.

[0122] For example, see Figure 7 ,exist Figure 4 In the segmentation result of the computational graph shown in the figure, the forward sub Figure 2It includes three operator combinations, starting with the operator combination consisting of deformation operator 1 and reordering operator 1. The operator combination consisting of deformation operator 1 and reordering operator 1 can be replaced with a pre-written fusion operator 1, which can directly complete the deformation + reordering operation without outputting the deformation result. Figure 2 The other two operator combinations in can also be optimized in the same way, which will not be described here.

[0123] For example, see Figure 8 ,exist Figure 4 In the segmentation result of the computational graph shown in the figure, the forward sub Figure 4 Including: an operator combination consisting of a reordering operator 4 and a deformation operator 4. The operator combination consisting of a reordering operator 4 and a deformation operator 4 is replaced with a pre-written fusion operator 2, which can directly complete the reordering + deformation operation without outputting the reordering result.

[0124] In the embodiment of the present application, the operator combination consisting of the deformation operator and the reordering operator is replaced by a fusion operator with the same function as the operator combination, which reduces the number of IO operations and the number of reading and writing video memory, thereby improving computing efficiency.

[0125] In some embodiments, if the computation subgraph includes a tensor core operator and a vector core operator, the computation subgraph satisfies the preset optimization condition, and the following subgraph optimization operations are performed on the computation subgraph:

[0126] Thread-local registers are allocated for data exchange between tensor core operators and vector core operators, where tensor core operators refer to operators executed by tensor compute units, such as matrix multiplication operators. Vector core operators refer to operators executed by vector compute units, such as addition operators. Tensor core operators and vector core operators are executed in parallel.

[0127] Specifically, since the number of thread-local registers is limited, when the interactive data between the tensor core operator and the vector core operator is large, it is difficult to pass all the interactive data at one time through the thread-local register. At this time, the interactive data can be divided into multiple data blocks, and multiple data blocks are passed through the thread-local register in a loop, with one data block passed each time the loop ends until all data blocks are calculated.

[0128] For example, see Fig. 9 ,exist Figure 6 The reverse Figure 5 In the middle, the opposite Figure 5Before optimization, addition operator 1 needs to wait for matrix multiplication operator 6 and matrix multiplication operator 7 to complete all calculations before it can perform calculations. The two can only be performed serially. Moreover, the calculation results of the matrix multiplication operator need to be output to the video memory for storage. Addition operator 1 needs to reload data from the video memory for calculation, resulting in poor performance.

[0129] In view of this, the present application performs the following optimization operations: matrix multiplication operator 6 and matrix multiplication operator 7 are executed by the tensor calculation unit, and addition operator 1 is executed by the vector calculation unit. The tensor calculation unit and the vector calculation unit can be calculated in parallel. In addition, matrix multiplication operator 6 and addition operator 1 exchange data through thread local registers; matrix multiplication operator 7 and addition operator 1 can also exchange data through thread local registers.

[0130] Since the number of thread local registers is limited, the amount of data transferred by the thread local registers each time is limited. Therefore, the matrix multiplication result output by the matrix multiplication operator 6 is divided into multiple data blocks, and the multiple data blocks are transferred to the addition operator 1 through the thread local registers in a loop. Similarly, the matrix multiplication result output by the matrix multiplication operator 7 is divided into multiple data blocks, and the multiple data blocks are transferred to the addition operator 1 through the thread local registers in a loop.

[0131] In each loop, a data block calculated by matrix multiplication operator 6 is saved in the thread local register, and a data block calculated by matrix multiplication operator 7 is saved in the thread local register. Two data blocks are taken out from the thread local register and input into the addition operator for accumulation to obtain an accumulation block. And so on, until all data blocks are calculated, the loop is stopped. The final accumulation result can be obtained by splicing the obtained multiple accumulation blocks.

[0132] In the embodiment of the present application, since the speed of reading and writing operations on thread local registers is much faster than the speed of reading and writing operations on video memory, the data reading and writing performance is greatly improved when the tensor core operator and the vector core operator exchange data through the thread local registers. At the same time, the tensor core operator and the vector core operator are executed in parallel to further improve the computing efficiency.

[0133] Based on the same technical concept, the embodiment of the present application provides a structural diagram of an optimization device for attention mechanism calculation, such as Fig.10 As shown, the optimization device 1000 for calculating the attention mechanism includes:

[0134] A segmentation module 1001 is used to segment the computation graph of the attention mechanism calculation to obtain multiple computation subgraphs, wherein each computation subgraph includes at least one operator, and the number of the multiple computation subgraphs is less than the number of operators included in the computation graph;

[0135] The execution module 1002 is used to compile each computational subgraph into a corresponding kernel function, and execute the kernel function corresponding to each of the computational subgraphs to obtain the target computational result of the attention mechanism calculation.

[0136] Optionally, the segmentation module 1001 is specifically used for:

[0137] The computation graph of the attention mechanism calculation is split to obtain multiple computation subgraphs and the execution order of the multiple computation subgraphs;

[0138] The execution module 1002 is specifically used for:

[0139] According to the execution order, each computational subgraph is compiled into a corresponding kernel function, and the kernel function corresponding to each computational subgraph is executed to obtain the target computational result of the attention mechanism calculation.

[0140] Optionally, the segmentation module 1001 is further used for:

[0141] Each computational subgraph is compiled into a corresponding kernel function, and the kernel function corresponding to each computational subgraph is executed. Before obtaining the target computational result calculated by the attention mechanism, a subgraph optimization operation is performed on the computational subgraphs among the multiple computational subgraphs that meet the preset optimization conditions.

[0142] Optionally, the segmentation module 1001 is specifically used for:

[0143] If the computation subgraph includes at least one operator combination for layout adjustment, the computation subgraph satisfies the preset optimization condition, and the following subgraph optimization operation is performed on the computation subgraph:

[0144] For at least one operator combination for layout adjustment included in the computation subgraph, the following operations are performed respectively: multiple operators included in one operator combination are replaced with a fused operator having the same function as that of the one operator combination.

[0145] Optionally, the one operator combination includes: a deformation operator and a reordering operator, wherein the deformation operator is executed before the reordering operator, or the deformation operator is executed after the reordering operator.

[0146] Optionally, the segmentation module 1001 is specifically used for:

[0147] If the computation subgraph includes a tensor core operator and a vector core operator, the computation subgraph satisfies the preset optimization conditions, and the following subgraph optimization operations are performed on the computation subgraph:

[0148] Thread local registers are allocated for data exchange between the tensor core operator and the vector core operator, wherein the tensor core operator refers to an operator executed by a tensor computing unit, and the vector core operator refers to an operator executed by a vector computing unit.

[0149] Optionally, the tensor core operator and the vector core operator are executed in parallel.

[0150] In an embodiment of the present application, the calculation graph of the attention mechanism calculation is segmented to obtain multiple calculation subgraphs. Each calculation subgraph is then compiled into a corresponding kernel function, and the kernel function corresponding to each calculation subgraph is executed to obtain the target calculation result of the attention mechanism calculation. Since the number of multiple calculation subgraphs obtained is less than the number of operators contained in the calculation graph, compared to compiling each operator in the calculation graph into a kernel function, the method of compiling each calculation subgraph in the calculation graph into a kernel function reduces the number of compiled kernel functions, thereby reducing the time overhead of kernel function emission during the attention mechanism calculation process, thereby improving the performance of the attention mechanism calculation.

[0151] Based on the same technical concept, the embodiment of the present application provides a computer device, such as Fig.11 As shown, it includes at least one artificial intelligence chip 100 and a memory 1101 connected to the at least one artificial intelligence chip 100. The specific connection medium between the artificial intelligence chip 100 and the memory 1101 is not limited in the embodiment of the present application. Fig.11 In the example, the artificial intelligence chip 100 and the memory 1101 are connected via a bus. The bus can be divided into an address bus, a data bus, a control bus, etc.

[0152] In an embodiment of the present application, the memory 1101 stores instructions that can be executed by at least one artificial intelligence chip 100. At least one artificial intelligence chip 100 can execute the steps of the above-mentioned optimization method for attention mechanism calculation by executing the instructions stored in the memory 1101.

[0153] Among them, the artificial intelligence chip 100 is the control center of the computer device, and various interfaces and lines can be used to connect various parts of the computer device, and the attention mechanism calculation can be optimized by running or executing instructions stored in the memory 1101 and calling data stored in the memory 1101. Optionally, the artificial intelligence chip 100 may include one or more processing units, and the artificial intelligence chip 100 may integrate an application processor and a modem processor, wherein the application processor mainly processes the operating system, user interface, and application programs, and the modem processor mainly processes wireless communications. It is understandable that the above-mentioned modem processor may not be integrated into the artificial intelligence chip 100. In some embodiments, the artificial intelligence chip 100 and the memory 1101 can be implemented on the same chip, and in some embodiments, they can also be implemented separately on independent chips.

[0154] The artificial intelligence chip 100 can be a general-purpose processor, such as a central processing unit (CPU), a digital signal processor, an application-specific integrated circuit (ASIC), a field programmable gate array or other programmable logic device, a discrete gate or transistor logic device, a discrete hardware component, and can implement or execute the methods, steps, and logic block diagrams disclosed in the embodiments of the present application. A general-purpose processor can be a microprocessor or any conventional processor, etc. The steps of the method disclosed in the embodiments of the present application can be directly embodied as a hardware processor, or can be executed by a combination of hardware and software modules in the processor.

[0155] Memory 1101, as a non-volatile computer-readable storage medium, can be used to store non-volatile software programs, non-volatile computer executable programs and modules. Memory 1101 may include at least one type of storage medium, such as flash memory, hard disk, multimedia card, card-type memory, random access memory (Random Access Memory, RAM), static random access memory (Static Random Access Memory, SRAM), programmable read-only memory (Programmable Read Only Memory, PROM), read-only memory (Read Only Memory, ROM), electrically erasable programmable read-only memory (Electrically Erasable Programmable Read-Only Memory, EEPROM), magnetic storage, disk, optical disk, etc. Memory 1101 is any other medium that can be used to carry or store a desired program code in the form of an instruction or data structure and can be accessed by a computer device, but is not limited thereto. The memory 1101 in the embodiment of the present application can also be a circuit or any other device that can realize a storage function, for storing program instructions and / or data.

[0156] Based on the same inventive concept, an embodiment of the present application provides a computer-readable storage medium, which stores a computer program that can be executed by a computer device. When the computer program runs on the computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0157] Based on the same inventive concept, an embodiment of the present application provides a computer program product, which includes a computer program stored on a computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer device, the computer device executes the steps of the optimization method for the above-mentioned attention mechanism calculation.

[0158] It will be appreciated by those skilled in the art that embodiments of the present invention may be provided as methods or computer program products. Therefore, the present invention may take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Moreover, 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 disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0159] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks 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 device or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0160] These computer program instructions may also be stored in a computer readable memory capable of directing a computer device or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer readable memory produce a manufactured product including an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0161] These computer program instructions can also be loaded onto a computer device or other programmable data processing device so that a series of operation steps are executed on the computer device or other programmable device to produce a process implemented by the computer device, thereby providing instructions for implementing the process. Figure 1 A process or multiple processes and / or boxes Figure 1 The steps for the functions specified in one or more boxes.

[0162] Although the preferred embodiments of the present invention have been described, those skilled in the art may make other changes and modifications to these embodiments once they have learned the basic creative concept. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications that fall within the scope of the present invention.

[0163] Obviously, those skilled in the art can make various changes and modifications to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalents, the present invention is also intended to include these modifications and variations.

Claims

1. A method for optimizing attention mechanism calculation, characterized in that: include: Splitting the computation graph of the attention mechanism calculation to obtain a plurality of computation subgraphs, wherein each computation subgraph includes at least one operator, and the number of the plurality of computation subgraphs is less than the number of operators included in the computation graph; Compile each computational subgraph into a corresponding kernel function, and execute the kernel function corresponding to each computational subgraph to obtain the target computational result of the attention mechanism calculation.

2. The method according to claim 1, characterized in that The computation graph of the attention mechanism calculation is segmented to obtain multiple computation subgraphs, including: The computation graph of the attention mechanism calculation is split to obtain multiple computation subgraphs and the execution order of the multiple computation subgraphs; The step of compiling each computational subgraph into a corresponding kernel function and executing the kernel function corresponding to each computational subgraph to obtain the target computation result of the attention mechanism calculation includes: According to the execution order, each computational subgraph is compiled into a corresponding kernel function, and the kernel function corresponding to each computational subgraph is executed to obtain the target computational result of the attention mechanism calculation.

3. The method according to claim 1, characterized in that Before compiling each computation subgraph into a corresponding kernel function and executing the kernel function corresponding to each computation subgraph to obtain the target computation result of the attention mechanism calculation, the method further includes: For the computation subgraphs that meet the preset optimization conditions among the multiple computation subgraphs, a subgraph optimization operation is performed.

4. The method according to claim 3, characterized in that The performing of subgraph optimization operation on the computation subgraphs that meet the preset optimization conditions among the multiple computation subgraphs includes: If the computation subgraph includes at least one operator combination for layout adjustment, the computation subgraph satisfies the preset optimization condition, and the following subgraph optimization operation is performed on the computation subgraph: For at least one operator combination for layout adjustment included in the computation subgraph, the following operations are performed respectively: multiple operators included in one operator combination are replaced with a fused operator having the same function as that of the one operator combination.

5. The method according to claim 4, characterized in that The operator combination includes: a deformation operator and a reordering operator, wherein the deformation operator is executed before the reordering operator, or the deformation operator is executed after the reordering operator.

6. The method according to claim 3, characterized in that The performing of subgraph optimization operation on the computation subgraphs that meet the preset optimization conditions among the multiple computation subgraphs includes: If the computation subgraph includes a tensor core operator and a vector core operator, the computation subgraph satisfies the preset optimization conditions, and the following subgraph optimization operations are performed on the computation subgraph: Thread local registers are allocated for data exchange between the tensor core operator and the vector core operator, wherein the tensor core operator refers to an operator executed by a tensor computing unit, and the vector core operator refers to an operator executed by a vector computing unit.

7. The method according to claim 6, characterized in that The tensor core operator and the vector core operator are executed in parallel.

8. A computer device comprising a memory, an artificial intelligence chip, and a computer program stored in the memory and executable on the artificial intelligence chip, characterized in that: When the artificial intelligence chip executes the computer program, the steps of the method as claimed in any one of claims 1 to 7 are implemented.

9. A computer-readable storage medium, characterized in that: It stores a computer program executable by a computer device. When the computer program is run on the computer device, the computer device executes the steps of the method according to any one of claims 1 to 7.

10. A computer program product, characterized in that The computer program product comprises a computer program stored on a computer-readable storage medium, wherein the computer program comprises program instructions, and when the program instructions are executed by a computer device, the computer device is caused to execute the steps of the method according to any one of claims 1 to 7.

Images