A Sub-model-based Scheduling Method for AI Compilers

Abstract

The present invention discloses a sub-model-based scheduling method for an AI compiler. The present invention encodes the structural information output by a pre-trained ASTPB model and applies it to a TVM model, so that the TVM model can take into account the structural characteristics of the sub-model during the scheduling of the sub-model. At the same time, the present invention trains a multi-head self-attention model and a feedforward network with a smaller number of parameters, so as to quickly obtain the predicted running time. Through the above two points, the parameter scheduling of the sub-model can be realized more quickly, so that the running time of the sub-model is appropriate and the scheduling efficiency is improved.

Description

Technical Field

[0001] The present invention belongs to the field of artificial intelligence, and in particular relates to a sub-model-based scheduling method for an AI compiler. Background Art

[0002] In recent years, deep learning has been widely used in industrial production environments. Before deploying a deep learning model, performance tuning to adapt the model to specific hardware is an indispensable step. In the early stages of machine learning model development, such a task is often manually optimized by experienced developers, a process that includes control flow orchestration, manual fusion of operators, transformation of equivalent computing techniques, and so on.

[0003] With the continuous development of AI, a more specialized tool such as deep learning compiler (AI compiler) has many available engineering implementations. Given the diversity of hardware, the main task of deep learning compilers is to effectively map calculations to a variety of hardware through multi-level intermediate representations (IR) and implement a variety of optimization methods in the process. Common deep learning compilers include TVM, TensorFlow, MLIR, Pytorch, etc.

[0004] In a traditional compilation process, the program is first compiled into an abstract syntax tree (AST), and then converted into an intermediate representation through a series of processing. In the AI ​​compiler compilation process, the code is first converted into an abstract syntax tree. Since the computational process of a machine learning model is essentially a directed acyclic graph (DAG), after the AI ​​compiler processes the AST of part of the model code, the model code part will eventually be converted into one or more DAGs, which provide a more comprehensive computational representation of the model.

[0005] According to the development of AI compilers, the history of AI compilers can be divided into several stages. The early AI compilers were mainly focused on upper-level implementations. During this period, dynamic graphs and static graphs were separated from each other, and computational graphs and operators were also separated from each other. The boundaries between graphs and operators were fixed, and the optimization methods were very limited. The computational graphs at this stage used expression optimization of conventional compilers, and there were some simple operator fusions at the operator level. After a period of development, the boundaries between graphs and operators were opened, and more optimization methods could be implemented in terms of functions. The architecture also subdivided the dedicated framework of the AI ​​compiler. Based on the previous foundation, the current AI compiler is working on the combination of dynamic graphs and static graphs, unified expression of graphs and operators, automatic parallel optimization, automatic scheduling, and other issues.

[0006] In actual production environments, after engineers have trained a usable neural network, an important step is still required, namely, optimization for specific hardware. The optimized model can often achieve a speedup ratio of at least 1.2 in execution efficiency compared to the unoptimized model.

[0007] This process is usually manually optimized by experienced engineers and requires a lot of attempts and work. In addition, commonly used models such as deep neural networks (DNNs), which play an important role in autonomous driving, augmented reality, translation, and other artificial intelligence applications, have proprietary kernel libraries provided by vendors (such as cuDNN, MKL-DNN, etc.). However, these kernel libraries require a lot of engineering work to manually tune for different hardware platforms and operators, which limits the development of new operators and dedicated accelerators. On the one hand, the AI ​​compiler utilizes various optimization methods of traditional compilers, and on the other hand, it implements common optimization methods for neural networks and implements automatic search solutions to find potential optimization methods.

[0008] This automatic search can be further divided into tuning at the operator level and tuning at the subgraph and the entire network. But the common point is that as the size of the network increases, the number of searches (search space) required will increase exponentially. Therefore, in actual engineering, subgraphs are often divided before searching to separate multiple sub-search tasks. But even if this is done, the search time required is still quite long. Take the common resnet-18 network as an example, it contains 18 convolutional layers and 11.7M parameters. It takes more than 15 hours to use TVM's auto-schedule tuning search on such a small network under the hardware conditions of i7-12700 and Nvidia GTX3090.

[0009] The long-term optimization poses a great challenge to the project cycle and timely deployment. Summary of the invention

