Automatic program optimization method, device and storage medium

Abstract

The application relates to an automatic program optimization method and device based on a flow network generation model and a storage medium, wherein the method comprises the following steps: S1, an initial tensor program is acquired, a plurality of candidate tensor programs with the same logical function are obtained through calculation graph extraction, subgraph segmentation and transformation based on the acquired initial tensor program, and a data set composed of a plurality of samples is constructed based on the obtained candidate tensor programs, wherein the sample comprises a binary tuple composed of a calculation subgraph and a candidate tensor program, and the hardware execution time corresponding to the binary tuple; S2, a plurality of samples are selected from the data set, and a GFlowNet sampling model is trained offline; and S3, the trained GFlowNet sampling model is used to optimize a program to be optimized. Compared with the prior art, the GFlowNet adopts a non-iterative probabilistic sampling strategy, generates a plurality of high-performance programs, apportions the sampling overhead from one mode to another mode, and accelerates the convergence speed.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of program optimization, and in particular to an automatic program optimization method and device based on a flow network generation model and a storage medium. BACKGROUND

[0002] With the increasing popularity of deep learning in daily life, the deployment environment of deep learning is becoming more and more diverse. Due to the large differences between the hardware architectures of various deployment environments (such as smart wearable devices, mobile terminals, desktop terminals, server terminals, and cloud terminals), the direct migration effect of models in different environments is poor. Traditional manual program optimization is high in cost and time-consuming, and cannot cope with the increasingly diverse deployment environment of deep learning.

[0003] Automatic program optimization has the advantages of automation, low cost, and short time consumption, and is suitable for the increasingly diverse deployment environment. In automatic program optimization, there are two core elements that determine the final optimization performance, one is the search space, and the other is the search strategy. The search space is a set of tensor programs with the same logical function but different actual execution times, and the search strategy is to find the optimal tensor program in the given search space.

[0004] Currently, the search strategy is divided into three categories, including machine learning search (such as bundle search, simulated annealing, genetic algorithm, etc.), reinforcement learning search (such as Q-Learning, etc.), and Monte Carlo search (such as Monte Carlo tree search, etc.). Traditional machine learning adopts a learning-based online search strategy, which has the problems of low efficiency, low performance, and long time consumption. The reinforcement learning strategy can only search for a single optimal program with the maximum return, and has the problems of low efficiency and poor diversity. The Monte Carlo search adopts an iterative method, which cannot handle the search between well-separated patterns in high-dimensional space, and has the problems of low efficiency, poor diversity, and long time consumption. SUMMARY

[0005] The purpose of the present application is to provide an automatic program optimization method and device based on a flow network generation model and a storage medium, which adopts a non-iterative probabilistic sampling strategy of GFlowNet to generate multiple high-performance programs, evenly distributes the sampling overhead from one pattern to another, speeds up the convergence speed, and has the characteristics of diversity, avoiding overfitting.

[0006] The purpose of the present application can be achieved by the following technical solutions:

[0007] An automatic program optimization method based on a flow network generation model, comprising:

[0008] Step S1: obtaining an initial tensor program, and based on the obtained initial tensor program, obtaining a plurality of candidate tensor programs with the same logical function through computational graph extraction, subgraph segmentation and transformation, and constructing a data set composed of a plurality of samples based on the obtained candidate tensor programs, wherein the sample includes a binary tuple composed of a computational subgraph and a candidate tensor program, and a hardware execution time corresponding to the binary tuple;

[0009] Step S2: selecting a plurality of samples from the data set, and training a GFlowNet sampling model offline, wherein the training process of a single sample includes:

[0010] initializing the parameters that can be optimized in the candidate tensor program as undetermined parameters, and taking the program corresponding to the undetermined parameters with different values as the state of sampling, and taking the position and value of the determined undetermined parameters as the action of sampling;

[0011] using the GFlowNet sampling model to obtain a plurality of programs through multiple samplings until all the undetermined parameters are determined as specific values, and obtaining the predicted value of the cost of each program through the trained cost model;

[0012] optimizing the GFlowNet sampling model based on the predicted value of the cost of each program;

[0013] Step S3: using the trained GFlowNet sampling model to optimize the program to be optimized.

[0014] The step S1 specifically includes:

[0015] Step S11: obtaining an initial tensor program, and extracting a computational graph based on the obtained initial tensor program;

