Performance optimization method and device for Triton compiler

Abstract

The invention discloses a performance optimization method and device for a Triton compiler. The method comprises the steps that kernel features of all operators are obtained, and all the operators are classified; then, corresponding configuration sets are constructed for the operators of all types, the configuration sets comprise multiple configuration schemes, and the configuration schemes are specifically configuration combinations of data block volumes, the number of thread bundles and the number of stages; and finally, determining the optimal configuration of each operator in the configuration set corresponding to each operator based on a machine learning model. The operator configuration scheme can be quickly and accurately determined, and the performance of a compiler is improved.

Description

Technical Field

[0001] This invention belongs to the field of computer technology, specifically relating to a performance optimization method and apparatus for the Triton compiler. Background Technology

[0002] With the rapid development of deep learning technology, a large number of new deep learning frameworks and hardware devices have emerged. The combination of frameworks and hardware has also shown an exponential growth trend. Faced with diverse deep learning frameworks, traditional manual optimization methods are difficult to guarantee that the model will achieve the best performance on different hardware platforms. In order to enable deep learning models to perform at their best on different hardware platforms, existing technologies provide deep learning compilers. Deep learning compilers connect deep learning models with hardware devices through automatic code generation and compilation optimization techniques, convert the model into a unified intermediate representation, and then perform compilation optimization to generate code adapted to specific acceleration devices, thereby improving the inference and training efficiency of the model.

[0003] However, existing deep learning compilers, such as the Triton compiler, are domain-specific languages ​​and their accompanying deep learning compilers, specifically designed for the development and compilation of efficient deep learning operators. One of their goals is to lower the barrier to GPU programming, enabling developers lacking GPU programming experience to quickly obtain high-performance code. The core solution of the Triton compiler is to simplify parallel programming through high-level abstraction, while leveraging compiler optimizations to ensure computational performance. Furthermore, improving the performance of the Triton compiler is crucial. Current technologies typically improve the execution efficiency of Triton compiler operators by introducing automatic tuning mechanisms. These mechanisms allow developers to predefine multiple kernel configurations and dynamically evaluate the performance of each configuration at runtime, thereby determining the optimal parameter combination and kernel allocation scheme for the Triton compiler. However, these automatic tuning mechanisms rely excessively on manually adding parameter configurations, reducing the portability of Triton compiler operators. Existing automatic tuning mechanisms are not precise enough to allow the Triton compiler to achieve optimal performance, thus reducing the processing performance of Triton compiler operators. Moreover, the automatic tuning mechanism may be time-consuming before determining the optimal parameter combination and kernel allocation scheme.

[0004] Therefore, how to quickly and accurately determine the configuration scheme of Triton compiler operators to improve the performance of Triton compiler is a technical problem that needs to be solved by those skilled in the art. Summary of the Invention

[0005] The purpose of this invention is to solve the technical problem that the Triton compiler in the prior art has low accuracy and may take too long when determining the configuration scheme of operators.

[0006] To achieve the above technical objectives, in one aspect, the present invention provides a performance optimization method for the Triton compiler, the method comprising:

[0007] Obtain the kernel features of each operator and classify each operator;

[0008] A corresponding configuration set is constructed for each type of operator. The configuration set includes multiple configuration schemes, which are specifically combinations of data block volume, number of thread bundles, and number of stages.

[0009] The optimal configuration of each operator is determined from the configuration set corresponding to each operator based on the machine learning model;

[0010] After constructing corresponding configuration sets for each type of operator, the method further includes reusing the intermediate file of the first processed operator into other operators of the same type and with the same block size, specifically including:

[0011] The operator that is processed first is taken as the first processing operator, and an index entry is established between the first block volume corresponding to the first processing operator and the saved intermediate file and saved to the corresponding type index table;

[0012] Obtain the volume of the second block corresponding to the next operator of the same type to be processed;

[0013] Determine whether there is a first block volume that is the same as the second block volume in the corresponding type index table;

[0014] If so, the corresponding intermediate file will be reused in the next operator of the same type based on the index item of the first block volume with the same volume;

[0015] If not, update the second block volume corresponding to the next operator of the same type to the first block volume, and at the same time, create an index entry between the first block volume corresponding to the next operator of the same type and its saved intermediate file and save it to the corresponding type index table.

[0016] Furthermore, the types of operators include pointwise computation type, reduction computation type, and dot product computation type.

[0017] Furthermore, obtaining the configuration set corresponding to the operator construction for each type of operator construction specifically includes:

[0018] When the operator is of type pointwise computation, the corresponding configuration set is constructed after adjusting the tensor shape to be greater than or equal to the smallest power of 2 of the current value, the data block volume to be a power of 2 and falling within the interval [256, 1024], the number of thread bundles to be within the interval [1, 2, 3, 4], and the number of stages to be 1.

[0019] When the operator type is a reduction computation type, a corresponding configuration set is constructed based on the operator dimension and data block volume;

[0020] When the operator is of type dot product calculation, a corresponding configuration set is constructed based on the shared memory usage.

[0021] Furthermore, after constructing corresponding configuration sets for each type of operator, the method also includes reusing the intermediate file of the first processed operator into other operators of the same type and with the same block size, specifically including:

[0022] The operator that is processed first is taken as the first processing operator, and an index entry is established between the first block volume corresponding to the first processing operator and the saved intermediate file and saved to the corresponding type index table;

[0023] Obtain the volume of the second block corresponding to the next operator of the same type to be processed;

[0024] Determine whether there is a first block volume that is the same as the second block volume in the corresponding type index table;

[0025] If so, the corresponding intermediate file will be reused in the next operator of the same type based on the index item of the first block volume with the same volume;

[0026] If not, update the second block volume corresponding to the next operator of the same type to the first block volume, and at the same time, create an index entry between the first block volume corresponding to the next operator of the same type and its saved intermediate file and save it to the corresponding type index table.