[0010] The present invention provides a sub-model-based scheduling method for an AI compiler, which can improve the search efficiency and speed during AI compilation tuning, so as to achieve the purpose of reducing the total tuning consumption time.

[0011] The present invention provides a sub-model-based scheduling method for an AI compiler, comprising:

[0012] S1. Use TVM to split the target model into multiple sub-models and corresponding Sketches. Select the first set of bodies based on the Sketches corresponding to the sub-models to build the initial population. Send the initial population to the hardware for execution to obtain the actual running time.

[0013] S2. Construct a cost model, wherein the cost model includes a pre-trained ASTPB model with frozen parameters, and a multi-head self-attention model and a feedforward network to be trained, obtain corresponding structural information encoding based on the selected sub-model through the pre-trained ASTPB model, obtain attention features based on the numerical attribute features of the first individual set through the multi-head self-attention model, fuse the attention features and the structural information encoding to obtain fused features, obtain the predicted running time based on the fused features through the feedforward network, and train the cost model through the loss function based on the predicted running time and the actual running time;

[0014] S3, predicting the predicted running time of each individual in the newly generated second individual set through the trained cost model, selecting individuals from the second individual set based on the predicted running time to construct a new population, and sending the new population to the hardware to execute the actual running time;

[0015] S4. Repeat steps S2-S3 until the number of repetitions exceeds the threshold or the actual running time is lower than the running time threshold, stop the repetition, and obtain the parameter combination of the sub-model to complete the scheduling of the sub-model.

[0016] Preferably, the multi-head self-attention model and feedforward network of the cost model in step S2 are replaced with the lightGBM model, the numerical attribute features and structural information encoding of the selected sub-model are input into the lightGBM model to obtain the predicted running time, and the cost model is trained through the loss function based on the predicted running time and the actual running time.

[0017] Preferably, selecting individuals from the second set of individuals to construct a new population based on the predicted running time includes:

[0018] S3.1, sorting the second individual set based on the predicted running time to obtain an individual sequence, and selecting multiple individuals from the individual sequence through the individual generation strategy to construct a new population;

[0019] The individual generation strategy includes:

[0020] constructing a selection factor based on the predicted run time standard deviation, wherein the selection factor increases as the value of the predicted run time standard deviation decreases;

[0021] Each time an individual is selected, m individuals are first selected from the individual sequence, and then n individuals are skipped until there are no individuals in the individual sequence that can be selected after o selections, and a new set of individuals is obtained;

[0022] The number of individuals m selected each time is based on the selection factor, the number of individuals, the proportion of individuals selected in the current time, and the number of times the current selection is made;

[0023] The number of individuals n skipped each time is based on the selection factor, the number of individuals, 1-the proportion of individuals selected in the current time, and the number of current selections;

[0024] The m value and n value both decrease as the value of the selection factor increases;

[0025] S3.2. When the new set of individuals is not lower than the threshold number of individuals for constructing a population, a new population is obtained. When the new set of individuals is lower than the threshold number of individuals for constructing a population, step S3.1 is repeated based on the predicted running time of each individual in the regenerated third set of individuals to screen out new individuals from the third set of individuals, and the new individuals are added to the new set of individuals to obtain a new population.

[0026] Preferably, the pre-trained ASTPB model includes:

[0027] The multiple submission programs on the PTA website are used as data sets, and the correctness of the submitted programs is used as the true value;

[0028] Parse the submitted program into an abstract syntax tree, use word2vec to encode the semantic information of each node of the abstract syntax tree to obtain first encoding information, and perform weighted averaging on the data of each node of the first encoding information to obtain second encoding information;

[0029] Encode the position information of each node to obtain a position vector, and add the position vector and the second encoded information to obtain a vector representation of the tree structure, wherein the position information includes the depth of the node in the abstract syntax tree, the position of the node in the sibling nodes at the same depth, and the position of the parent node of the node in the tree structure;

[0030] The ASTPB model includes a convolution module, a residual module, an encoder, a decoder, a connector and a simple neural network, and the simple neural network includes a Sigmoid activation function, a linear layer and a Softmax activation function;

[0031] A new tree is obtained by hierarchically fusing the vector representation of the tree structure through a convolution module, and each node of the new tree is assembled into a node sequence through routing traversal. The node sequence is compressed through a residual module to obtain a fused position vector, and the fused position vector is sequentially passed through an encoder and a decoder to obtain a decoding feature, and the decoding feature is connected with the compression result of the fused position vector based on the connector compression, and the connection result is mapped through a Sigmoid activation function and a linear layer to obtain a structural information encoding, and a prediction value is obtained based on the structural information encoding through a Softmax activation function again, and the first residual module and the connector are both residual blocks of the TiDE model;

