Method and device for calling kernel function by Triton operator and electronic equipment

Abstract

The invention discloses a method and device for calling a kernel function by a Triton operator and electronic equipment, and relates to the technical field of artificial intelligence. The method comprises the steps that a runtime structure of a Triton operator is constructed, if a kernel function is called for the first time and is a single core, the kernel function is compiled through JITFunction, if the kernel function is called for the first time and is a multi-core, compiling parameters with the best overall performance are obtained through overall adjustment, if the kernel function is not called for the first time, the kernel function is obtained through a mapping table in LibEntry, and if the kernel function is called for the second time, the kernel function is called for the first time and is a multi-core. Compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions are written in the mapping table. Practice shows that according to the scheme provided by the invention, the end-to-end time overhead can be reduced by about 70% in small-scale function calling. Moreover, the parameter tuning space of the multi-core operator is expanded through a global tuning mode, the degree of freedom of configuration parameters is improved, and the calculation efficiency of reduction operators is improved.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence technology, and in particular to a method, device and electronic equipment for calling a kernel function by a Triton operator. Background Art

[0002] Triton is an efficient and easy-to-use parallel programming language that is oriented towards thread block granularity programming. It can write high-performance operator code on a variety of AI chips such as GPUs and DSA (Domain Specific Architecture), accelerating the reasoning and training process of large models. Triton is based on the Python language environment and follows PyTorch's interface definition for tensors. It can input and output tensor data of the torch.Tensor type. The task of the Triton compiler is to compile the operator function from an abstract high-level language to a machine language through multiple layers of intermediate representation, generate binary core files in real time, and save them in the user cache directory, using hash values ​​to mark and identify the core.

[0003] After the Triton compiler completes the compilation of the operator core, whenever an operator call occurs, the operator core will be searched and obtained through the runtime mechanism. The Triton compiler provides three types of runtime packaging of operator cores, which exist in the outer layer of the Triton function in the form of decorators. From the outside to the inside, they are: Autotuner (automatic parameter tuner), which selects the best operator core parameters through actual execution and performance measurement methods; Heuristics (heuristic parameter tuner), which calculates core parameters by passing and calling heuristic functions; JITFunction (just-in-time compilation function), which maintains the runtime core mapping table, uses compilation parameters, call parameters, environment parameters and other information to store and read operator cores, and the runtime mapping table is automatically lost when the process ends.

[0004] like Figure 1 As shown in the figure, in the Triton operator definition, a JITFunction object must be wrapped outside the operator function; while Autotuner and Heuristics are more flexible, and can be used without or with multiple objects wrapped outside the JITFunction. When an operator call occurs, the Triton runtime will call the running functions of all objects in sequence from the outside to the inside in the order of object packaging to complete parameter acquisition and operator core search.

[0005] Test experiments show that even if the GPU performance of the Triton operator is consistent with that of the native PyTorch operator, the time consumption of the Triton operator measured on the CPU will be higher than that of the PyTorch operator; the smaller the scale of the tensor being calculated, the more obvious the time consumption of the Triton operator lags behind. Through further experimental monitoring and analysis, it was found that the time overhead of the Triton runtime is the main reason for the high CPU time consumption of the Triton operator, among which the JITFunction causes the most time consumption.

[0006] During the runtime of JITFunction, there are several sources of time consumption: (1) JITFunction will call Python components to obtain various parameters from the parameters passed to the Triton operator function, even if the information of these parameters is already included in the compiled operator core binary file; (2) JITFunction will call the inspect.bind (the name of the Python language function interface) interface to bind the function's incoming parameters to the function's registered parameter list, complete the strict correspondence between the two, and encapsulate all parameters into a custom KernalArg (Triton type name) object, which is the largest runtime overhead; (3) JITFunction will classify the parameters and calculate the key values, and parameters of different categories may intersect; (4) JITFunction searches the corresponding operator core binary file from the mapping table based on the key value and returns it. Except for the fourth overhead, the first three items are unnecessary behaviors, which increase the CPU time at runtime. Taking tensor addition as an example, the Triton kernel function needs to process six parameters in the JITFunction stage: input tensor A, input tensor B, output tensor C, tensor size N, scalar operand alpha, and thread block size S. Assume that the default value of alpha is 1, the default value of S is 512, and S is a fixed parameter compiled into the core. When the user calls, the input tensor A, input tensor B, output tensor C, and scalar N are passed. The autotuner Autotuner is used to adjust S to 256, and alpha is not passed. The parameter processing process can be as follows: Figure 2 Therefore, the existing Triton runtime mechanism takes a long time.