[0016] Step S12: segmenting the extracted computational graph into a plurality of computational subgraphs according to operation instructions, wherein the operation instructions are composed of one or more basic operations;

[0017] Step S13: transforming each computational subgraph using a conversion operator to obtain a plurality of candidate tensor programs with the same logical function;

[0018] Step S14: based on each candidate tensor program of each computational subgraph, constructing a data set composed of a plurality of samples.

[0019] The basic operation at least includes matrix addition, matrix multiplication, convolution, linear rectifier function operation and normalization exponential function operation.

[0020] The conversion operator at least includes cutting, rearranging, fusing, parallelizing and vectorizing.

[0021] An input of the cost model is normalized features, and an output is hardware execution time, wherein the normalized features are extracted based on a candidate tensor program in a sample.

[0022] The normalized features are obtained by:

[0023] Based on the candidate tensor program, a scheduling primitive is extracted, wherein the scheduling primitive is content corresponding to a conversion operator;

[0024] Based on the extracted scheduling primitive, a conversion operation feature is parsed, wherein the conversion operation feature includes a conversion operator, and a numerical and character identifier corresponding to the conversion operator;

[0025] The conversion operation feature is normalized to obtain normalized features.

[0026] The predicted value of the cost of each program optimizes the GFlowNet sampling model, specifically: the predicted value of the cost of each program optimizes the GFlowNet sampling model so that the sampling probability of each state of the GFlowNet sampling model is proportional to the predicted value of the cost of the final program.

[0027] The step S3 specifically includes:

[0028] Step S31: initializing the parameters that can be optimized in the program to be optimized as undetermined parameters;

[0029] Step S32: using the trained GFlowNet sampling model to sample multiple times to obtain multiple programs;

[0030] Step S33: obtaining the predicted value of the cost of each program based on the cost model;

[0031] Step S34: selecting the program with the smallest predicted value of the cost as the optimized result.

[0032] An automatic program optimization device based on a flow network generation model, comprising a memory, a processor, and a program stored in the memory, wherein the processor implements the method as described above when executing the program.

[0033] A storage medium having a program stored thereon, wherein the program is executed to implement the method as described above.

[0034] Compared with the prior art, the present application has the following beneficial effects:

[0035] 1. GFlowNet adopts a non-iterative probability sampling strategy to generate multiple high-performance programs, which can evenly distribute the sampling overhead from one mode to another mode, accelerate the convergence speed, and have the characteristics of diversity, thereby avoiding overfitting.

[0036] 2、The technology trains GFlowNet offline, then uses GFlowNet to sample to generate multiple high-performance tensor programs, and selects the optimal program from them, thereby obtaining an optimal program with high performance, diversity, and robustness.

[0037] 3、Strong network generalization, such as GPT, BERT, ResNet, and MoE network models.

[0038] 4、Can be applied to various hardware architecture platforms, with strong hardware generalization, such as CPU, GPU, TPU, and NPU AI dedicated accelerators. BRIEF DESCRIPTION OF DRAWINGS

[0039] Figure 1 The main step flowchart of the method of the present application is shown in the figure.

[0040] Figure 2 The technical route of automatic program optimization is shown in the figure.

[0041] Figure 3 The schematic diagram of constructing a data set is shown in the figure.

[0042] Figure 4 The schematic diagram of offline training of a cost model is shown in the figure.

[0043] Figure 5 The schematic diagram of offline training of a GFlowNet sampling model is shown in the figure.

[0044] Figure 6 The schematic diagram of GFlowNet sampling to generate an optimal program is shown in the figure. DETAILED DESCRIPTION

[0045] The present application will be described in detail below in conjunction with the drawings and specific embodiments. The present embodiment is implemented on the basis of the technical solution of the present application, and detailed implementation and specific operation processes are given, but the protection scope of the present application is not limited to the following embodiments.

[0046] As shown in Figure 2 , there are two core elements in automatic program optimization that determine the final optimization performance, one is the search space, and the other is the search strategy. The search space is a set of tensor programs with the same logical function but different actual execution times, and the search strategy is to find the optimal tensor program in the given search space.

[0047] To solve the problems existing in the prior art, specifically, to realize offline and low time consumption, we construct a dataset for offline training of the cost model and the GFlowNet probability sampling model. For different tensor programs, GFlowNet only needs to sample once to generate high-performance programs without online learning. GFlowNet uses non-iterative probability sampling to generate multiple high-performance programs and selects the optimal program from them, thereby realizing the optimal program with high performance, diversity and robustness. Therefore, we design the following four stages:

[0048] (1) Constructing a dataset

[0049] (2) Offline training of the cost model

[0050] (3) Offline training of the GFlowNet sampling model

[0051] (4) GFlowNet sampling to generate the optimal program

[0052] Specifically, an automatic program optimization method based on a flow network generation model, as shown in Figure 1 , comprises:

[0053] Step S1: obtaining an initial tensor program, and based on the obtained initial tensor program, a plurality of candidate tensor programs with the same logical function are obtained after computational graph extraction, subgraph segmentation and transformation, and a dataset composed of a plurality of samples is constructed based on the obtained candidate tensor programs, wherein the sample comprises a binary tuple composed of a computational subgraph and a candidate tensor program, and a hardware execution time corresponding to the binary tuple;

[0054] To realize offline and low time consumption, we need to construct the corresponding dataset and train the offline GFlowNet automatic program optimization search strategy. As shown in Figure 2 , the corresponding training dataset is constructed, and the specific steps are as follows.

[0055] (1): obtaining an initial tensor program, and based on the obtained initial tensor program, extracting a computational graph, specifically, extracting network structure information, i.e. computational graph, from the initial tensor program according to certain rules. As shown in Figure 3As shown, the final extracted computational graphs include BERT-Tiny, BERT-Base, BERT-Medium, BERT-Large, DCGAN, DenseNet, Inception V3, MobileNet V3, ResNet50, ResNext50, and Wide ResNet50, etc. Each computational graph further includes various different sizes, such as BERT-Base-1, 64, BERT-Base-1, 128, BERT-Base-4, 64, and BERT-Base-4, 128, etc.

[0056] (2) The extracted computational graphs are divided into multiple computational subgraphs according to operation instructions, wherein the operation instructions are composed of one or more basic operations. Specifically, the extracted computational graphs are divided into various subgraphs according to certain rules. For example, Figure 3 As shown, the subgraphs include various basic operations such as matrix addition, matrix multiplication, convolution, linear rectifier function operation (ReLU), and normalization exponential function operation (Softmax). In addition, the subgraphs also include fusion operations of various basic operations, such as matrix multiplication + ReLU, matrix multiplication + Softmax, and convolution + matrix addition, etc. Similarly, a subgraph also includes various different sizes, such as 1x64, 1x256, 4x64, and 8x128, etc. of matrix multiplication.

[0057] (3) For various subgraphs, different numbers and orders of conversion operators are applied to obtain multiple candidate tensor programs according to certain rules. For example, Figure 3 As shown, for the same matrix multiplication, different conversion operators such as split, reorder, fuse, parallel, and vectorize are applied in a certain order to obtain multiple tensor programs with the same logical function but different actual execution times.

[0058] (4) For a specific subgraph and a specific candidate tensor program, they are combined to form a final executable program. We execute the corresponding executable program by using the target hardware and obtain the execution time of the hardware. Finally, a corresponding data set is constructed, with the input being the specific subgraph and the specific program, and the label being the execution time of the hardware.

[0059] Step S2: Select multiple samples from the data set to train the GFlowNet sampling model offline, wherein the training process of a single sample includes:

[0060] The parameters that can be optimized in the candidate tensor program are initialized as undetermined parameters, and the programs corresponding to different values of the undetermined parameters are regarded as sampled states, and the positions and values of the undetermined parameters are determined as sampled actions;

[0061] After the GFlowNet sampling model is trained, a plurality of programs are obtained after multiple samplings until all the undetermined parameters are determined as specific values, and the predicted values of the costs of the programs are obtained by using the trained cost model;

[0062] The GFlowNet sampling model is optimized based on the predicted values of the costs of the programs;

[0063] Specifically, since the traditional learning-based online search strategy needs to search in the search space every time for different input tensor programs to obtain a better tensor program, the efficiency of multiple executions is low, the performance is low, the time consumption is long, and the diversity is poor. GFlowNet is a network flow model based on probabilistic sampling. The model uses a non-iterative probabilistic sampling strategy to abstract the programs of various undetermined parameters as states and to abstract the positions and values of the undetermined parameters as actions. Therefore, as shown in Figure 5 , we use the dataset and the offline trained cost model to train the GFlowNet sampling model offline. For different tensor programs, the GFlowNet sampling is used to generate a plurality of high-performance tensor programs, and the optimal program is selected from the plurality of high-performance tensor programs, so that the optimal program with high performance, diversity and robustness is realized. The specific operation is as follows.