[0032] A cross loss function is constructed through the predicted value and the true value, and the ASTPB model is trained through the cross loss function based on the dataset to obtain the pre-trained ASTPB model.

[0033] Preferably, parsing the submission program into an abstract syntax tree includes:

[0034] The submission program includes java source code or C code, and javalang is used to process the Java source code, or pycparser is used to process the C code to obtain an abstract syntax tree.

[0035] Preferably, in step S3, the method for generating the second individual set includes obtaining the second individual set by selecting, crossover and mutating individuals in the first individual set.

[0036] Preferably, the attention feature is obtained based on the numerical attribute features of the selected sub-models through the multi-head self-attention model, including:

[0037] The Q, K and V values ​​of the multi-head self-attention model are all taken from the numerical attribute characteristics of the sub-model. The results output by the multi-head attention model are sequentially subjected to the first fully connected layer, activation function, second fully connected layer, regularization and normalization to obtain the attention features.

[0038] Preferably, the predicted running time is obtained based on the fused features through a feed-forward network, including:

[0039] The feedforward network includes a third full connection, a Relu activation function, a fourth full connection and a Softmax layer, and the attention features and structural information encoding are sequentially input into the third full connection, the Relu activation function, the fourth full connection and the Softmax layer to obtain the predicted running time.

[0040] Preferably, the cost model is trained by a loss function based on the predicted running time and the actual running time, and the loss function is a mean square error loss function.

[0041] Compared with the prior art, the present invention has the following beneficial effects:

[0042] The present invention encodes the structural information output by the pre-trained ASTPB model and applies it to the TVM model, so that the TVM model can take into account the structural characteristics of the sub-model during the scheduling of the sub-model. At the same time, the present invention trains the multi-head self-attention model and the feedforward network with a smaller number of parameters, so as to quickly obtain the predicted running time. Through the above two points, the parameter scheduling of the sub-model can be realized more quickly, so that the sub-model running time is appropriate and the scheduling efficiency is improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] Figure 1A flow chart of a sub-model-based scheduling method for an AI compiler provided in a specific embodiment of the present invention;

[0044] Figure 2 A flowchart of a sub-model-based scheduling method for an AI compiler provided in a specific embodiment of the present invention. DETAILED DESCRIPTION

[0045] The present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be pointed out that the embodiments described below are intended to facilitate the understanding of the present invention and do not have any limiting effect on the present invention.

[0046] In order to speed up the scheduling time of the AI ​​compiler and improve efficiency, the present invention adds the pre-trained ASTPB model to the TVM model. In the process of searching parameters, the structure of the sub-model can be considered to have the possibility of effective scheduling. At the same time, the cost model is constructed by the ASTPB model with frozen parameters, the multi-head self-attention model with a small number of parameters, and the feedforward network. The training can be completed quickly and the predicted running time can be obtained, so that the scheduling of the sub-model can be completed efficiently.

[0047] like Figure 1 and Figure 2 As shown, a specific embodiment of the present invention provides a sub-model-based scheduling method for an AI compiler, including:

[0048] Step 1: Use TVM to divide the target model to be scheduled into multiple sub-models, generate a corresponding Sketch for each sub-model, and generate a fixed number of first volume sets according to the Sketch of the sub-model. P n , to build the initial population, and send the initial population to the hardware for execution to obtain the actual running time.

[0049] Step 2. Construct a cost model and train the cost model based on the predicted running time and the actual running time: the cost model includes a pre-trained ASTPB model with frozen parameters, a multi-head self-attention model and a feedforward network. The corresponding structural information encoding is obtained based on the selected sub-model through the pre-trained ASTPB model with frozen parameters, and the attention feature is obtained based on the numerical attribute features of the first individual set through the multi-head self-attention model. The predicted running time is obtained through the feedforward network based on the attention feature and the structural information encoding. The cost model is trained through the loss function based on the predicted running time and the actual running time.