[0007] The Triton runtime has a built-in Autotuner automatic parameter tuner. Developers set several groups of configuration parameters to be selected in the Autotuner instance, and specify the key-value parameters that affect the selection of configuration parameters, which are marked in the form of decorators before the kernel function definition. When the kernel function is called, Autotuner selects the configuration parameters with the best performance from these groups of configuration parameters by compiling, running, and measuring the time consumption, and passes them to the subsequent runtime process to complete the compilation of the kernel function. The optimal configuration parameters are also saved in the cache of the Autotuner instance, forming a mapping between key-value parameters and configuration parameters. When the kernel function is called again using the same key-value parameters, Autotuner will automatically find the configuration parameters that match the mapping in the cache and pass them without re-tuning. In addition, Autotuner also provides an optional pruning function. When the number of configuration parameters set by the developer is too large, compiling and running them one by one will significantly consume tuning time and affect execution efficiency; the pruning function can delete some configuration parameters with poor expected performance before tuning according to the rules specified by the developer, reduce the parameter space, and improve the tuning speed.

[0008] Triton's built-in Autotuner only adjusts parameters for the Triton core functions it decorates. However, some functions need to be implemented in combination with multiple Triton core functions, which is a scenario that a single tuner for kernel functions cannot handle. Multiple core functions are interconnected through intermediate variables, so that the parameter selection of one core function may affect the index space and parameter selection of other core functions; and the best configuration parameters obtained by independently adjusting the parameters of each core function may not lead to the best performance of the function as a whole. Therefore, due to the functional limitations of Triton's native Autotuner, the functions of multi-core functions can often only fix the intermediate variables and then tune the kernel functions separately, and it is impossible to obtain the globally optimal configuration parameters. Taking tensor accumulation as an example, when performing global accumulation of larger-scale tensors, in order to fully utilize parallelism, a two-stage core is required to process. Such as Figure 3 As shown, in the first stage, the input tensor is divided into several blocks, each block is input into a core 1, accumulated into a tensor element, and the tensor elements output by all cores form a new tensor as an intermediate variable; in the second stage, the intermediate variable is input into core 2, accumulated into a tensor element, and forms the final output. In this process, the intermediate variable connects core 1 and core 2, and its size will directly affect the parallelism of core 1 and the task scale of core 2, thereby merging the tuning tasks of the two cores into one problem. Therefore, the tuning of the accumulation operator needs to adjust the size of the intermediate variable, the various parameters of core 1, and the various parameters of core 2 for the best overall performance at the same time, which is obviously beyond the capabilities of Autotuner. Therefore, the existing Triton automatic parameter tuner has certain functional limitations. Summary of the invention

[0009] In order to solve the problems existing in the prior art, the present invention provides the following technical solutions.

[0010] A first aspect of the present invention provides a method for a Triton operator to call a kernel function, comprising: Constructing the runtime structure of the Triton operator, the constructed runtime structure includes decorators from the inside out: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; Determine whether the kernel function is called for the first time based on the collected registration parameters and the parameters passed in when calling; If the kernel function is called for the first time and it is a single core, the kernel function is compiled through JITFunction, and the compilation parameters and the compiled kernel function are written into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the comprehensive performance of each core is taken as the goal, and the intermediate variables and the parameters of each core are adjusted as a whole to obtain the compilation parameters with the best global performance, and the best compilation parameters and the corresponding kernel function are written into the mapping table of LibEntry; If the kernel function is not called for the first time, the kernel function is obtained through a mapping table in LibEntry, where the mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

[0011] Preferably, judging whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during calling includes: In LibEntry, the incoming parameters during the call are mapped to the collected registration parameters and packaged, and a key value is obtained using the packaged parameters; Check whether the key value obtained from the mapping table exists. If not, the kernel function is called for the first time; if it exists, it is not the first time the kernel function is called; if the kernel function is called for the first time, a mapping is established between the key value and the compiled kernel function and written into the mapping table.