[0027] Furthermore, the intermediate file is specifically the Trtion IR generation and optimization intermediate file.

[0028] Furthermore, before determining the optimal configuration of each operator from the configuration set corresponding to each operator based on the machine learning model, the method further includes:

[0029] If the number of configuration schemes in the configuration set corresponding to the operator is not greater than a preset threshold, then a preset tool is invoked to measure the execution time of the operator under each configuration scheme;

[0030] The configuration scheme with the shortest execution time is taken as the optimal configuration for the corresponding operator.

[0031] Furthermore, the determination of the optimal configuration for each operator based on the machine learning model from the configuration set corresponding to each operator specifically includes:

[0032] Obtain the static features of the operator under each configuration in the configuration set;

[0033] The static features under each configuration are input into the machine learning model to obtain multiple execution time prediction results;

[0034] The configuration corresponding to the shortest execution time prediction result is taken as the optimal configuration of the operator.

[0035] Furthermore, the static features include tensor shape, operator type, data type, number of memory access instructions, number of loops, and number of loop iterations.

[0036] Furthermore, the machine learning model is specifically the XGBoost model.

[0037] On the other hand, the present invention also provides a performance optimization device for the Triton compiler, the device comprising:

[0038] The classification module is used to obtain the kernel features of each operator and classify the operators.

[0039] The configuration module is used to build corresponding configuration sets for various types of operators. The configuration set includes multiple configuration schemes, which are specifically configuration combinations of data block volume, number of thread bundles, and number of stages.

[0040] The determination module is used to determine the optimal configuration of each operator from the configuration set corresponding to each operator based on a machine learning model;

[0041] In this module, after constructing corresponding configuration sets for each type of operator, the configuration module is also used to reuse the intermediate file of the operator processed first to other operators of the same type and with the same block size.

[0042] This invention provides a performance optimization method and apparatus for the Triton compiler. Compared with existing technologies, this method first obtains the kernel features of each operator and classifies them; then, it constructs a corresponding configuration set for each type of operator, the configuration set including multiple configuration schemes, specifically combinations of data block size, thread bundle number, and stage number; finally, based on a machine learning model, it determines the optimal configuration for each operator in the corresponding configuration set. This method can quickly and accurately determine the configuration scheme for operators, improving compiler performance. Attached Figure Description

[0043] To more clearly illustrate the technical solutions in the embodiments or prior art of this specification, the drawings used in the description of the embodiments or prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments recorded in this specification. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0044] Figure 1 The diagram shown is a flowchart illustrating the performance optimization method of the Triton compiler provided in the embodiments of this specification.

[0045] Figure 2 The diagram shown is a structural schematic of the performance optimization device for the Triton compiler provided in the embodiments of this specification.

[0046] Figure 3 The diagram shown is a flowchart illustrating the performance optimization method of the Triton compiler in a specific application scenario of the embodiments of this specification.

[0047] Figure 4 The diagram shown illustrates the process of determining the optimal configuration using a machine learning model in a specific application scenario according to an embodiment of this specification.

[0048] Figure 5 The diagram shown is a hierarchical connection diagram of the machine learning model in the embodiments of this specification. Detailed Implementation

[0049] To enable those skilled in the art to better understand the technical solutions in this specification, the technical solutions in the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments in this application, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0050] like Figure 1The diagram illustrates a performance optimization method for the Triton compiler provided in this specification. While this specification provides the method operation steps or apparatus structure shown in the following embodiments or figures, based on conventional methods or without creative effort, the method or apparatus may include more or fewer operation steps or module units after partial merging. In steps or structures where there is no logically necessary causal relationship, the execution order of these steps or the module structure of the apparatus are not limited to the execution order or module structure shown in the embodiments or figures of this specification. When the method or module structure is applied in actual devices, servers, or terminal products, it can be executed sequentially or in parallel according to the method or module structure shown in the embodiments or figures (e.g., in a parallel processor or multi-threaded processing environment, or even in a distributed processing or server cluster implementation environment).

[0051] This application proposal specifically includes two aspects:

[0052] First, it supports automated parameter configuration, optimizing operators in three stages: parameter configuration setting, compilation, and optimal configuration search. Based on the mapping relationship between kernel computation characteristics and parameter configuration, it constructs the optimal kernel parameter configuration and implements intermediate representation reuse under the same block size configuration during compilation. At the same time, it adjusts the optimal parameter configuration search strategy to reduce the impact of system fluctuations on parameter selection.

[0053] Secondly, the hyperparameter-tuned machine learning model is integrated into the Triton compiler, and the trained model achieves efficient prediction with optimal parameter configuration, significantly reducing the overall compilation overhead.

[0054] The performance optimization methods for the Triton compiler provided in the embodiments of this specification can be applied to deep learning frameworks, such as... Figure 1 As shown, the method specifically includes the following steps:

[0055] Step S101: Obtain the kernel features of each operator and classify each operator.

[0056] The types of operators include pointwise computation type, reduction computation type, and dot product computation type.

[0057] By compiling to obtain operator kernel features, and based on the relationship between operator kernel features and parameter configuration, and according to program features such as whether the operator contains reduction operations and loop structures, the main operator kernels are divided into four types: pointwise computation (regular), reduction computation, loop reduction, and dot product computation, as shown in Table 1 below. The instructions contained in the operators of different operator types at the operator level and the intermediate file ttir level are as follows.

[0058] Table 1

[0059] ;

[0060] The operator kernel is compiled for the first time to obtain the contents of the ttir intermediate file. Since the kernel cannot be executed directly due to the lack of effective parameter configuration during the first compilation, a default configuration strategy is adopted, which sets the thread block size to 64 and initializes num_warps and num_stages to 4 and 2 respectively to construct a temporary Config object and trigger the compilation process.