[0064] (1) First, a specific program of a specific subgraph is obtained from the dataset, and then the parameters that can be optimized are initialized as undetermined parameters according to the program. As shown in Figure 5 , the initialized state contains various undetermined parameters, such as the loop values of the loop variables in the for loop.

[0065] (2) Then, the GFlowNet predicts the probabilities of different actions, and then the GFlowNet obtains the next action corresponding to the current state according to the probabilistic sampling, which determines the positions and values of the undetermined parameters. As shown in Figure 5 , the first sampling of state 0 obtains state 1,1, which determines the loop value of loop variable j0 as 8, and the mth sampling of state 0 obtains state 1,m, which determines the loop value of loop variable i1 as 32. Finally, after multiple samplings, the final state, i.e., the candidate tensor program, is obtained, in which all the undetermined parameters become specific parameters.

[0066] (3) Finally, as shown in Figure 5As shown, the corresponding subgraph is obtained from the dataset, and then the normalized features are extracted from the subgraph and the sampled program, and input into the cost model to obtain the corresponding prediction value. The GFlowNet probability sampling strategy is adjusted according to the prediction value, so that the sampling probability of each state of GFlowNet is proportional to the prediction value of the final program. Therefore, the final program of GFlowNet probability sampling has characteristics such as high performance, diversity and robustness.

[0067] wherein the input of the cost model is the normalized feature, and the output is the hardware execution time, wherein the normalized feature is extracted based on the candidate tensor program in the sample, and the acquisition process of the normalized feature includes:

[0068] Based on the candidate tensor program, the scheduling primitive is extracted, wherein the scheduling primitive is the content corresponding to the conversion operator;

[0069] Based on the extracted scheduling primitive, the conversion operation feature is parsed, wherein the conversion operation feature includes the conversion operator, and the numerical and character identifiers corresponding to the conversion operator;

[0070] The conversion operation feature is normalized to obtain the normalized feature.

[0071] Specifically, since it is very time-consuming to obtain the execution time of the actual hardware running program, we use the cost model to evaluate the execution time of the program, which greatly shortens the program evaluation time. As shown, Figure 4 we build a dataset to train the cost model offline, which is used to efficiently and quickly evaluate the execution time. The specific operation is shown as follows.

[0072] (1) First, the subgraph and the corresponding program are obtained from the dataset, and then the scheduling primitive is extracted using the preprocessor. As shown, Figure 4 the scheduling primitive contains various conversion operators, such as split, reorder, fuse, parallel, and vectorize, etc. For the same subgraph, different programs correspond to different scheduling primitives.

[0073] (2) For the extracted scheduling primitive, the extractor is used to further extract the corresponding features. As shown, Figure 3 after extraction, only three types of core features are retained, including conversion operators (such as split, reorder, and fuse, etc.), numerical and character identifiers (such as i, j, i.0, and j.1, etc.).

[0074] (3) For the extracted features, the features are normalized by the postprocessor to obtain the normalized features. As shown, Figure 3As shown, the conversion operator is converted into a corresponding one-hot code using a one-hot encoding table, such as split into [1, 0, 0, 0…], reorder into [0, 1, 0, 0…], and the like, and then the character identifier is converted into a corresponding numerical code using a character label table, such as i into 1, i.0 into 2, j into 6, j.1 into 8, and the like.

[0075] (4) Finally, as shown in Figure 4 , the cost model inputs the normalized features and outputs the corresponding predicted value for different hardware platforms. During the entire offline training process, we calculate the loss of the predicted value and the label, and use the loss to supervise the prediction result of the cost model, so that the predicted value of the cost model and the label loss are minimized.

[0076] Step S3: using the trained GFlowNet sampling model to optimize the program to be optimized, specifically including:

[0077] Step S31: initializing the parameters that can be optimized in the program to be optimized as undetermined parameters;

[0078] Step S32: using the trained GFlowNet sampling model to sample multiple times to obtain multiple programs;

[0079] Step S33: obtaining the predicted value of the cost of each program based on the cost model;

[0080] Step S34: selecting the program with the smallest predicted value of the cost as the optimized result.