[0050] In one specific embodiment, the present invention extracts a first individual set. P n The characteristics of the numerical attribute characteristics F b ,Should Fb Including the number of various floating-point operations: floating-point addition, floating-point subtraction, floating-point multiplication and some mathematical operation functions, etc.; the number of various integer operations: such as the floating-point operations mentioned above; vector-related features: the number of innermost loops vectorized (SIMD), the number of vectorized operations and the length product; loop unrolling-related features: including the maximum length of the unrolling, etc.; CUDA thread binding related: the length, addition, and product of blockId and threadId, etc.; computational density: defined as the number of floating-point operations / number of bytes of access space; a series of cache access features; memory application features; other features: the maximum number of loop unrolling operations, etc.

[0051] The specific embodiment of the present invention converts the numerical attribute feature F b Input the multi-head self-attention model to get the attention feature T attnout When scheduling the algorithm, after each hardware evaluation data is obtained, a neural network with fewer parameters using a multi-head attention mechanism is trained. The multi-head attention mechanism is as follows:

[0052]

[0053]

[0054]

[0055] Where n is the number of longs, d k is the dimension of Q, K, V, W 0 is the output weight matrix, Respectively i The weight matrix of queries, keys, and values ​​for each header.

[0056] The specific embodiment of the present invention uses a multi-head attention mechanism to extract features. In engineering implementation, our Q, K, and V values ​​are all taken as F b . Then use the linear layer and Relu activation function, the calculation is as follows:

[0057]

[0058]

[0059]

[0060]

[0061]

[0062] in, and is the parameter matrix of the first and second fully connected layers, with the same input and output dimensions. It is layer normalization, which normalizes a single sample in a batch.

[0063] In a specific embodiment, this embodiment provides a method for pre-training an ASTPB model, including:

[0064] We use multiple submitted programs from the PTA (Programming Teaching Assistant) website as data sets, and use Pycparser or Javalang to generate an abstract syntax tree (AST) of the submitted program, taking the correctness of the submitted program as the truth value. Each node in the abstract syntax tree contains specific information about the corresponding tag in the source code. In each node, there are attributes used to define the node features, and a training sample set is constructed based on the abstract syntax trees corresponding to multiple sub-models.

[0065] This embodiment uses word2vec to encode the semantic information of each node of the abstract syntax tree to obtain the first encoding information , the multiple word vector information in each node of the first encoding information is accumulated and averaged to obtain the second encoding information for:

[0066]

[0067] in, k is the index of the node, j is the index of word vector information.

[0068] The ASTPB method provided in the specific embodiment of the present invention focuses on using position binding as a more efficient and effective method for capturing position information in an abstract syntax tree. ASTPB uses the following form for position binding. The position information of each node is represented by a symbol: ,in, is the depth of the node in the AST, is the sequence number of the node in this layer of AST, is the depth of the node's parent node in the AST, It is the sequence number of the parent node of the node in the current AST layer.

[0069] The node location information PB is encoded to obtain a location vector, so that the dimension of the location vector is the same as the numerical attribute vector. The formula is as follows:

[0070]

[0071]

[0072] in, value Represents each parameter value in the position information, represents the encoding value at the even position in the encoding vector for each parameter value, x dim represents the dimension of the numerical attribute vector, t Indicates position count, value 0, 1, 2..., X dim / 8, concatenate the encoding vectors of the four parameter values ​​to obtain the position vector of the position information PB .

[0073] In this embodiment, the position vector PB Add the second coded information to get the vector representation of the tree structure for:

[0074]

[0075] The vector representation of the tree structure captures the relationship between nodes, taking into account both the properties of individual nodes and their relative positions (locally and globally) in the tree.

[0076] The ASTPB model provided in this embodiment includes a convolution module, a residual module, an encoder, a decoder, a connector and a simple neural network, and the simple neural network includes a Sigmoid activation function, a linear layer and a Softmax activation function.

[0077] In order to emphasize the local relationship between nodes, this embodiment uses the convolution method to transform the tree structure into a new tree structure. The vector representation of the tree structure is hierarchically fused to obtain a new tree, where each node of the new tree Obtained through the following methods

[0078]

[0079]

[0080] in, lfn represents the depth of the layer convolution, l Indicates the layer index, Represents the vector representation of each node after layer convolution, is the convolution operation l The sth convolution kernel of the layer, cn Indicates l The number of nodes in the layer, is the convolution operation l The initial vector of the sth node in the layer, For the lThe vector representation after the layer convolution, in the ASTPB model, the embodiment of the present invention generally selects lfn =1 or lfn =2, thus keeping the convolution layer at 2-3 layers can capture local information while maintaining low overhead.