[0012] Preferably, the step of matching the incoming parameters during the call to the collected registration parameters includes: Match the incoming parameters to the registered parameters one by one according to the position and keyword. For parameters that are not passed in but have default values, use the default values ​​to supplement them, and other parameters that are not passed in are temporarily left empty.

[0013] Preferably, obtaining a key value using the packaged parameters includes: The packaged parameters are divided into specialized parameters, non-specialized parameters and constant parameters; constant parameters are parameters marked as constants in the function declaration; specialized parameters are parameters that are not marked in the function declaration and are compiled into the kernel function; non-specialized parameters are parameters marked in the JITFunction decorator and are not compiled into the kernel function; A key value is obtained using the data type and value of the specialized parameter, the data type of the non-specialized parameter, and the value of the constant parameter.

[0014] Preferably, the overall performance of each core is taken as the goal, and the intermediate variables and the parameters of each core are adjusted as a whole at the same time to obtain the compilation parameters with the best global performance, which includes: The parameters include task block parameters and hardware parallel parameters, wherein the task block parameters include the size of tensors processed by the program in a thread block and the size of tensors processed in a loop body; the hardware parallel parameters include the number of thread warps allocated by the program in a stream processor and the number of stages for soft pipelining on storage; Compare the core performance produced by each set of parameters, and select the one with the shortest running time and the best performance as the output result.

[0015] Preferably, the following method is used to determine whether the kernel function is a single core: If the LibEntry decorator wraps a function that has been decorated with JITFunction, and the inner layer is an instance of JITFunction, the core function is single-core; if the LibEntry decorator wraps an ordinary Python function, the core function is multi-core.

[0016] Preferably, the method further comprises: If the key value is not found in the mapping table in the current thread, the status of the thread lock is queried; If the thread lock is disabled, the current thread first enables the thread lock protection mapping table and then compiles the kernel function; If the thread lock is enabled, the current thread queries repeatedly until the thread lock is disabled or the key value is found.

[0017] A second aspect of the present invention provides a device for a Triton operator to call a kernel function, comprising: A runtime structure construction module is used to construct the runtime structure of the Triton operator. The constructed runtime structure includes decorators from the inside to the outside: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; A judgment module is used to judge whether the kernel function is called for the first time based on the collected registration parameters and the input parameters during the call; The compilation module called for the first time is used to compile the kernel function through JITFunction if the kernel function is called for the first time and it is a single core, and write the compilation parameters and the compiled kernel function into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the intermediate variables and the parameters of each core are adjusted as a whole at the same time with the comprehensive performance of each core as the goal to obtain the compilation parameters with the best global performance, and write the best compilation parameters and the corresponding kernel function into the mapping table of LibEntry; The compilation module for non-first call is used to obtain the kernel function through the mapping table in LibEntry if the kernel function is not called for the first time. The mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

[0018] A third aspect of the present invention provides a memory storing a plurality of instructions, wherein the instructions are used to implement the method of the Triton operator calling a kernel function as described in the first aspect.

[0019] The fourth aspect of the present invention provides an electronic device, comprising a processor and a memory connected to the processor, the memory storing a plurality of instructions, the instructions being loaded and executed by the processor so that the processor can execute the method of calling a kernel function by a Triton operator as described in the first aspect.

[0020] The beneficial effects of the present invention are as follows: the method, device and electronic device for calling a kernel function by the Triton operator provided by the present invention can reduce the end-to-end time overhead of the operation when the kernel function is called by constructing a new runtime structure, which is mainly manifested in three aspects: (1) LibEntry only captures part of the parameters passed in by the user at runtime, without obtaining all the registered parameters, because other parameters that need to be obtained through the automatic tuner and the heuristic parameter tuner have been recorded in the mapping table, and can directly access the mapping table to read, saving more parameter processing time. (2) LibEntry saves the time of calling the built-in inspect.bind binding function of Python, and only needs to use the position order of the passed parameters, pass keywords and default values ​​to complete the parameter correspondence, and does not need to register to the proprietary category KernelArg, reducing additional overhead. (3) LibEntry divides function parameters into three non-intersecting types, and only calculates the key value once for each parameter, avoiding redundant key value calculations and compressing the key value processing time as much as possible. Practice shows that the solution provided by the present invention is more effective in small-scale function calls, and the end-to-end time overhead can be reduced by about 70%. Moreover, global tuning has expanded the parameter tuning space of multi-core operators and improved the freedom of configuration parameters. Compared with tuning each core individually, it pays more attention to the global combination performance of operators and improves the computational efficiency of reduction operators. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 This is a schematic diagram of the structure of Triton when it is running; Figure 2 This is a schematic diagram of the parameter processing method when Triton is running; Figure 3 Schematic diagram of tensor accumulation method in the first and second stages; Figure 4 A schematic diagram of the process flow of the Triton operator calling the kernel function of the present invention; Figure 5 A schematic diagram of the implementation process of the method for calling a kernel function by the Triton operator of the present invention; Figure 6 A schematic diagram of the functional structure of the device for calling the kernel function by the Triton operator described in the present invention. DETAILED DESCRIPTION