[0061] By analyzing the instruction patterns in the ttir file, an automated classification of operator types was achieved: if the ttir contains the tt.reduce instruction, the operator is marked as a reduction type; otherwise, it is determined to be a pointwise computation type; for reduction type operators, if the scf.for instruction is further detected, it is further subdivided into a cyclic reduction type; by identifying the tt.dot instruction in the ttir, it is marked as a dot product computation type.

[0062] Step S102: Construct a corresponding configuration set for each type of operator. The configuration set includes multiple configuration schemes, specifically a combination of configurations for data block volume, number of thread bundles, and number of stages.

[0063] Specifically, the data block size, or blocksize, essentially refers to the data being divided into blocks. In the Triton compiler, the kernel processes data from different blocks each time it executes, determining the data size processed by each thread block. The number of stages, or num_stages, is the number of stages the compiler should use in the software pipeline loop, which is very useful for matrix multiplication workloads on SM80+ GPUs. Setting the number of stages primarily optimizes TensorCore computations by dividing the computation process into multiple stages and employing a pipelined parallel strategy to overlap data loading and computation operations, thereby improving computational throughput.

[0064] In this application, before constructing the corresponding configuration set for each type of operator, the method further includes determining whether the optimal configuration corresponding to the operator is stored in the cache. If it is stored, the optimal configuration scheme is directly called.

[0065] In this embodiment of the application, when the operator is a pointwise computation type, the corresponding configuration set is constructed after adjusting the tensor shape to be greater than or equal to the smallest power of 2 of the current value, the data block volume to be a power of 2 and falling within the range of [256, 1024], the number of thread bundles to be within the range of [1, 2, 3, 4], and the number of stages to be 1.

[0066] When the operator type is a reduction computation type, a corresponding configuration set is constructed based on the operator dimension and data block volume;

[0067] When the operator is of type dot product calculation, a corresponding configuration set is constructed based on the shared memory usage.

[0068] Specifically, the Triton compiler identifies hardware information such as operator type and thread structure, configures and combines key parameters such as block size, number of thread bundles and number of stages, and uses a brute-force combination method to generate candidate Config parameter configuration items for different types of computing kernels; that is, first determine the block size, and then combine it with multiple thread bundle numbers and number of stages to form multiple parameter configuration items, thus completing the construction of a reasonable kernel configuration scheme.

[0069] For pointwise computation types, the focus is on thread block size and memory access mode to improve data locality; the shape of the input tensor is obtained and adjusted to be a power of 2 greater than or equal to the current value; the block size is calculated based on the adjusted tensor shape; when the block size is a power of 2 and falls within the range of [256, 1024], the performance is better; the number of thread bundles can be selected in the range of [1, 2, 4, 8]; the number of stages is mainly applicable to dot product computation types and is set to 1 by default in pointwise computation; the permutation and combination of the above parameters can generate a set of candidate Config objects and their corresponding parameter configurations.

[0070] For point-by-point computation type operators, after applying the above-mentioned constraints to each parameter, multiple configuration schemes are generated through brute-force combination to form a configuration set.

[0071] Reduction calculations require consideration of both block size and the number of thread bundles. Looping reduction calculations require determining the block size based on the loop upper bound to improve computational efficiency. For reduction types without loop structures, if the reduction dimension is small and can load the entire row of data at once, the block size of the current dimension is directly set to the tensor shape of that dimension. When the reduction dimension is large, since it's impossible to load all data at once, the reduction calculation will evolve into a looping reduction type. In this case, the block size of the reduction dimension needs to be set to the loop step size, and the tensor shape of the reduction dimension should be used as the upper bound of the loop. A one-dimensional operator has one dimension and only needs one blocksize configured; a two-dimensional operator has two dimensions and requires two blocksizes configured.

[0072] Blocksize refers to the size of the block. Dimension refers to whether the thread block is arranged in one, two, or three dimensions. The tensor shape value is the arrangement and amount of data. That is, a shape value of (256, 128, 64) is equivalent to three-dimensional data. The first dimension has 256 elements, and the total amount of data is 256 * 128 * 64.

[0073] For operators of the reduction computation type, the tensor shape is determined based on the operator dimension and the data block volume, and then multiple configuration schemes are generated by directly combining the required parameters.

[0074] For dot product calculations, it is necessary to fully utilize the hardware features of TensorCore and optimize the number of stages and data partitioning strategies. The operators for dot product calculations are specifically designed to handle matrix multiplication operations. Their input matrix tensor shapes are [M, K] and [K, N], and their output matrix shape is [M, N]. First, the block sizes such as bsM, bsN, and bsK are initialized according to the shapes of the input and output tensors. Then, based on the shared memory capacity limit of the target hardware, the value of num_stages is dynamically adjusted to ensure that the shared memory usage is within the allowable range.

[0075] Because dot product calculations rely on TensorCores for efficient matrix multiplication and accumulation operations, data must be transferred from global memory to shared memory before being loaded into the TensorCores. Therefore, the capacity of shared memory becomes a key limiting factor. Shared memory usage... It can be calculated using the formula:

[0076]

[0077] Where bsM, bsN, and bsK represent the block size of the matrix in dimensions M, N, and K, respectively; dataType represents the tensor data type, such as fp32, fp16, etc., i.e., the memory size occupied by each element; the num_stages parameter typically ranges from 1 to 4 and is used to control the number of parallel stages in the pipeline, but its specific value needs to be verified by formula to see if it exceeds the capacity limit of shared memory; the parameter configuration construction algorithm for dot product calculation is implemented as follows:

[0078] Input: Kernel-passed parameters

[0079] Output: Dot product type parameter configuration

[0080] 1. CreatDotConfigs():

[0081] 2. getShapeMNK(args)