[0081] The specific embodiment of the present invention uses the convolution module to perform routing traversal on the new tree after the layer convolution to obtain the vector representation of the traversed node, and form a node sequence V composed of the vector representation route , the specific routing uses the pre-order traversal of the tree structure similar to human reading habits to perform routing operations. After routing, a matrix sequence encoded by a piece of code after all the above steps can be obtained. This matrix sequence can be equivalent to a time series of indefinite length, that is, the node sequence V route .

[0082] Then, the residual module of the Tide model is used to extract the feature representation of the node sequence, which specifically includes compressing the node sequence to obtain the fused position vector. ResidualBlock () The compression process obtains the fused position vector V fp for:

[0083]

[0084]

[0085]

[0086]

[0087] Among them, W1, W2, W3 represent weight parameters, b1, b2, b3 represent bias parameters. The first fusion position vector is passed through the encoder and decoder in turn to obtain the decoding feature V decoder .

[0088] In a specific embodiment of the present invention, the decoded features and the fused position vector are compressed by a connector and subjected to Sigmoid activation and linear layer mapping to obtain structural information coding, and a prediction value is obtained based on the structural information coding through a Softmax activation function again.

[0089] The structural information encoding provided in this embodiment V sc for:

[0090]

[0091] Among them, it represents the connector implementation ResidualBlock ()Compression process.

[0092] It should also be noted that for other classification tasks, the ASTPB model can continue to use the Softmax layer to output a single prediction value y, which is related to the specific task and is expressed as:

[0093] y = Softmax( V sc ).

[0094] A cross loss function is constructed through the predicted value and the true value, and the ASTPB model is trained through the cross loss function based on the dataset to obtain the pre-trained ASTPB model.

[0095] In a specific embodiment, the predicted running time is obtained through a feedforward network based on the attention feature and the structural information encoding, including:

[0096] The feedforward network includes a third full connection, a Relu activation function, a fourth full connection and a Softmax layer, and the attention features and structural information encoding are sequentially input into the third full connection, the Relu activation function, the fourth full connection and the Softmax layer to obtain the predicted running time.

[0097] Specifically, this embodiment uses a simple feedforward neural network to extract and predict the attention features and structural information encoding. The purpose of using a simple feedforward neural network is to reduce the model calculation overhead as much as possible. The specific implementation process is shown in the following formula:

[0098]

[0099]

[0100]

[0101]

[0102]

[0103] in, and is the parameter matrix of the third and fourth fully connected layers. The layer converts the predicted values ​​to the interval (0,1).

[0104] In a specific embodiment, the cost model is trained by a loss function based on the predicted running time and the actual running time. loss MES for:

[0105]

[0106] Where D is the total number of iterations, For thed The actual running time obtained by sending the population obtained in the iteration to the hardware for execution is For the d The predicted running time of the cost model generation in epochs.

[0107] In a specific embodiment, the multi-head self-attention model and feedforward network of the cost model are replaced with the lightGBM model, the numerical attribute features and structural information encoding of the selected sub-model are input into the lightGBM model to obtain the predicted running time, and the cost model is trained through the loss function based on the predicted running time and the actual running time.

[0108] Step 3: Use the trained cost model to predict the predicted running time of each individual in the newly generated second individual set based on the features of each individual in the second individual set, select individuals from the second individual set based on the predicted running time to build a new population, and send the new population to the hardware for execution to obtain the corresponding actual running time.

[0109] In a specific embodiment, in step three, the method for generating the second individual set includes obtaining the second individual set by selecting, crossover and mutation of individuals in the first individual set.

[0110] Since a fixed number of individuals with parameter combinations are repeatedly generated during the generation of the population, a cost model is used to predict the performance of these parameter combinations on the hardware after these individuals are generated. Subsequently, based on these prediction results, a portion of the individuals with better performance are selected from these individuals to join the population for the next hardware evaluation. How to select a certain amount of generated individuals to join the population based on the prediction results is of specific significance. If only individuals with excellent prediction values ​​are selected to join the hardware evaluation each time, then on the one hand, there may be a small proportion of negative samples in the training set, resulting in a poor training effect of the model during the scheduling process. On the other hand, the scheduling process may fall into a local optimum too early, and converge too early to stop scheduling when there is still room for scheduling. At the same time, there will be obvious characteristics for the parameter combinations generated at different scheduling stages. In the early stage of scheduling, the predicted values ​​of these parameter combinations will be relatively scattered. In the later stage of scheduling, the predicted values ​​of these parameter combinations will become concentrated. Based on the above problems, the specific embodiments of the present invention propose the following parameter combination selection strategy:

[0111] Step (1): In this embodiment, the second individual set is sorted based on the predicted running time to obtain an individual sequence, and multiple individuals are selected from the individual sequence through an individual generation strategy to construct a new population.

[0112] In one embodiment, the specific steps of the individual generation strategy are:

[0113] This embodiment constructs a selection factor based on the predicted run time standard deviation, and the selection factor increases as the value of the predicted run time standard deviation decreases.

[0114] Specifically, the selection factor provided in this example is for:

[0115]

[0116]

[0117] in, is the standard deviation of the predicted run time, is a hyperparameter, where , is the prediction running time of the e-th individual, is the mean of the prediction running time, and N is the number of individuals.

[0118] In this embodiment, each time an individual is selected, m individuals are first selected from the individual sequence, and then n individuals are skipped, until there are no individuals in the individual sequence that can be selected after o selections, and a new individual set is obtained.

[0119] The number of individuals m selected each time is based on the selection factor, the number of individuals, the proportion of individuals selected in the current time, and the number of current selections.

[0120] Specifically, the number of individuals selected in the pth time Ch ( p )for:

[0121]

[0122] The number of individuals n skipped each time is based on the selection factor, the number of individuals, 1-the proportion of individuals selected in the current time, and the number of current selections.

[0123] Specifically, the number of individuals skipped for the pth time nCh ( p )for:

[0124]

[0125] Among them, the Ch ( p )and nCh ( p ) decreases as the value of the selection factor increases. is the proportion of individuals selected each time to the total number of individuals, where When selecting a parameter combination from the set of generated parameter combinations, first select Ch (0) added to the set, then skippednCh (0) elements, then select Ch (1) Add to the set, and repeat this process until Ch ( p ) = 0 to stop.

[0126] Using the above individual generation strategy, when the standard deviation decreases, the weight of the high prediction value of the sample set finally obtained from the sequence will gradually decrease, ensuring the existence of negative samples while ensuring that enough high-prediction individuals are obtained.

[0127] Step (2): When the new individual set is not less than the threshold number of individuals for constructing the population, a new population is obtained. When the new individual set is less than the threshold number of individuals for constructing the population, step (1) is repeated based on the predicted running time of each individual in the regenerated third individual set to screen out new individuals from the third individual set, and the new individuals are added to the new individual set to obtain a new population.

[0128] Step 4: Repeat Step 2 to Step 3 until it exceeds the repetition number threshold, that is, the iteration number upper limit is reached or the actual running time is lower than the running time threshold, then stop repeating and obtain the parameter combination of the sub-model to complete the scheduling of the sub-model.

Claims