[0022] In order to better understand the above technical solution, the above technical solution will be described in detail below in conjunction with the accompanying drawings and specific implementation methods.

[0023] The method provided by the present invention can be implemented in the following terminal environment, and the terminal may include one or more of the following components: a processor, a memory, and a display screen. The memory stores at least one instruction, and the instruction is loaded and executed by the processor to implement the method described in the following embodiment.

[0024] The processor may include one or more processing cores. The processor uses various interfaces and lines to connect various parts of the entire terminal, and executes various functions of the terminal and processes data by running or executing instructions, programs, code sets or instruction sets stored in the memory, and calling data stored in the memory.

[0025] The memory may include random access memory (RAM) or read-only memory (ROM). The memory may be used to store instructions, programs, codes, code sets or instructions.

[0026] The display screen is used to display the user interface of each application.

[0027] In addition, those skilled in the art can understand that the structure of the above terminal does not constitute a limitation on the terminal, and the terminal may include more or fewer components, or combine certain components, or arrange the components differently. For example, the terminal also includes components such as a radio frequency circuit, an input unit, a sensor, an audio circuit, and a power supply, which will not be described in detail here.

[0028] Embodiment 1 like Figure 4 As shown, an embodiment of the present invention provides a method for a Triton operator to call a kernel function, comprising: S101, constructing a runtime structure of the Triton operator, wherein the constructed runtime structure includes decorators from the inside out: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; S102, judging whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during calling; S103, if the kernel function is called for the first time and it is a single core, the kernel function is compiled through JITFunction, and the compilation parameters and the compiled kernel function are written into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the comprehensive performance of each core is taken as the goal, the intermediate variables and the parameters of each core are adjusted as a whole at the same time to obtain the compilation parameters with the best global performance, and the best compilation parameters and the corresponding kernel function are written into the mapping table of LibEntry; S104, if the kernel function is not called for the first time, the kernel function is obtained through a mapping table in LibEntry, wherein the mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

[0029] In the runtime structure of the constructed Triton operator, the decorator can also include Heuristics and Autotuner. If the decorators Heuristics and Autotuner are included, when the kernel function is called for the first time, the kernel function is compiled through Autotuner, Heuristics and JITFunction in sequence.

[0030] In the embodiment of the present invention, LibEntry has four functions: initialization function, key value calculation function, global tuning function and run call function, which are defined outside the core function in the form of decorators. If the function is a Triton single core function, the LibEntry decorator also needs to be outside all Triton native runtime decorators (including Autotuner, Heuristics and JITFunction). Figure 1 In comparison, the Triton kernel function decorated with LibEntry will add a LibEntry layer to the outermost layer of the original runtime structure.

[0031] When defining an operator function decorated by LibEntry, LibEntry will first initialize some variables, including creating an empty mapping table, setting thread locks, etc. The empty mapping table is used to independently maintain the compiled Triton kernel function and the corresponding constant parameters outside the Triton native runtime mechanism. The thread lock is used to ensure thread safety and avoid repeated compilation in a multi-threaded environment. Secondly, LibEntry will collect various information about the function and identify whether the function it wraps is a single-core Triton function or a multi-core Python function: for a single-core Triton function, LibEntry will obtain the parameter name list and decorator category attributes from the second outer decorator, and obtain parameter specialization options from JITFunction; for a multi-core function, LibEntry will obtain the interface definition parameter list, and collect the incoming candidate parameter information to build a global tuner to select the parameter combination with the best overall performance for the function.