[0082] 3. dataType = getDataType(args)

[0083] 4. bsM = shapeToBlocksize(M) # Set the block size according to the matrix dimensions MNK

[0084] 5. bsN = shapeToBlocksize(N)

[0085] 6. bsK = setBlocksizeKList()

[0086] 7. warps, stages = setValueList() # Set Config parameters

[0087] 8. for k in bsK:

[0088] 9. for s in stages:

[0089] 10. SharedMem = (bsM+bsN)*dataType*k*s

[0090] 11. if SharedMem <= getActualSharedMem(): # Whether the actual shared memory capacity is exceeded.

[0091] 12. configs=CreateConfig(bsM,bsN,k,warps,s)

[0092] 13. dotConfigs.push(configs) # Combines the Config parameter configuration list

[0093] 14. Return dotConfigs;

[0094] Specifically, for operators of the dot product type, the generated configuration scheme is constrained by the amount of shared memory used. Multiple configuration schemes are generated while the amount of shared memory used does not exceed the capacity limit of shared memory.

[0095] After constructing corresponding configuration sets for each type of operator, the method further includes reusing the intermediate file of the first processed operator into other operators of the same type and with the same block size, specifically including:

[0096] The operator that is processed first is taken as the first processing operator, and an index entry is established between the first block volume corresponding to the first processing operator and the saved intermediate file and saved to the corresponding type index table;

[0097] Obtain the volume of the second block corresponding to the next operator of the same type to be processed;

[0098] Determine whether there is a first block volume that is the same as the second block volume in the corresponding type index table;

[0099] If so, the corresponding intermediate file will be reused in the next operator of the same type based on the index item of the first block volume with the same volume;

[0100] If not, update the second block volume corresponding to the next operator of the same type to the first block volume, and at the same time, create an index entry between the first block volume corresponding to the next operator of the same type and its saved intermediate file and save it to the corresponding type index table.

[0101] The intermediate file is specifically the intermediate file for generating and optimizing Trtion IR.

[0102] Specifically, for the same operator, when the block size of different configuration items is the same, the TritonIR ​​intermediate files generated by their compilation are completely identical.

[0103] The compiler's compilation process is divided into multiple stages, each of which generates corresponding intermediate files: Abstract Syntax Tree (AST) generation, Triton IR generation and optimization (TTIR), Triton GPU IR generation and optimization (TTGIR), LLVM IR generation and optimization (LLIR), PTX code generation (PTX), and CUDA binary code generation (CUBIN). In multiple configuration items with the same block size, only one TTIR intermediate file needs to be generated. Subsequent configuration items can directly reuse this intermediate file, thereby avoiding repeated execution of the AST generation and TTIR optimization processes and reducing the overall compilation cost.

[0104] The operator name, target hardware architecture, and configuration parameters are integrated into a unique hash value, which is then used as the name of the cache directory. This allows for the quick location and loading of intermediate files such as ttir corresponding to specific configuration items.

[0105] The compilation reuse optimization process involves: first, identifying configurations in the same blocks, indexing the first occurrence of configuration items and their saved intermediate files, and finally changing the compilation process for reuse.

[0106] The chunk size in the Config configuration is extracted as a key index item, and a mapping relationship is established between it and the generated hash value. During the compilation process, when processing new configuration items, the system checks whether their chunk size already exists in the index table.

[0107] If it exists, the intermediate ttir file in the cache directory pointed to by the hash value corresponding to the block size is directly reused, and the conversion from ttir to ttgir is used for subsequent compilation;

[0108] If it does not exist, the block size of the current configuration item and its generated hash value are added to the index table, and all intermediate files in the compilation stage are saved for later reuse.

[0109] Each parameter configuration item must be looked up and compared using the index table before compilation;

[0110] When the index table is empty or no matching entry is found for the current configuration item's chunk size in the table, the system will execute the complete compilation process, adding the generated chunk size and its corresponding hash value to the index table, and saving the compiled intermediate files to the cache directory. Conversely, if a configuration item with the same chunk size already exists in the index table, the system will still generate the hash value for the current configuration item, but will not add it to the index table; simultaneously, the system will extract the hash value corresponding to the cached chunk size from the index table and generate a cache path based on that hash value, thus directly reusing the saved intermediate files.

[0111] By creating an index table and identifying block sizes, the reuse of ttir intermediate files was achieved, and the compilation process was optimized accordingly.

[0112] First, copy the cached ttir files to the newly generated cache directory in the current configuration. Then, modify the compilation stage to skip the original first two stages, namely the abstract syntax tree generation and Triton IR generation and optimization stages. Finally, load the ttir file code in the cache directory directly into the compilation module, and continue the subsequent compilation process from the ttir to ttgir conversion part.

[0113] Step S103: Determine the optimal configuration of each operator from the configuration set corresponding to each operator based on the machine learning model.

[0114] The process of determining the optimal configuration for each operator from the configuration set corresponding to each operator based on a machine learning model specifically includes:

[0115] Obtain the static features of the operator under each configuration in the configuration set;

[0116] The static features under each configuration are input into the machine learning model to obtain multiple execution time prediction results;

[0117] The configuration corresponding to the shortest execution time prediction result is taken as the optimal configuration of the operator.

[0118] The static features include tensor shape, operator type, data type, number of memory access instructions, number of loops, and number of loop iterations. The machine learning model is specifically the XGBoost model.

[0119] XGBoost is a machine learning algorithm based on Gradient Boosting Decision Tree (GBDT). It constructs a strong learner by integrating multiple weak learners (usually decision trees). The model's input is a sample feature vector Xi containing m features, and the output is the predicted value of the sample. For regression problems, this is a continuous value; for classification problems, it is converted into class probabilities using a logistic function. For example... Figure 5 The diagram shows the hierarchical connections of the machine learning model, namely the XGBoost model. XGBoost is an "additive model" in which the prediction is a weighted sum of the predictions of K trees.