1. A sub-model-based scheduling method for an AI compiler, characterized in that: include: S1. Use TVM to split the target model into multiple sub-models and corresponding Sketches. Select the first set of bodies based on the Sketches corresponding to the sub-models to build the initial population. Send the initial population to the hardware for execution to obtain the actual running time. S2. Construct a cost model, wherein the cost model includes a pre-trained ASTPB model with frozen parameters, and a multi-head self-attention model and a feedforward network to be trained, obtain corresponding structural information encoding based on the selected sub-model through the pre-trained ASTPB model, obtain attention features based on the numerical attribute features of the first individual set through the multi-head self-attention model, fuse the attention features and the structural information encoding to obtain fused features, obtain the predicted running time based on the fused features through the feedforward network, and train the cost model through the loss function based on the predicted running time and the actual running time; S3, predicting the predicted running time of each individual in the newly generated second individual set through the trained cost model, selecting individuals from the second individual set based on the predicted running time to construct a new population, and sending the new population to the hardware for execution to obtain the corresponding actual running time; S4, repeating steps S2-S3 until the number of repetitions exceeds the threshold or the actual running time is lower than the running time threshold, stopping the repetition, and obtaining the parameter combination of the sub-model to complete the scheduling of the sub-model; Pre-trained ASTPB model, including: The multiple submission programs on the PTA website are used as data sets, and the correctness of the submitted programs is used as the true value; Parse the submitted program into an abstract syntax tree, use word2vec to encode the semantic information of each node of the abstract syntax tree to obtain first encoding information, and perform weighted averaging on the data of each node of the first encoding information to obtain second encoding information; Encode the position information of each node to obtain a position vector, and add the position vector and the second encoded information to obtain a vector representation of the tree structure, wherein the position information includes the depth of the node in the abstract syntax tree, the position of the node in the sibling nodes at the same depth, and the position of the parent node of the node in the tree structure; The ASTPB model includes a convolution module, a residual module, an encoder, a decoder, a connector and a simple neural network, and the simple neural network includes a Sigmoid activation function, a linear layer and a Softmax activation function; The vector representation of the tree structure is hierarchically fused through the convolution module to obtain a new tree. The nodes of the new tree are organized into a node sequence through routing traversal. The node sequence is compressed through the residual module to obtain a fused position vector. The fused position vector is passed through the encoder and decoder in turn to obtain decoding features. The decoding features and the compression results of the fused position vector based on the connector are connected. The connection results are mapped through the Sigmoid activation function and the linear layer to obtain the structural information encoding. The Softmax activation function is used again to obtain the predicted value based on the structural information encoding. The residual module and the connector are both residual blocks of the TiDE model. A cross loss function is constructed through the predicted value and the true value, and the ASTPB model is trained through the cross loss function based on the dataset to obtain the pre-trained ASTPB model.

2. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: The multi-head self-attention model and feedforward network of the cost model in step S2 are replaced with the lightGBM model, the numerical attribute features and structural information encoding of the selected sub-model are input into the lightGBM model to obtain the predicted running time, and the cost model is trained through the loss function based on the predicted running time and the actual running time.

3. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: Individuals are selected from the second set of individuals based on the predicted running time to construct a new population, including: S3.1, sorting the second individual set based on the predicted running time to obtain an individual sequence, and selecting multiple individuals from the individual sequence through the individual generation strategy to construct a new population; The individual generation strategy includes: constructing a selection factor based on the predicted run time standard deviation, wherein the selection factor increases as the value of the predicted run time standard deviation decreases; Each time an individual is selected, m individuals are first selected from the individual sequence, and then n individuals are skipped until there are no individuals in the individual sequence that can be selected after o selections, and a new set of individuals is obtained; The number of individuals m selected each time is based on the selection factor, the number of individuals, the proportion of individuals selected in the current time, and the number of times the current selection is made; The number of individuals n skipped each time is based on the selection factor, the number of individuals, 1-the proportion of individuals selected in the current time, and the number of current selections; The m value and n value both decrease as the value of the selection factor increases; S3.

2. When the new set of individuals is not lower than the threshold number of individuals for constructing a population, a new population is obtained. When the new set of individuals is lower than the threshold number of individuals for constructing a population, step S3.1 is repeated based on the predicted running time of each individual in the regenerated third set of individuals to screen out new individuals from the third set of individuals, and the new individuals are added to the new set of individuals to obtain a new population.

4. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: Parse the submission program into an abstract syntax tree, including: The submission program includes java source code or C code, and javalang is used to process the Java source code, or pycparser is used to process the C code to obtain an abstract syntax tree.

5. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: In step S3, the method for generating the second individual set includes obtaining the second individual set by selecting, crossover and mutation of individuals in the first individual set.

6. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: The attention features are obtained based on the numerical attribute features of the selected sub-models through the multi-head self-attention model, including: The Q, K and V values ​​of the multi-head self-attention model are all taken from the numerical attribute characteristics of the sub-model. The results output by the multi-head attention model are sequentially subjected to the first fully connected layer, activation function, second fully connected layer, regularization and normalization to obtain the attention features.

7. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: The predicted running time is obtained based on the fused features through the feedforward network, including: The feedforward network includes a third full connection, a Relu activation function, a fourth full connection and a Softmax layer, and the attention features and structural information encoding are sequentially input into the third full connection, the Relu activation function, the fourth full connection and the Softmax layer to obtain the predicted running time.

8. The sub-model-based scheduling method for an AI compiler according to claim 1, characterized in that: The cost model is trained based on the predicted running time and the actual running time through a loss function, where the loss function is a mean square error loss function.

Images