[0032] In one embodiment of the present invention, judging whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during the call includes: In LibEntry, the incoming parameters during the call are mapped to the collected registration parameters and packaged, and a key value is obtained using the packaged parameters; Check whether the key value obtained from the mapping table exists. If not, the kernel function is called for the first time; if it exists, it is not the first time the kernel function is called; if the kernel function is called for the first time, a mapping is established between the key value and the compiled kernel function and written into the mapping table.

[0033] Specifically, the following method can be used to map the incoming parameters during the call to the collected registration parameters: map the incoming parameters to the registration parameters one by one according to the position and keyword, use the default values ​​to supplement the parameters that are not passed in but have default values, and leave the other parameters that are not passed in temporarily empty. Using this method, LibEntry can capture only some of the parameters passed in by the user at runtime without obtaining all the registration parameters, because other parameters that need to be obtained through the automatic tuner and heuristic tuner have been recorded in the configuration mapping table, and can be directly accessed to read the mapping table, saving more parameter processing time.

[0034] Furthermore, the method of obtaining a key value by using the packaged parameters may include: dividing the packaged parameters into specialized parameters, non-specialized parameters and constant parameters; wherein the constant parameters are parameters marked as constants in the function declaration; the specialized parameters are parameters that are not marked in the function declaration and are compiled into the kernel function; the non-specialized parameters are parameters that are marked in the JITFunction decorator and are not compiled into the kernel function; and obtaining a key value by using the data type and value of the specialized parameters, the data type of the non-specialized parameters and the value of the constant parameters. The key value is in the form of a tuple: wherein for the constant parameter, its value is directly added to the key value; for the specialized parameter, its data type and information on whether the address is aligned are added to the key value; for the non-specialized parameter, only its data type is added to the key value.

[0035] In the embodiment of the present invention, after obtaining a key value, the calculated key value can be checked from the mapping table to see whether it exists. If it does not exist, the kernel function is called for the first time; if it exists, the kernel function is not called for the first time; if the kernel function is called for the first time, the key value is mapped to the compiled kernel function and written into the mapping table. This method is applicable to both single-core functions and multi-core functions.

[0036] In a preferred embodiment of the present invention, the overall performance of each core is taken as the goal, and the intermediate variables and the parameters of each core are adjusted as a whole at the same time to obtain the compilation parameters with the best global performance, which may include: The parameters include task block parameters and hardware parallel parameters, wherein the task block parameters include the size of tensors processed by the program in a thread block and the size of tensors processed in a loop body; the hardware parallel parameters include the number of thread warps allocated by the program in a stream processor and the number of stages for soft pipelining on storage; Compare the core performance produced by each set of parameters, and select the one with the shortest running time and the best performance as the output result.

[0037] In another preferred embodiment of the present invention, the following method can be used to determine whether the kernel function is a single core: If the LibEntry decorator wraps a function that has been decorated with JITFunction, and the inner layer is an instance of JITFunction, then the core function is single-core; if the LibEntry decorator wraps a normal Python function, then the core function is multi-core. Therefore, at runtime, LibEntry can determine whether the function is a single-core function by identifying whether the function has been decorated with Triton's JITFunction and whether it is an instance of the JITFunction class.

[0038] The method for calling a kernel function by a Triton operator provided by the present invention may further include the following steps: if the key value is not found in the mapping table in the current thread, the state of the thread lock is queried; if the thread lock is in a disabled state, the current thread first enables the thread lock protection mapping table, and then compiles the kernel function; if the thread lock is in an enabled state, the current thread repeatedly queries until the thread lock is disabled or the key value is found.

[0039] In a specific embodiment, when a single-core Triton function is called, LibEntry will first process the parameters passed in by the user from the calling interface to the Triton function call. According to the parameter name information obtained from JITFunction in the initialization phase, LibEntry maps the passed-in parameters to the registered parameters one by one according to the position and keyword, and adds the default values ​​for the parameters that are not passed in. Then, LibEntry classifies the set of all parameters into specialized parameters, non-specialized parameters, and constant parameters, where both specialized parameters and non-specialized parameters will be used to call the compiled kernel function, while constant parameters have been compiled into the kernel function and no longer participate in the call. For example, in the example of tensor addition, the specialized parameters include the pointers and tensor size N of the input tensor A, input tensor B, and output tensor C, the non-specialized parameters include the scalar operand alpha, and the constant parameters include the thread block size S. Next, LibEntry will call the key value calculation function to encode and calculate the specialized parameters, non-specialized parameters, and constant parameters into a set of key values, through which the current call requirements of the Triton function can be mapped to the only optimal kernel function. The key value is composed of the data type and value of specialized parameters, the data type of non-specialized parameters, the value of constant parameters, etc., with almost no information redundancy.