[0120] In machine learning models, the tree ensemble layer consists of k regression trees in sequence. Each new tree learns from the residuals of the predictions of all previous trees, i.e., the difference between the prediction and the actual result. The internal structure of the three trees consists of a root node, internal nodes, and leaf nodes. Nodes are used to divide samples into different branches based on a certain threshold of a feature. Leaf nodes are the final partitioning units, and each leaf node contains a continuous score (weight) w. The weight of the leaf node to which a sample finally falls is taken as the prediction contribution of that tree. During connection, the trees are generated sequentially, and the new tree fits the current residuals of the model. During prediction, the leaf weights of all trees are added together to obtain the final prediction value.

[0121] The training process for a machine learning model is as follows:

[0122] 1. Start from the root node: Place all training samples in the root node.

[0123] 2. Enumerate features and split points: For samples in the current node, traverse each feature and its possible split points (e.g., each value after sorting by feature value).

[0124] 3. Calculate the split gain: Evaluate the "gain" brought by each split point. The higher the gain, the more the overall loss function of the model is reduced after the split. The gain is calculated by measuring the changes in the statistics of the first-order gradient (g) and second-order gradient (h) of the samples within the node before and after the split.

[0125] 4. Select the best split: Select the feature and split point with the greatest gain, and split the current node into two child nodes, left and right.

[0126] 5. Recursive splitting: Repeat steps 2-4 for the generated left and right child nodes until the stopping condition is met (such as reaching the maximum depth of the tree, having too few node samples, or the gain being less than the set threshold).

[0127] The loss function of a machine learning model during training is as follows: the training objective is to minimize a target function that includes a regularization term. This target function consists of two parts:

[0128] ;

[0129] Loss term (L): Measures the error between the model's predicted value ŷi and the true label yi, such as mean squared error, log loss, etc.

[0130] Regularization term (Ω): Controls the complexity of the model and prevents overfitting.

[0131] In XGBoost, the complexity of a tree is defined as:

[0132] ,

[0133] Where T is the number of leaf nodes and wj is the leaf weight.

[0134] This prompts the model to choose simpler trees (fewer leaves, smaller weights). In the above formula, Obj(θ) is the overall objective function, representing the total objective that the XGBoost model needs to minimize during training. θ represents all the parameters of the model (in XGBoost, this refers to the structure of all trees and the weights of the leaves). fk: the k-th tree, which represents the k-th weak learner in the additive model, i.e., the k-th CART (regression) tree. The XGBoost model is obtained by sequentially adding K trees (f1, f2, ..., fK). γ and λ: regularization hyperparameters used to control the strength of regularization.

[0135] When searching for the optimal tree structure, XGBoost performs a second-order Taylor expansion of the loss function, which allows it to more accurately approximate the true loss function and naturally introduces the first-order gradient (g) and second-order gradient (h) to calculate the gain. This is one of the reasons why it converges faster and more accurately. In addition, XGBoost further prevents overfitting by reducing the contribution of each tree through shrinkage rate and by using strategies such as column sampling.

[0136] Specifically, by combining machine learning models with compiler optimizations, the trained model can directly predict the optimal parameter configuration, avoiding the cumbersome process of compiling, measuring, and selecting multiple configuration items in traditional methods. Figure 4 The diagram illustrates the process of determining the optimal configuration for a machine learning model in a specific application scenario.

[0137] Determine if the optimal parameters are saved in the cache;

[0138] If yes, then compile directly using the optimal parameter configuration; otherwise, compile for the first time to obtain kernel features.

[0139] Generate a config list based on type and shape, and combine it with kernel features to form data to be predicted;

[0140] XGBoost model prediction, obtain the list index corresponding to the minimum time;

[0141] Obtain the optimal configuration and save it to the cache;

[0142] Compile using the optimal parameter configuration.

[0143] The impact of operator kernel features and hardware environment on performance was analyzed, appropriate data features were selected, and redundant features that did not significantly contribute to the prediction results were removed.

[0144] The construction of the dataset for the machine learning model includes:

[0145] Using various typical neural network models such as ResNet50 and VGG16, Triton operators are generated using PyTorch's Inductor compiler to cover diverse computational types.

[0146] Based on the analysis of operator program characteristics, the dataset construction extracts features from multiple dimensions, including operator type, operands, block partitioning strategy, and instruction granularity. These features include operator type, tensor shape, data type, number of parameters, number of Load / Store instructions, loop structure (e.g., loop count, iteration count, upper bound), number of Load instructions within a loop, number of Triton Language (TL) instructions, and floating-point operations. Simultaneously, key hardware features affecting operator performance, such as the number of computation units, thread bundle size, and shared memory capacity, are included in the feature set. Finally, the kernel execution time under different parameter configurations is selected as the dataset's label value.

[0147] The dataset contains four main operator types: pointwise computation, reduction computation, cyclic reduction computation, and dot product computation. Each type accounts for approximately 25%, and the ratio of the training set to the validation set is 8:2.

[0148] Based on the characteristics of the training dataset and the fact that the predicted output value is continuous, the selected model is a nonlinear regression prediction model in supervised learning. Considering the small dataset and the high requirement for prediction accuracy, the XGBoost model was selected as the training model through K-fold cross-validation.

[0149] Some hyperparameters, such as n_estimators, max_depth, min_child_weight, gamma, and learning_rate, were selected, and the search range was set accordingly. A grid search was then performed to find the optimal hyperparameter values, as shown in Table 2 below. The grid search evaluates different combinations of given hyperparameters and finds the best parameter values ​​from them. It has a high computational cost but provides a comprehensive search and is suitable for small-scale parameter spaces.