[0081] As shown in Figure 6 , in the inference stage, we load the GFlowNet sampling model, input the initial program, and GFlowNet samples n times to generate n high-performance programs, and then selects the optimal program. Compared with the online search strategy based on learning, this offline GFlowNet probability sampling strategy only needs to sample multiple times at one time to generate multiple high-performance tensor programs, and select the optimal program from them, thereby realizing the optimal program with high performance, diversity and robustness.

[0082] If the above functions are realized in the form of software function units and sold or used as independent products, they can be stored in a computer readable storage medium. Based on such understanding, the technical solutions of the present application, essentially or in part, or parts of the technical solutions can be embodied in the form of a software product. The computer software product is stored in a storage medium, including several instructions for causing a computer device (which can be a personal computer, a server, or a network device, etc.) to perform all or part of the steps of the methods described in the various embodiments of the present application. The foregoing storage medium includes: U disk, mobile hard disk, read-only memory (ROM, Read-Only Memory), random access memory (RAM, Random Access Memory), magnetic disk or optical disk, and various media that can store program codes.

Claims

1. An apparatus for automatic program optimization based on a flow network generation model, comprising a memory, a processor, and a program stored in the memory, wherein, The processor implements the following when executing the program: Step S1: obtaining an initial tensor program, and based on the obtained initial tensor program, obtaining a plurality of candidate tensor programs with the same logical function through computational graph extraction, subgraph segmentation and transformation, and constructing a data set composed of a plurality of samples based on the obtained candidate tensor programs, wherein the sample includes a binary tuple composed of a computational subgraph and a candidate tensor program, and a hardware execution time corresponding to the binary tuple; Step S2: selecting a plurality of samples from the data set, and training the GFlowNet sampling model offline, wherein the training process of a single sample includes: initializing the parameters that can be optimized in the candidate tensor program as undetermined parameters, and taking the program corresponding to the undetermined parameters with different values as the state of sampling, and taking the position and value of the determined undetermined parameters as the action of sampling; using the GFlowNet sampling model to obtain a plurality of programs through multiple samplings until all undetermined parameters are determined as specific values, and obtaining the predicted value of the cost of each program through the trained cost model; optimizing the GFlowNet sampling model based on the predicted value of the cost of each program; Step S3: optimizing the program to be optimized using the trained GFlowNet sampling model; The step S1 specifically includes: Step S11: obtaining an initial tensor program, and extracting a computational graph based on the obtained initial tensor program; Step S12: dividing the extracted computational graph into a plurality of computational subgraphs according to operation instructions, wherein the operation instructions are composed of one or more basic operations; Step S13: transforming each computational subgraph using a conversion operator to obtain a plurality of candidate tensor programs with the same logical function; Step S14: constructing a data set composed of a plurality of samples based on each candidate tensor program of each computational subgraph; The basic operation at least includes matrix addition, matrix multiplication, convolution, linear rectifier function operation and normalization exponential function operation; The conversion operator at least includes segmentation, rearrangement, fusion, parallelization and vectorization; The input of the cost model is normalized features, and the output is hardware execution time, wherein the normalized features are extracted based on the candidate tensor program in the sample; The process of obtaining the normalized features includes: Based on the candidate tensor program, the scheduling primitives are extracted, wherein the scheduling primitives are the contents corresponding to the conversion operator; Based on the extracted scheduling primitives, the conversion operation features are parsed, wherein the conversion operation features include the conversion operator, and the numerical and character identifiers corresponding to the conversion operator; The conversion operation features are normalized to obtain normalized features, including: converting the conversion operator into corresponding one-hot encoding using a one-hot encoding table, and converting the character identifier into corresponding numerical encoding using a character marker table.

2. The automatic program optimization apparatus based on a flow network generation model according to claim 1, wherein, The GFlowNet sampling model is optimized based on the predicted value of the cost of each program, specifically: the GFlowNet sampling model is optimized based on the predicted value of the cost of each program so that the sampling probability of each state of the GFlowNet sampling model is proportional to the predicted value of the cost of the final program.

3. The apparatus of claim 1, wherein, The step S3 specifically includes: Step S31: initialize the parameters that can be optimized in the program to be optimized as undetermined parameters; Step S32: use the trained GFlowNet sampling model to perform multiple samplings to obtain multiple programs; Step S33: obtain the predicted value of the cost of each program based on the cost model; Step S34: select the program with the smallest predicted value of the cost as the result of optimization.

Images