[0040] After obtaining the key value, LibEntry searches the mapping table corresponding to the current device to see if the key value exists, resulting in two situations.

[0041] In the first case, if the key value does not exist, it means that the call requirement is the first time and there is no pre-compiled kernel function. Therefore, it is necessary to return to the Triton native runtime mechanism, such as Figure 5 The miss in the code will enter path 1. The native runtime mechanism will complete parameter tuning, heuristic calculation, just-in-time compilation and operation in the order of decorator packaging, execute the compiled kernel function and return it to LibEntry. LibEntry writes the return value kernel function and the constant parameters it uses into the mapping table and finally exits. In the second case, if the key value exists, it means that the current call has been compiled, so Figure 5 Path 2 in the figure. LibEntry can directly take out the compiled kernel function and constant parameters from the mapping table for reuse, calculate the parallel scale of the thread block with the constant parameters and the parameters passed by the user, and finally pass it to the kernel function to complete the operation and calculation.

[0042] When a multi-core function is called, LibEntry is still executed first. Figure 5 The process of path 2 in . First, the parameter list information collected during the initialization phase is combined to map the incoming parameters to the registered parameters by position and keyword. Parameters that are not passed in but have default values ​​are supplemented with default values, and other parameters that are not passed in are temporarily left vacant. Then, LibEntry calculates the key value from the incoming parameters according to the tuning key value definition set when the function is initialized, and queries the mapping table to see if there is a corresponding entry. If the mapping hits, the cached parameters are used to call and execute the multi-core function. If not, the tuning process of path 3 is entered.

[0043] In path 3, the global tuner needs to traverse the configuration parameters, compile the entire function, and measure the performance of the multi-core function in order to select the parameter configuration with the best global performance and write it back to the mapping where the key value is located. Figure 3 The two-stage tensor accumulation calculation shown will be optimized in the form of a multi-core combination, and the parameters of core 1, the size parameters of the intermediate variables, and the parameters of core 2 will be adjusted at the same time, with the comprehensive performance of each core as the goal to complete the global optimal parameter selection.

[0044] A more special case is that when running LibEntry with multiple threads, the reading and writing of the mapping table needs to ensure thread safety. Therefore, when LibEntry does not find the key value in the mapping table, it is necessary to query the status of the thread lock. If the thread lock is disabled, it means that no other thread is currently writing the mapping table. In this case, the thread will first enable the thread lock and protect the mapping table, and then compile and write using the native runtime mechanism to ensure that when the mapping table is locked, no other thread interferes with the write operation of the current thread; and if the thread lock is enabled, it means that other threads are compiling the kernel and writing the mapping table, then the current thread needs to query repeatedly until the thread lock is disabled or the key value is found. At this time, the mapping table is either released or the kernel function that the thread needs to find is already stored. This ensures that multi-threaded function calls will not be compiled repeatedly or conflicts will occur when writing to the mapping table, saving unnecessary runtime and compilation overhead.

[0045] Embodiment 2 like Figure 6 As shown, another aspect of the present invention also includes a functional module architecture that is completely consistent with the aforementioned method flow, that is, an embodiment of the present invention also provides a device for a Triton operator to call a kernel function, including: The runtime structure construction module 201 is used to construct the runtime structure of the Triton operator. The constructed runtime structure includes decorators from the inside to the outside: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; A judgment module 202 is used to judge whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during the call; The compilation module 203 called for the first time is used to compile the kernel function through JITFunction if the kernel function is called for the first time and it is a single core, and write the compilation parameters and the compiled kernel function into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the intermediate variables and the parameters of each core are adjusted as a whole at the same time with the comprehensive performance of each core as the goal to obtain the compilation parameters with the best global performance, and write the best compilation parameters and the corresponding kernel function into the mapping table of LibEntry; The non-first call compilation module 204 is used to obtain the kernel function through the mapping table in LibEntry if the kernel function is not called for the first time. The mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