[0150] Table 2

[0151] ;

[0152] The XGBoost model is deployed and predicted within the compiler.

[0153] Once the XGBoost runtime environment is configured, the saved .dat format model can be loaded by calling the load_model method, and then used in the compiler. The model is loaded at triton / runtime / autotuner::run.

[0154] The data to be predicted is organized. Since the compiler cannot directly obtain features such as operator computation logic and parallel structure, the ttir intermediate file of the operator is obtained during the first compilation. Program features are then extracted from the ttir file. The kernel program features indirectly obtained through ttir and the Config parameter configuration list automatically generated based on operator type and other information together constitute the feature data. This feature data is then organized according to the model training set format to form the dataset to be predicted.

[0155] Mean squared error (MSE), mean absolute error (MAE), and coefficient of determination R² are used to evaluate the predictive ability of the model. MSE and MAE measure the difference between the predicted and actual values; the smaller the value, the smaller the deviation between the predicted and actual values. R² measures the goodness of fit of the regression model to the data, and its value ranges between [0, 1]. The closer the value is to 1, the better the fit, and the closer it is to 0, the worse the fit.

[0156] The calculated mean squared error is 0.029, the mean absolute error is 0.063, and the R² is 0.87. However, as the number of samples increases, the value of R² usually increases as well. Therefore, it can only be used as a quantitative evaluation indicator and cannot be directly used to accurately measure the model's predictive ability. An adjusted coefficient of determination is introduced, which takes into account the influence of the number of feature variables in the model. The adjusted coefficient of determination is calculated to be 0.85, indicating that the model optimized by grid search has better predictive performance.

[0157] Model deployment involves extracting operator features, generating a dataset to be predicted, and loading a pre-trained XGBoost model to complete the prediction with the optimal parameter configuration.

[0158] Extract static features such as tensor shape, operator type, data type, number of memory access instructions, number of loops, and number of loop iterations. Among these, features such as tensor shape, operator type, and data type can be directly obtained from the input parameters and operator type identification process, while other features need to be extracted from the ttir intermediate file. The number of memory access instructions can be achieved by counting the load and store operation nodes in ttir. The feature extraction of loop structures requires parsing the control flow graph and identifying the number of nested loops and the number of iterations.

[0159] Based on the model training set format, the operator program and Config parameter feature data are organized to form the dataset to be predicted.

[0160] Configure the XGBoost runtime environment in the system, call the load_model method to load the model in the compiler, and use the generated dataset to be predicted to complete the prediction with the best configuration;

[0161] The system filters out the prediction execution time from the list of prediction results and records the list index of the prediction data. Then, it retrieves the best configuration item from the Config parameter configuration list using the list index to complete the function of predicting the best configuration.

[0162] like Figure 3 The diagram shows a flowchart illustrating the performance optimization method of the Triton compiler in a specific application scenario. The method includes:

[0163] Determine if the optimal parameters are stored in the cache;

[0164] If yes, then compile directly using the optimal parameter configuration; otherwise, compile for the first time to obtain kernel features.

[0165] Automatically generate configurations based on type and shape;

[0166] Determine if there are configurations with the same chunk size. If so, reuse intermediate files to continue compilation; otherwise, compile the configuration items; evaluate the best configuration.

[0167] The optimal configuration is saved to the cache;

[0168] Compile using the optimal parameter configuration.

[0169] In this embodiment of the application, before determining the optimal configuration of each operator in the configuration set corresponding to each operator based on the machine learning model, the method further includes:

[0170] If the number of configuration schemes in the configuration set corresponding to the operator is not greater than a preset threshold, then a preset tool is invoked to measure the execution time of the operator under each configuration scheme;

[0171] The configuration scheme with the shortest execution time is taken as the optimal configuration for the corresponding operator.

[0172] Specifically, before using a machine learning model to determine the optimal configuration, it is possible to first determine whether the number of configuration schemes corresponding to the operator is not greater than a preset threshold. If so, it means that the time taken for the operator to execute the current number of configuration schemes is not long. Preset tools can be used to measure the execution time of the operator under each configuration scheme, and the configuration scheme with the shortest execution time can be directly taken as the optimal configuration of the corresponding operator, without the need to use a machine learning model for prediction.

[0173] In specific application scenarios, the Triton compiler will execute the machine code generated by compiling each configuration item in turn, and use a specific search method to evaluate the performance of the configuration items. The measurement results will be used to compare the performance of different configuration items and finally select the best configuration.

[0174] In the automatic tuning mechanism, the do_bench performance testing tool is used to accurately measure the execution time of kernel functions. However, multiple experiments have shown that when using the do_bench method to perform multiple time measurements on the same operator under the same configuration list, the measurement results fluctuate significantly, which significantly affects the final search results.

[0175] The configuration was evaluated using the PyTorch Profiler tool, which is a core tool in the PyTorch framework for performance analysis and optimization of deep learning models.

[0176] Building upon the Profiler framework, pre-execution and statistical analysis functions have been added to implement an optimal parameter search time measurement algorithm and stably obtain measurement results. First, each parameter configuration is pre-compiled, and the generated intermediate files and compilation results are saved. Then, the Profiler tool is used to accurately measure the kernel execution time under each configuration, selecting the configuration with the shortest execution time as the optimal parameter. Next, the optimal configuration and its corresponding intermediate files are cached to avoid the overhead of repeated compilation and optimal configuration search. Finally, the final compilation process is completed under the determined optimal parameter configuration, generating high-performance target code.

[0177] Initialization requires two key parameters: `blocksize` and `key`. `blocksize` is a list specifying the names of the parameters whose block sizes need to be set in the configuration parameters; `key` is also a list, with each element representing the actual tensor shape corresponding to `blocksize`. During initialization, these two parameters are recorded in the `self.blocksize` and `self.key` properties, respectively, providing a basis for subsequent parameter configuration generation.