[0046] Further, in the judgment module, judging whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during the call includes: In LibEntry, the incoming parameters during the call are mapped to the collected registration parameters and packaged, and a key value is obtained using the packaged parameters; Check whether the key value obtained from the mapping table exists. If not, the kernel function is called for the first time; if it exists, it is not the first time the kernel function is called; if the kernel function is called for the first time, a mapping is established between the key value and the compiled kernel function and written into the mapping table.

[0047] The step of matching the incoming parameters during the call to the collected registration parameters includes: Match the incoming parameters to the registered parameters one by one according to the position and keyword. For parameters that are not passed in but have default values, use the default values ​​to supplement them, and other parameters that are not passed in are temporarily left empty.

[0048] Further, obtaining a key value by using the packaged parameters includes: The packaged parameters are divided into specialized parameters, non-specialized parameters and constant parameters; constant parameters are parameters marked as constants in the function declaration; specialized parameters are parameters that are not marked in the function declaration and are compiled into the kernel function; non-specialized parameters are parameters marked in the JITFunction decorator and are not compiled into the kernel function; A key value is obtained using the data type and value of the specialized parameter, the data type of the non-specialized parameter, and the value of the constant parameter.

[0049] Furthermore, in the compilation module called for the first time, the intermediate variables and the parameters of each core are adjusted as a whole at the same time with the comprehensive performance of each core as the goal, and the compilation parameters with the best global performance are obtained, including: The parameters include task block parameters and hardware parallel parameters, wherein the task block parameters include the size of tensors processed by the program in a thread block and the size of tensors processed in a loop body; the hardware parallel parameters include the number of thread warps allocated by the program in a stream processor and the number of stages for soft pipelining on storage; Compare the core performance produced by each set of parameters, and select the one with the shortest running time and the best performance as the output result.

[0050] Furthermore, the following method is used to determine whether the kernel function is a single core: If the LibEntry decorator wraps a function that has been decorated with JITFunction, and the inner layer is an instance of JITFunction, the core function is single-core; if the LibEntry decorator wraps an ordinary Python function, the core function is multi-core.

[0051] Furthermore, the device also includes a thread lock module, which is used to: If the key value is not found in the mapping table in the current thread, the status of the thread lock is queried; If the thread lock is disabled, the current thread first enables the thread lock protection mapping table and then compiles the kernel function; If the thread lock is enabled, the current thread queries repeatedly until the thread lock is disabled or the key value is found.

[0052] The device can be implemented by the method of using the Triton operator to call the kernel function provided in the first embodiment. The specific implementation method can be found in the description of the first embodiment and will not be repeated here.

[0053] The present invention also provides a memory storing a plurality of instructions, wherein the instructions are used to implement the method of the Triton operator calling a kernel function as described in the first embodiment.

[0054] The present invention also provides an electronic device, comprising a processor and a memory connected to the processor, wherein the memory stores a plurality of instructions, and the instructions can be loaded and executed by the processor so that the processor can execute the method of calling a kernel function by a Triton operator as described in Example 1.

[0055] Although preferred embodiments of the present invention have been described, additional changes and modifications may be made to these embodiments by those skilled in the art once the basic inventive concepts are known. Therefore, the appended claims are intended to be interpreted as including the preferred embodiments and all changes and modifications that fall within the scope of the present invention. Obviously, those skilled in the art may make various changes and modifications to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalents, the present invention is also intended to include these modifications and variations.

Claims

1. A method for calling a kernel function using a Triton operator, characterized in that: include: Constructing the runtime structure of the Triton operator, the constructed runtime structure includes decorators from the inside out: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; Determine whether the kernel function is called for the first time based on the collected registration parameters and the parameters passed in when calling; If the kernel function is called for the first time and it is a single core, the kernel function is compiled through JITFunction, and the compilation parameters and the compiled kernel function are written into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the comprehensive performance of each core is taken as the goal, and the intermediate variables and the parameters of each core are adjusted as a whole to obtain the compilation parameters with the best global performance, and the best compilation parameters and the corresponding kernel function are written into the mapping table of LibEntry; If the kernel function is not called for the first time, the kernel function is obtained through a mapping table in LibEntry, where the mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