[0178] Based on the above-described performance optimization method for the Triton compiler, one or more embodiments of this specification also provide a platform or terminal for performance optimization of the Triton compiler. This platform or terminal may include devices, software, modules, plug-ins, servers, clients, etc., using the methods described in the embodiments of this specification, combined with necessary hardware implementation devices. Based on the same innovative concept, the systems in one or more embodiments provided in this specification are as described in the following embodiments. Since the implementation schemes and methods for solving the system problem are similar, the specific system implementations in the embodiments of this specification can refer to the implementation of the aforementioned methods. Repeated descriptions will not be repeated. The terms "unit" or "module" used below can refer to a combination of software and / or hardware that implements a predetermined function. Although the systems described in the following embodiments are preferably implemented in software, hardware implementations, and a combination of software and hardware, are also possible and contemplated.

[0179] Specifically, Figure 2 This is a schematic diagram of the module structure of one embodiment of the performance optimization device for the Triton compiler provided in this specification, as shown below. Figure 2 As shown, the performance optimization devices for the Triton compiler provided in this specification include:

[0180] Classification module 201 is used to obtain the kernel features of each operator and classify each operator;

[0181] Configuration module 202 is used to construct corresponding configuration sets for various types of operators. The configuration set includes multiple configuration schemes, and the specific configuration scheme is a combination of configurations for data block volume, number of thread bundles, and number of stages.

[0182] The determination module 203 is used to determine the optimal configuration of each operator based on the configuration set corresponding to each operator using a machine learning model;

[0183] Specifically, after constructing corresponding configuration sets for each type of operator, the configuration module 202 is further configured to reuse the intermediate file of the first processed operator in other operators of the same type and with the same block size, including:

[0184] The operator that is processed first is taken as the first processing operator, and an index entry is established between the first block volume corresponding to the first processing operator and the saved intermediate file and saved to the corresponding type index table;

[0185] Obtain the volume of the second block corresponding to the next operator of the same type to be processed;

[0186] Determine whether there is a first block volume that is the same as the second block volume in the corresponding type index table;

[0187] If so, the corresponding intermediate file will be reused in the next operator of the same type based on the index item of the first block volume with the same volume;

[0188] If not, update the second block volume corresponding to the next operator of the same type to the first block volume, and at the same time, create an index entry between the first block volume corresponding to the next operator of the same type and its saved intermediate file and save it to the corresponding type index table.

[0189] It should be noted that the system described above may include other implementation methods based on the description of the corresponding method embodiments. The specific implementation methods can be referred to the description of the corresponding method embodiments above, and will not be elaborated here.

[0190] This application also provides an electronic device, including:

[0191] processor;

[0192] Memory used to store the processor's executable instructions;

[0193] The processor is configured to perform the methods provided in the embodiments described above.

[0194] The electronic device provided in this application stores executable instructions for a processor in a memory. When the processor executes these instructions, it first acquires the kernel features of each operator and classifies them. Then, it constructs a corresponding configuration set for each type of operator. The configuration set includes multiple configuration schemes, specifically combinations of data block size, thread bundle number, and stage number. Finally, based on a machine learning model, it determines the optimal configuration for each operator within its corresponding configuration set. This allows for rapid and accurate determination of operator configuration schemes, improving compiler performance.

[0195] The foregoing has described specific embodiments of this specification. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps recited in the claims may be performed in a different order than that shown in the embodiments and may still achieve the desired result. Furthermore, the processes depicted in the drawings do not necessarily require the specific or sequential order shown to achieve the desired result. In some embodiments, multitasking and parallel processing are possible or may be advantageous.

[0196] The methods or apparatus described in the embodiments provided in this specification can implement business logic through a computer program and record it on a storage medium. The storage medium can be read and executed by a computer to achieve the effects of the solutions described in the embodiments of this specification, such as:

[0197] Obtain the kernel features of each operator and classify each operator;

[0198] A corresponding configuration set is constructed for each type of operator. The configuration set includes multiple configuration schemes, which are specifically combinations of data block volume, number of thread bundles, and number of stages.

[0199] The optimal configuration for each operator is determined from the configuration set corresponding to each operator based on a machine learning model.

[0200] The storage medium can include physical devices for storing information, typically digitizing the information and then storing it using electrical, magnetic, or optical methods. The storage medium can include: devices that store information using electrical energy, such as various types of memory, like RAM and ROM; devices that store information using magnetic energy, such as hard disks, floppy disks, magnetic tapes, magnetic core memory, bubble memory, and USB flash drives; and devices that store information using optical methods, such as CDs or DVDs. Of course, there are other readable storage media, such as quantum memories and graphene memories.

[0201] The embodiments in this specification are not limited to conforming to industry communication standards, standard computer resource data update and data storage rules, or the situations described in one or more embodiments of this specification. Slightly modified implementations based on certain industry standards or custom methods or embodiments can also achieve the same, equivalent, or similar, or predictable, implementation effects as described above. Embodiments that utilize these modified or modified methods for data acquisition, storage, judgment, and processing still fall within the scope of optional implementations of the embodiments in this specification.

[0202] The controller can be implemented in any suitable manner. For example, it can take the form of a microprocessor or processor and a computer-readable medium storing computer-readable program code (e.g., software or firmware) executable by the (micro)processor, logic gates, switches, application-specific integrated circuits (ASICs), programmable logic controllers, and embedded microcontrollers. Examples of controllers include, but are not limited to, the following microcontrollers: ARC 625D, Atmel AT91SAM, Microchip PIC18F26K20, and Silicon Labs C8051F320. A memory controller can also be implemented as part of the control logic of the memory. Those skilled in the art will also recognize that, in addition to implementing the controller in purely computer-readable program code form, the same functionality can be achieved by logically programming the method steps to make the controller take the form of logic gates, switches, ASICs, programmable logic controllers, and embedded microcontrollers. Therefore, such a controller can be considered a hardware component, and the means included therein for implementing various functions can also be considered as structures within the hardware component. Alternatively, the means for implementing various functions can be considered as both software modules implementing the method and structures within the hardware component.

[0203] The device embodiments described above are merely illustrative. For example, the division of units is only a logical functional division, and in actual implementation, there may be other division methods. For example, multiple units or plug-ins may be combined or integrated into another system, or some features may be ignored or not executed. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be through some interfaces, indirect coupling or communication connection between devices or units, and may be electrical, mechanical, or other forms.

[0204] These computer program instructions can also be loaded onto a computer or other programmable resource data updating device, causing a series of operational steps to be performed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions that execute on the computer or other programmable device for implementing the process. Figure 1 One or more processes and / or boxes Figure 1 The steps of the function specified in one or more boxes.

[0205] The various embodiments in this specification are described in a progressive manner. Similar or identical parts between embodiments can be referred to mutually. Each embodiment focuses on describing the differences from other embodiments. In particular, system embodiments are basically similar to method embodiments, so the description is relatively simple; relevant parts can be referred to the descriptions in the method embodiments. In the description of this specification, the terms "one embodiment," "some embodiments," "example," "specific example," or "some examples," etc., refer to specific features, structures, materials, or characteristics described in connection with that embodiment or example, which are included in at least one embodiment or example of this specification. In this specification, the illustrative expressions of the above terms do not necessarily refer to the same embodiment or example. Furthermore, the specific features, structures, materials, or characteristics described can be combined in any suitable manner in one or more embodiments or examples. Moreover, without contradiction, those skilled in the art can combine and integrate the different embodiments or examples described in this specification and the features of different embodiments or examples.

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

Claims

1. A performance optimization method for the Triton compiler, characterized in that, The method includes: Obtain the kernel features of each operator and classify each operator; A corresponding configuration set is constructed for each type of operator. The configuration set includes multiple configuration schemes, which are specifically combinations of data block volume, number of thread bundles, and number of stages. The optimal configuration of each operator is determined from the configuration set corresponding to each operator based on the machine learning model; After constructing corresponding configuration sets for each type of operator, the method further includes reusing the intermediate file of the first processed operator into other operators of the same type and with the same block size, specifically including: The operator that is processed first is taken as the first processing operator, and an index entry is established between the first block volume corresponding to the first processing operator and the saved intermediate file and saved to the corresponding type index table; Obtain the volume of the second block corresponding to the next operator of the same type to be processed; Determine whether there is a first block volume that is the same as the second block volume in the corresponding type index table; If so, the corresponding intermediate file will be reused in the next operator of the same type based on the index item of the first block volume with the same volume; If not, update the second block volume corresponding to the next operator of the same type to the first block volume, and at the same time, create an index entry between the first block volume corresponding to the next operator of the same type and its saved intermediate file and save it to the corresponding type index table.

2. The performance optimization method for the Triton compiler as described in claim 1, characterized in that, The types of operators include pointwise computation type, reduction computation type, and dot product computation type.

3. The performance optimization method for the Triton compiler as described in claim 2, characterized in that, The process of obtaining the configuration set corresponding to the operator construction for each type of operator construction specifically includes: When the operator is of type pointwise computation, the corresponding configuration set is constructed after adjusting the tensor shape to be greater than or equal to the smallest power of 2 of the current value, the data block volume to be a power of 2 and falling within the interval [256, 1024], the number of thread bundles to be within the interval [1, 2, 3, 4], and the number of stages to be 1. When the operator type is a reduction computation type, a corresponding configuration set is constructed based on the operator dimension and data block volume; When the operator is of type dot product calculation, a corresponding configuration set is constructed based on the shared memory usage.

4. The performance optimization method for the Triton compiler as described in claim 1, characterized in that, The intermediate file is specifically the intermediate file for generating and optimizing Trtion IR.

5. The performance optimization method for the Triton compiler as described in claim 1, characterized in that, Before determining the optimal configuration of each operator from the configuration set corresponding to each operator based on the machine learning model, the method further includes: If the number of configuration schemes in the configuration set corresponding to the operator is not greater than a preset threshold, then a preset tool is invoked to measure the execution time of the operator under each configuration scheme; The configuration scheme with the shortest execution time is taken as the optimal configuration for the corresponding operator.

6. The performance optimization method for the Triton compiler as described in claim 1, characterized in that, The process of determining the optimal configuration for each operator from the configuration set corresponding to each operator based on a machine learning model specifically includes: Obtain the static features of the operator under each configuration in the configuration set; The static features under each configuration are input into the machine learning model to obtain multiple execution time prediction results; The configuration corresponding to the shortest execution time prediction result is taken as the optimal configuration of the operator.

7. The performance optimization method for the Triton compiler as described in claim 6, characterized in that, The static features include tensor shape, operator type, data type, number of memory access instructions, number of loops, and number of loop iterations.

8. The performance optimization method for the Triton compiler as described in claim 6, characterized in that, The machine learning model mentioned is specifically the XGBoost model.

9. A performance optimization device for a Triton compiler, characterized in that, The device includes: The classification module is used to obtain the kernel features of each operator and classify the operators. The configuration module is used to build corresponding configuration sets for various types of operators. The configuration set includes multiple configuration schemes, which are specifically configuration combinations of data block volume, number of thread bundles, and number of stages. The determination module is used to determine the optimal configuration of each operator from the configuration set corresponding to each operator based on the machine learning model; In this module, after constructing corresponding configuration sets for each type of operator, the configuration module is also used to reuse the intermediate file of the operator processed first to other operators of the same type and with the same block size.

Images