2. The method for calling a kernel function using a Triton operator as claimed in claim 1, wherein: The step of judging whether the kernel function is called for the first time according to the collected registration parameters and the input parameters during the call includes: In LibEntry, the incoming parameters during the call are mapped to the collected registration parameters and packaged, and a key value is obtained using the packaged parameters; Check whether the key value obtained from the mapping table exists. If not, the kernel function is called for the first time; if it exists, it is not the first time the kernel function is called; if the kernel function is called for the first time, a mapping is established between the key value and the compiled kernel function and written into the mapping table.

3. The method for calling a kernel function using a Triton operator as claimed in claim 2, wherein: The step of matching the incoming parameters during the call to the collected registration parameters includes: Match the incoming parameters to the registered parameters one by one according to the position and keyword. For parameters that are not passed in but have default values, use the default values ​​to supplement them, and other parameters that are not passed in are temporarily left empty.

4. The method for calling a kernel function using a Triton operator as claimed in claim 2, wherein: The method of obtaining a key value by using the packaged parameters includes: The packaged parameters are divided into specialized parameters, non-specialized parameters and constant parameters; constant parameters are parameters marked as constants in the function declaration; specialized parameters are parameters that are not marked in the function declaration and are compiled into the kernel function; non-specialized parameters are parameters marked in the JITFunction decorator and are not compiled into the kernel function; A key value is obtained using the data type and value of the specialized parameter, the data type of the non-specialized parameter, and the value of the constant parameter.

5. The method for calling a kernel function using a Triton operator as claimed in claim 1, wherein: The above-mentioned compilation parameters with the best global performance obtained by adjusting the intermediate variables and the parameters of each core at the same time with the comprehensive performance of each core as the goal include: The parameters of each core include task block parameters and hardware parallel parameters, wherein the task block parameters include the size of tensors processed by the program in a thread block and the size of tensors processed in a loop body; the hardware parallel parameters include the number of thread warps allocated by the program in a stream processor and the number of stages of soft pipelining on storage; Compare the core performance generated by the parameters of each set of cores, and select the set with the shortest running time and the best performance as the output result.

6. The method for calling a kernel function using a Triton operator as claimed in claim 1, wherein: The following method is used to determine whether the kernel function is a single core: If the LibEntry decorator wraps a function that has been decorated with JITFunction, and the inner layer is an instance of JITFunction, the core function is single-core; if the LibEntry decorator wraps an ordinary Python function, the core function is multi-core.

7. The method for calling a kernel function using a Triton operator as claimed in claim 2, wherein: The method further comprises: If the key value is not found in the mapping table in the current thread, the status of the thread lock is queried; If the thread lock is disabled, the current thread first enables the thread lock protection mapping table and then compiles the kernel function; If the thread lock is enabled, the current thread queries repeatedly until the thread lock is disabled or the key value is found.

8. A device for calling a kernel function using a Triton operator, characterized in that: include: A runtime structure construction module is used to construct the runtime structure of the Triton operator. The constructed runtime structure includes decorators from the inside to the outside: JITFunction and LibEntry; wherein LibEntry creates an empty mapping table, sets a thread lock, and collects registration parameters from other decorators during initialization; A judgment module is used to judge whether the kernel function is called for the first time based on the collected registration parameters and the input parameters during the call; The compilation module called for the first time is used to compile the kernel function through JITFunction if the kernel function is called for the first time and it is a single core, and write the compilation parameters and the compiled kernel function into the mapping table of LibEntry; if the kernel function is called for the first time and it is multi-core, the intermediate variables and the parameters of each core are adjusted as a whole at the same time with the comprehensive performance of each core as the goal to obtain the compilation parameters with the best global performance, and write the best compilation parameters and the corresponding kernel function into the mapping table of LibEntry; The compilation module for non-first call is used to obtain the kernel function through the mapping table in LibEntry if the kernel function is not called for the first time. The mapping table contains compilation parameters and corresponding kernel functions or optimal compilation parameters and corresponding kernel functions.

9. A memory, characterized in that: Multiple instructions are stored, and the instructions are used to implement the method of calling a kernel function by a Triton operator as described in any one of claims 1 to 7.

10. An electronic device, characterized in that: It includes a processor and a memory connected to the processor, the memory stores a plurality of instructions, and the instructions can be loaded and executed by the processor so that the processor can execute the method of calling a kernel function of a Triton operator as described in any one of claims 1 to 7.

Images