Cross-language function calling method and device based on just-in-time compilation and storage medium

Abstract

The application provides a cross-language function calling method and device based on instant compilation and a storage medium, and belongs to the technical field of computers. The method comprises the following steps: in a runtime based on a first language end, a specialization signature is calculated according to a Triton instant compilation function and input parameters; based on the specialization signature and a hash table from the specialization signature to a CompiledKernel maintained in a cache of the Triton instant compilation function, if the required CompiledKernel is queried, the corresponding CompiledKernel is selected; otherwise, a Triton compiler is triggered for compilation, and the compiled CompiledKernel is placed in the cache of the Triton instant compilation function; parameters used for calling the required CompieldKernel are selected from the input parameters, and the corresponding CompiledKernel is called. The application embodiment can use the basic framework required by the triton jit function in the first programming language, so that the developer can develop an operator library based on the triton jit function implementation, and a C++ interface is provided. The application embodiment reduces the runtime overhead on the basis of retaining the jit running mechanism of the triton language.

Description

TECHNICAL FIELD

[0001] The present application belongs to the technical field of computer, and particularly relates to a cross-language function calling method and device based on instant compilation and a storage medium. BACKGROUND

[0002] In order to adapt to the flexible requirements of deep learning research and deployment, at present, there are many embedded Domain-Specific Languages (embedded DSL) or libraries based on the python language in the embedded field, which can enable users to write kernel functions suitable for accelerators such as GPUs in a language similar to python. Triton language (Triton Lang) is one of the languages, which only provides a usage mode embedded in the python end, therefore, the operators developed by using the Triton language generally only provide a Python interface. An operator generally takes calling one or more kernel functions as the core, and is assisted by necessary parameter checking and preprocessing, starting configuration calculation and other logic, and is wrapped into a more user-friendly Wrapper function, which can generally directly use the Tensor type object provided by the commonly used high-dimensional array library as a parameter (for example, the array type of torch, jax, etc.). However, since the logic between the Wrapper and the calling Kernel is executed on the Python end, the low efficiency of the Python interpreter may become a limiting factor affecting the overall performance. Even if the Kernel startup has asynchrony, the Kernel is controlled after being started, returns to the host thread and does not have to wait for the Kernel to complete execution, but if the Wrapper consumes too much time and even exceeds the execution time of the Kenel, the GPU will still be idle.

[0003] The main usage of Triton language is just-in-time compilation: a function written in the language needs to extract necessary information from incoming parameters, construct a specialized signature, and compile the required binary file by the compiler in real time, and then execute it by the runtime. Since Triton language is implemented as an embedded DSL of Python, the jit runtime is also often implemented in the Python end, and the runtime is responsible for providing enough information to trigger the compiler for just-in-time compilation, and executing the compilation result. The runtime may also be affected by the low efficiency of the Python interpreter. In addition, the compiler of Triton language also only provides an interface on the Python end. Therefore, the related technical solutions have the problems of: the runtime of Triton language depends on the Python interpreter and execution environment, which is not conducive to providing interfaces of other languages for functions written in Triton language; and the low efficiency of the Python interpreter affects the overall execution efficiency. SUMMARY

[0004] The present application provides a cross-language function calling method and device based on just-in-time compilation and a storage medium, aiming to solve the problems of the prior art that the runtime of Triton language depends on the Python interpreter and execution environment, which is not conducive to providing interfaces of other languages for functions written in Triton language; and the low efficiency of the Python interpreter affects the overall execution efficiency.

[0005] According to a first aspect of the present application, a cross-language function calling method based on just-in-time compilation is provided, comprising: in a runtime based on a first language end, calculating a specialized signature according to a Triton just-in-time compilation function and incoming parameters; based on the specialized signature and a hash table from the specialized signature to CompiledKernel maintained in the cache of the Triton just-in-time compilation function, if the required CompiledKernel is queried, selecting the corresponding CompiledKernel; otherwise, triggering the Triton compiler to compile, and placing the compiled CompiledKernel into the cache of the Triton just-in-time compilation function; and selecting parameters for calling the required CompieldKernel from the incoming parameters of the Triton just-in-time compilation function, and calling the corresponding CompiledKernel according to the required CompiledKernel and the selected CompiledKernel parameters.

[0006] In an optional embodiment, the triggering Triton compiler to compile comprises: calling a Triton compilation interface through a Python interpreter embedded in a runtime of a first language end; and passing the specialized signature to the Triton compilation interface, using a Triton compiler to compile, and wrapping a CompiledKernel into the compilation result. Wherein, the Triton compiler is configured with a file cache mechanism, according to a CompiledKernel to be compiled, a path in a file system where a corresponding compilation result is stored is specified, if the compilation result exists in the specified path, the path is returned, otherwise, the compilation is triggered first, and the compilation result is stored in the specified path.

[0007] In an optional embodiment, the cross-language function calling method further comprises managing instantiation of the Triton just-in-time compilation function through a multiple-instance mode, comprising: creating a global static hash mapping table for managing instantiated Triton just-in-time compilation function objects; combining a file path where the Triton just-in-time compilation function is located and a function name to generate a unique identifier as a key of the hash mapping, and the corresponding Triton just-in-time compilation function as a value of the hash mapping; when instantiation of the Triton just-in-time compilation function is requested, querying the hash mapping in the global static hash mapping table through the unique identifier, if the Triton just-in-time compilation function corresponding to the unique identifier is not in the query hash mapping, creating the corresponding Triton just-in-time compilation function and storing it in the query hash mapping; and managing cache of the Triton just-in-time compilation function through reuse of instances of the Triton just-in-time compilation function.

[0008] In an optional embodiment, the cross-language function calling method further comprises configuring a function parameter distribution rule of the Triton just-in-time compilation function based on a static signature rule, comprising: assigning a corresponding enumeration value to each parameter of the Triton just-in-time compilation function, so as to use a series of enumeration values as a static signature of the Triton just-in-time compilation function; generating a static signature object when the Triton just-in-time compilation function object is constructed by parsing the corresponding Triton function code; and generating the specialized signature according to the static signature object and distributing the incoming parameters to select parameters for calling the CompiledKernel when the Triton just-in-time compilation function object is called.

[0009] In an optional implementation, the enumerated values include CONSTEXPR, SPECIALIZED and NON-CONSTEXPR. Among them, CONSTEXPR identifies a compile-time constant parameter, SPECIALIZED identifies a non-constant parameter whose corresponding value affects the specialization signature, and NON-CONSTEXPR identifies a non-constant parameter whose corresponding value does not affect the specialization signature. The processing logic for the enumerated values is that for a parameter marked as CONSTEXPR, the value of the corresponding argument is embedded in the CompiledKernel generated by the compiler as part of the specialization signature, and for a parameter marked as SPECIALIZED, a feature is extracted from the corresponding argument according to a preset condition at compile time as part of the specialization signature, and for a parameter marked as NON-CONSTEXPR, the value of the argument does not participate in the generation of the specialization signature, and the parameter of the CompiledKernel generated by the compiler includes the parameter marked as NON-CONSTEXPR.

[0010] In an optional implementation, the cross-language function calling method further includes implementing polymorphic calling of the Triton just-in-time compilation function based on a variable parameter template, including: defining a calling operator of the Triton just-in-time compilation function as a variable parameter template function, so that a Triton just-in-time compilation function object receives different numbers and types of parameter packages to support calling of different Triton functions and different parameter combinations of the same Triton function; at compile time, performing type deduction on each parameter to determine the type to which the parameter belongs; and based on the determined type, value of each parameter and the static signature of the Triton just-in-time compilation function, determining the specialization signature and selecting parameters for calling the CompiledKernel.

[0011] In an optional implementation, the cross-language function calling method further includes a process of determining the specialization signature and selecting parameters for calling the CompiledKernel based on a folding expression, including: using a lambda expression for processing parameters to process each parameter according to the static signature rule through a folding expression.

[0012] In an optional implementation, the cross-language function calling method further includes: calling a driver API of a device to execute a required Kernel.

[0013] According to a second aspect of the present application, a control device is provided, which includes a memory, a processor, and a computer program stored on the memory and executable on the processor, and the processor executes the computer program to implement the above-mentioned cross-language function calling method.

[0014] According to a third aspect of the present application, a machine readable storage medium is provided, having instructions stored thereon, which cause a machine to perform the cross-language function calling method described above.

[0015] By the technical solution described above, the cross-language function calling method provided by the embodiments of the present application, in the runtime based on the first language, calculates a specialized signature according to the Triton just-in-time compiled function and the input parameters; based on the specialized signature and the hash table from the specialized signature to the CompiledKernel maintained in the cache of the Triton just-in-time compiled function, if the required CompiledKernel is queried, the corresponding CompiledKernel is selected; otherwise, the Triton compiler is triggered to compile, and the compiled CompiledKernel is put into the cache of the Triton just-in-time compiled function; and the parameters used to call the required CompieldKernel are selected from the input parameters of the Triton just-in-time compiled function, and the corresponding CompiledKernel is called according to the required CompiledKernel and the selected CompiledKernel parameters. The embodiments of the present application can use the basic framework required by the triton jit function in the first programming language (for example, C++ programming language), so that the developer can develop the operator library implemented based on the triton jit function, and provide the C++ interface. The embodiments of the present application reduce the runtime overhead on the basis of retaining the jit running mechanism of the triton language.

[0016] Other features and advantages of the present application will be set forth in the following description, and in part will become apparent to those skilled in the art from the description, or can be learned by practice of the present application. The objects and other advantages of the present application can be realized and achieved by the structures and processes as described in the specification and drawings. BRIEF DESCRIPTION OF DRAWINGS

[0017] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the related art, the drawings needed to be used in the embodiments or related art description will be briefly introduced. Obviously, the drawings in the following description are some embodiments of the present application, and those skilled in the art can also obtain other drawings according to these drawings without creative labor.

[0018] Figure 1 is a flowchart of the cross-language function calling method based on just-in-time compilation provided by the embodiments of the present application.

[0019] Figure 2is a flowchart of an example cross-language function calling method provided by an embodiment of the present application.

[0020] Figure 3 is a Triton Just-in-time Compilation function instance acquisition flowchart provided in an embodiment of the present application.

[0021] Figure 4 is a schematic diagram of using Triton Just-in-time Compilation function to implement torch operator and provide python interface in an example embodiment of the present application. DETAILED DESCRIPTION

[0022] In order to make the purposes, technical solutions and advantages of the embodiments of the present application clearer, the technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are some of the embodiments of the present application, not all of the embodiments of the present application. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative work fall within the scope of protection of the present application.

[0023] As described above, since Triton language is implemented as an embedded DSL of Python, its Just-in-time Compilation runtime is also implemented in the Python end. The runtime is responsible for providing sufficient information to trigger the compiler for Just-in-time Compilation, and executing the compilation result. The runtime may also be affected by the low efficiency of the Python interpreter. In addition, the compiler of Triton language also only provides an interface on the Python end. Therefore, the related technical solutions have problems including that the runtime of Triton language depends on the Python interpreter and execution environment, which is not conducive to providing an interface of other languages for functions written in Triton language; and the low efficiency of the Python interpreter affects the overall execution efficiency.

[0024] To this end, an embodiment of the present application provides a cross-language function calling method based on Just-in-time Compilation. The method provides a scheme for using a function written in an external language (for example, Triton Lang) in a first language, that is, a low-level language (for example, C++ language) through Just-in-time Compilation, so that a developer can develop an operator library based on Triton Just-in-time Compilation function (TritonJITFunction) implementation, and provide a low-level language interface, and reduce the runtime overhead on the basis of retaining the jit runtime mechanism of Triton language.

[0025] Figure 1 is a flowchart of a cross-language function calling method based on Just-in-time Compilation provided by an embodiment of the present application. The method can include the following steps:

[0026] Step S110: In the runtime based on the first language end, a specialization signature is calculated according to the Triton JIT function and the input parameters.

[0027] wherein the Triton JIT function is a type defined in the runtime, is an encapsulation of a function written in the Triton language, and can be called by specifying a file path and a function name structure where the corresponding function is located. The CompiledKernel refers to an encapsulation of a compilation result of a function written in the Triton language, and can be called by specifying a file path and a function name structure where the compilation result is located.

[0028] In the embodiments of the present application, the torch operator is implemented using Triton JIT C++ runtime on the first language end (for example, the C++ end). Please refer to Figure 2 For example, in the runtime, a specialization signature is calculated according to the Triton JIT function (TritonJITFunction), the input parameters, the running environment, the compilation target, etc., and the corresponding CompiledKernel can be selected according to the specialization signature. In the Triton jit C++ runtime, it can be implemented as a C++ library, and the core types configured are TritonJITFunction and CompiledKernel.

[0029] In the embodiments of the present application, the cross-language function calling method can further include managing the instantiation of Triton JIT functions through a multi-instance mode, including: creating a global static hash mapping table for managing the instantiated Triton JIT function objects; combining the file path and the function name of the Triton JIT function to generate a unique identifier as the key of the hash mapping, and the corresponding Triton JIT function as the value of the hash mapping; when requesting to instantiate the Triton JIT function, querying the hash mapping in the global static hash mapping table through the unique identifier, if the Triton JIT function corresponding to the unique identifier is not in the query hash mapping, creating the corresponding Triton JIT function and storing it in the query hash mapping; and managing the cache of Triton JIT functions through reuse of Triton JIT function instances.

[0030] Please refer to Figure 3For example, regarding the construction of TritonJITFunction instance, on the C++ side, the file path where the function implemented in triton language is located and the function name are provided, and the corresponding TritonJITFunction instance is obtained. TritonJITFunction uses a Multiton mode to manage its instances. All TritonJITFunction instances exist in a global hash table with the key being the concatenation of the file path where the function is located and the function name, and the value being TritonJITFunction. When TritonJITFunction instance is obtained by providing the file path where the function implemented in triton language is located and the function name, the path and the function name are first concatenated into a key, and the hash table is queried. If the key does not exist in the hash table, a TritonJITFunction instance is first created and stored in the hash table, otherwise the reference of the TritonJITFunction instance obtained by the query is directly returned.

[0031] In the embodiment of the present application, for the call of TritonJITFunction instance, the input parameters include a fixed part and a variable part. The fixed part includes stream, started grid (which can include x, y, z dimensions), compilation options num_warps, num_stages, etc., which can be determined by the back end. The variable part is defined by the corresponding function implemented in triton language. Because the call method of TritonJITFunction is implemented as a variable parameter template (which will be described in detail below, and will not be described here), it can support variable parameters.

[0032] Step S120: Based on the specialized signature and the hash table from the specialized signature to the CompiledKernel maintained in the cache of Triton JIT function, if the required CompiledKernel is queried, the corresponding CompiledKernel is selected.

[0033] In the embodiment of the present application, the call of Triton JIT function is implemented by specializing and calling CompiledKernel. Please refer to Figure 2The information required for the instantiation of the CompiledKernel can be derived from the incoming parameters, the runtime environment, the compilation target, etc. in the example, and will be extracted from the sources in the following. The information used to determine the CompiledKernel is called the specialization signature. The TritonJITFunction object is preferably configured to maintain a hash table. Each TritonJITFunction instance maintains a corresponding cache, implemented as a hash table, where the Key is the specialization signature and the Value is the CompiledKernel, for fast lookup of the required CompiledKernel according to the specialization signature.

[0034] Step S130: Otherwise, trigger the Triton compiler to compile, and place the compiled CompiledKernel into the cache of the Triton JIT function.

[0035] In the preferred embodiment of the application, triggering the Triton compiler to compile can include: calling the Triton compilation interface through the Python interpreter embedded in the runtime of the first language side; and passing the specialization signature to the Triton compilation interface, using the Triton compiler to compile, and wrapping the compilation result into a CompiledKernel. Wherein, the Triton compiler is configured with a file caching mechanism, according to the CompiledKernel to be compiled, specifying the path in the file system where the corresponding compilation result is stored, if the compilation result exists in the specified path, return the path, otherwise, first trigger the compilation, and store the compilation result in the specified path.

[0036] Reference is made to Figure 2 In the embodiment of the application, the Runtime is implemented in C++, and the Compiler part can reuse the Triton compiler. The embodiment of the application can configure a separate compilation script to call the Triton compilation interface (triton.compile). For example, when the required CompiledKernel does not exist in the cache of the Triton JIT function, a separate compilation script is triggered. Since the Triton compilation interface called therein itself uses a file system-based caching mechanism, the path of the compilation result can be used to construct a CompiledKernel instance according to the path, and store it in the cache of the Triton JIT function.

[0037] Step S140: Select the parameters for calling the required CompieldKernel from the incoming parameters of the Triton JIT function, and call the corresponding CompiledKernel according to the required CompiledKernel and the selected CompiledKernel parameters.

[0038] Please refer to Figure 2 The parameters required for calling the selected CompiledKernel are selected from the incoming parameters when calling the Triton Just-In-Time compilation function, and the corresponding CompiledKernel is called to implement the execution of Triton Just-In-Time compilation on the C++ side.

[0039] The embodiment of the present application provides a cross-language calling method based on Just-In-Time compilation, which can be divided into Runtime part and Compiler part. Wherein, the core types provided by the Runtime are TritonJITFunction and CompiledKernel. TritonJITFunction (Triton Just-In-Time function) corresponds to the function written in triton language. Each TritonJITFunction instance maintains the corresponding cache, which is implemented as a hash map table of specialized signatures to CompiledKernel; TritonJITFunction is responsible for computing the specialized signature, CompiledKernel is selected, the parameters of the Kernel are selected and the Kernel is called. CompiledKernel can be configured as the encapsulation of the compilation result (for example, cubin binary code on the Nvidia GPU backend), which is responsible for calling the corresponding binary code through the driver API.

[0040] Wherein, the part of Compiler is configured to be implemented by calling a separate compilation script through the Python interpreter embedded in the runtime. The compilation script can reuse the triton.compile compilation interface on the Python side, and the Runtime provides the function body (fn), signature, constant parameters (constants), parameter specialization information (spec), compilation target (target) and compilation options (compile_options) and other information; compiled into binary code and returned to the path of the Runtime. The Runtime constructs the CompiledKernel instance according to the path and stores it in the cache of TritonJITFunction. The embodiment of the present application provides the basic framework required for the function written in triton in C++ programming language, so that the developer can develop the operator library based on the function written in triton, and provides the C++ interface. The embodiment of the present application can reduce the runtime overhead on the basis of retaining the jit running mechanism of triton language.

[0041] As described in the foregoing, please refer to Figure 3The embodiments of the present application can manage the CompiledKernel cache through two-layer instance multiplexing (TritonJITFunction instance multiplexing and CompiledKernel instance multiplexing within TritonJITFunction), to avoid frequent triggering of compilation.

[0042] In the embodiments of the present application, the same function written by Triton can be used in multiple places in a C++ function. If a separate TritonJITFunction instance is created at each place, it cannot be guaranteed that the cache managed thereby can be multiplexed. In order to multiplex the cache, the embodiments of the present application are configured to guarantee that the same TritonJITFunction instance is used at multiple places. To this end, the embodiments of the present application manage and create TritonJITFunction instances based on a Multiton pattern. Please refer to Figure 3 For example, in the embodiments of the present application, a global static hash map is configured to manage all instances created by the type. The hash map takes the file path and function name of the function as the key. As long as the file path and function name of the input are consistent, the same reference to the TritonJITFunction instance is obtained. The Multiton pattern can facilitate the manner of instance creation.

[0043] In optional embodiments of the present application, other manners that can achieve the above purpose can also include creating a corresponding global static object or a corresponding getter function for each TritonJITFunction, wherein a local static object is returned.

[0044] In the preferred embodiments of the present application, the cross-language function calling method can further include configuring a function parameter distribution rule of the Triton just-in-time compilation function based on a static signature rule, including: allocating a corresponding enumeration value to each parameter of the Triton just-in-time compilation function, to use a series of enumeration values as the static signature of the Triton just-in-time compilation function; generating a static signature object by parsing the corresponding Triton function code of the Triton just-in-time compilation function object when the Triton just-in-time compilation function object is constructed; and generating a specialized signature from the input parameters and distributing the input parameters to select the parameters for calling the CompiledKernel when the Triton just-in-time compilation function object is called according to the static signature object.

[0045] The preferred enumerated types of the embodiments of the present application can include CONSTEXPR, SPECIALIZED and NON-CONSTEXPR. Among them, CONSTEXPR identifies a compile-time constant parameter, SPECIALIZED identifies a non-constant parameter whose corresponding value affects the specialization signature, and NON-CONSTEXPR identifies a non-constant parameter whose corresponding value does not affect the specialization signature. Preferably, the processing logic for the enumerated values is as follows: for the parameters marked as CONSTEXPR, the value of the corresponding argument is part of the specialization signature, and the value of the argument is embedded in the CompiledKernel generated by the compilation; for the parameters marked as SPECIALIZED, the features are extracted from the corresponding argument according to the preset conditions at compile time, as part of the specialization signature; for the parameters marked as NON-CONSTEXPR, the value of the argument does not participate in the generation of the specialization signature, and the parameters of the CompiledKernel generated by the compilation include the parameters marked as NON-CONSTEXPR.

[0046] For example, the processing logic for the enumerated values is shown in Table 1.

[0047] Table 1: Enumerated values and their corresponding processing logic

[0048]

[0049]

[0050] When constructing a Triton just-in-time function, the corresponding Triton function code needs to be analyzed. For each parameter, if its type is marked as tl.constexpr, its corresponding enumerated value is CONSTEXPR; otherwise, according to whether the parameter is marked as do_not_specialize in the function decorator, if it is marked as do_not_specialize, the corresponding enumerated value is NON-CONSTEXPR, otherwise, the corresponding enumerated value is SPECIALIZED.

[0051] For example, the Triton function signature and function decorator can be represented as:

[0052] @triton.jit(do_not_specialize=[“n”])

[0053] def add(a_ptr, b_ptr, out_ptr, n, tile_size: tl.constexpr):

[0054] Since the type of the parameter tile_size is marked as tl.constexpr, its corresponding enumeration value is CONSTEXPR; since the parameter n is indicated as do_not_specialize, its corresponding enumeration value is NON-CONSTEXPR, and the rest of the parameters are SPECIALIZED.

[0055] In the preferred embodiment of the present application, the cross-language function calling method can further include implementing polymorphic calling of the Triton just-in-time compilation function based on the variable parameter template, which can include: defining the calling operator of the Triton just-in-time compilation function as a variable parameter template function, so that the operator receives different numbers and types of parameter packages. Among them, different Triton just-in-time compilation functions can accept different numbers and types of parameter packages; and the same Triton just-in-time compilation function can accept different types of parameter packages; at compile time, type deduction is performed on each parameter to determine the type to which each parameter belongs; and based on the determined type of each parameter, the value and the static signature of the Triton just-in-time compilation function, the specialization signature is determined, and the parameters used to call the compiled result CompiledKernel are selected.

[0056] Among them, variable means that different functions can be different, or the same function can also pass different types of parameter combinations. Triton function is the calling operator or calling method of Triton just-in-time compilation function, that is, operator().

[0057] In the embodiments of the present application, the number of parameters required by different Triton function calls is not the same, even if the number of parameters required by two triton jit functions is the same, the parameter distribution logic can be different, and even if it is the same Triton function, different combinations of parameters of different types can be passed in. Therefore, it is necessary to forward the parameters according to a specific logic to call the CompiledKernel. This logic can be implemented by templates or runtime code generation. The preferred embodiments of the present application use C++ variadic templates and expression folding to implement the operator method of TritonJITFunction to support polymorphism. Specifically, the calling method of TritonJITFunction is implemented as a variadic template (Variadic Template), which passes in parameters including fixed and variable parts. The fixed part includes stream, started grid (which can include x, y, z dimensions), compilation options num_warps, num_stages, etc., which are determined by the specific backend. The variable part is defined by the function implemented by the triton language. The part determined by the function implemented by the triton language is declared as a parameter pack (Parameter Pack).

[0058] In the preferred embodiments of the present application, the cross-language function calling method can also include determining the edited result specialization signature based on the folded expression, and the process of selecting the parameters for calling the CompiledKernel can include: using the lambda expression for processing parameters, processing each parameter according to the edited result static signature rule through the folded expression.

[0059] For example, by expanding the parameter pack through the folded expression, using the lambda expression, and processing each parameter according to the static signature rule. Wherein, the processing of each item can include extracting the parameter type, extracting the feature, parameter distribution (determining whether the parameter is used to call the CompiledKernel), etc.

[0060] Please refer to Figure 4 , in the above example, in the code logic, the lambda expression for processing parameters is defined, and then each parameter in the parameter pack is processed using the folded expression. The lambda can be configured to: construct a specialization signature; collect parameters for calling. The configuration logic is as follows:

[0061] 1) Construct a specialization signature: first initialize Signatures, a string vector, to collect the corresponding part of each parameter of the triton jit function in the specialization signature.

[0062] After the vector is initialized, it starts processing according to the static signature and the type of the passed parameters:

[0063] For parameters with enumeration values ​​of SPECIALIZED and NON-CONSTEXPR in the static signature, the format is "{type}{spec}". The rule for type is: if the actual parameter type is tensor, then type is "*{dtype}", where dtype is a string corresponding to the tensor data type; if the parameter type is not tensor, then it is a string corresponding to the parameter type. The rule for spec is: if the static signature of the parameter is SPECIALIZED and its type is an integer, then the format is ":16" or ":1", depending on whether its value is divisible by 16 or equal to 1; if the parameter is a tensor, then the format is ":16" or ":1", depending on whether the value of its data pointer is divisible by 16 or equal to 1. In other cases, spec is empty.

[0064] For a parameter whose enumeration value is CONSTEXPR in a static signature, the format is "{value}", where value is the result of formatting the parameter value.

[0065] Please refer to Figure 4 , such as the triton function corresponding to the called TritonJITFunction is:

[0066] @triton.jit(do_not_specialize=[″n″])

[0067] def add(a_ptr, b_ptr, out_ptr, n, tile_size: tl.constexpr):

[0068] The parameters passed in during the call are: a_ptr, b_ptr, and out_ptr are all float type Tensors, and their data pointers are divisible by 16, n is the value 1000 of int64_t type, tile_size is the value 512 of int64_t, and the extracted specialized signature is "*fp32:16, *fp32:16, *fp32:16, i64, 512".

[0069] 2) Collect parameters for calling the kernel: This process requires maintaining two vectors:

[0070] Pointers to Kenrel Args, void* vector, used to collect pointers to each kernel arg and then used for kernel launch.

[0071] Data Pointers, a void* vector, is used to collect data pointers for tensor type parameters in kernel arguments (kernel args). Because Pointers to Kernel Args includes pointers to data pointers, the lifecycle of the data pointers themselves must be guaranteed, so they are stored in a vector.

[0072] After the above two vectors are initialized, the following processing is performed according to the static signature and the type of the passed parameters:

[0073] Parameters marked as CONSTEXPR are not processed;

[0074] For parameters marked as SPECIALIZED, if their type is integer and their value is 1, no processing is performed;

[0075] In other cases, if the parameter type is tensor, first extract its data pointer and store it in Data Pointers, then take the pointer of the data pointer in Data pointers and store it in Pointers to Kenrel Args. Otherwise, directly take the pointer of the parameter and store it in Pointers to Kenrel Args.

[0076] In a preferred embodiment of the present application, the processing logic of the lambda expression arg_handle for processing parameters can be consistent with the Python-based jit Runtime behavior of Triton.

[0077] Continuing with the above example, since the enumeration values ​​corresponding to a_ptr, b_ptr, and out_ptr are SPECIALIZED and their type is Tensor, their respective data pointers are extracted and stored in Data Pointers, and the pointers to the corresponding data pointers in Data Pointers are taken and stored in Pointers to Kernel Args; for the parameter n, since its corresponding enumeration value is NON-CONSTEXPR, its pointer is directly taken and stored in Pointers to Kernel Args; for the parameter tile_size, since its corresponding enumeration value is CONSTEXPR, it indicates that it is a compile-time constant parameter and is not used to call CompiledKernel, so it is not processed.

[0078] In the preferred embodiments of the present application, the cross-language function calling method can further include: calling the driver API of the device to execute the required Kernel.

[0079] In the above example, after using lambda to process the parameters in the parameter package one by one, the values of the above three vectors are ready. The signatures are spliced to obtain the specialized signature, and the CompiledKernel is queried from the cache. If the required CompiledKernel is not in the cache, the compilation is triggered. Then, the driver API of the device (for example, cuLaunchKernel on nvidia gpu) is called to call the compiled CompiledKernel using Pointers to Kernel Args.

[0080] In the preferred embodiments of the present application, the torch operator implemented using Triton JIT C++ Runtime can include: writing triton jit function code using triton language; defining a function for processing Tensor in C++. Wherein, defining a function for processing Tensor in C++ can include: for the configuration of the Wrapper part, atenAPI can be used, and libtorch is depended; and triton jit function is called using TritonJITFunction.

[0081] In the preferred embodiments of the present application, the torch operator implemented also provides an entry for use on the python side, please refer to Figure 4 The custom TorchLibrary m can be defined by using the macro provided by torch (usage TORCH LIBRARY (library_name, m)), and then the custom operator can be defined by using m.def (op_name,...), and the operator implementation can be registered to the operator table of torch by using m.impl (op_name,...); that is, the operator can be called by using torch.ops.library_name.op_name, wherein, library_name is the library name, and op_name is the operator name.

[0082] Accordingly, the cross-language function calling method provided in the embodiments of the present application, in the runtime based on the first language, calculates a specialization signature according to a Triton just-in-time compiled function and incoming parameters; based on the specialization signature and a hash table from the specialization signature to a CompiledKernel maintained in the cache of the Triton just-in-time compiled function, if the required CompiledKernel is queried, the corresponding CompiledKernel is selected; otherwise, the Triton compiler is triggered to compile, and the compiled CompiledKernel is placed in the cache of the Triton just-in-time compiled function; and the parameters used to call the required CompieldKernel are selected from the incoming parameters of the Triton just-in-time compiled function, and the corresponding CompiledKernel is called according to the required CompiledKernel and the selected CompiledKernel parameters. The embodiments of the present application can use the basic framework required by the triton jit function in the first programming language (for example, the C++ programming language), so that the developer can develop an operator library based on the triton jit function implementation, and provide a C++ interface. The embodiments of the present application reduce the runtime overhead of TritonJITFunction on the basis of retaining the just-in-time compilation running mechanism of the triton language. And since the runtime is implemented on the C++ side, other parts of the operator (for example, metadata calculation, parameter preprocessing, output allocation, etc.) can be developed using C++, which can reduce the overall running overhead of the operator.

[0083] The embodiments of the present application can maintain the feature of using Triton language to write functions and calling them on the C++ side, without the need to re-implement the Triton language as a DSL embedded in C++ in order to use the functions of the Triton language on the C++ side. The implementation type of the function can include:

[0084] 1) Calculate the specialization signature and obtain the CompiledKernel (extract information from specific incoming parameters, such as the data type of the Tensor type parameter, the type of the non-Tensor parameter, the specialization information of part of the integer type parameter and the Tensors data pointer parameter, the value of the compiler constant parameter, and the specialization information of the specific CompiledKernel); trigger compilation and cache the compiled Kernel.

[0085] 2) Parameter distribution. Some parameters are used to compile a specific CompiledKernel, such as constant parameters and parameters and compilation options for specializing a Kernel, and some parameters are passed to the compiled Kernel; parameters with the value of the enumeration CONSTEXPR in the function static signature are constant parameters at compile time, which are used to specialize a specific CompiledKernel rather than parameters at runtime of the CompiledKernel; parameters with the value of the enumeration NON-CONSTEXPR in the function static signature indicate that the value of the parameter is not used to specialize the CompiledKernel. Parameters with the value of the enumeration SPECIALIZED in the function static signature are used to specialize the CompiledKernel.

[0086] 3) Calling the driver API to execute the editing result corresponding to the selected CompiledKernel.

[0087] Since the compiler and the runtime are separable in the just-in-time compilation mechanism, embodiments of the present application do not need to make substantial modifications to the Triton compiler, and can reuse the compilation interface (triton.compile) of the Triton language itself, and only need to supplement necessary auxiliary functions to enable the editing result to be triggered from the C++ runtime and used by the runtime.

[0088] For calling triton.compile, embodiments of the present application embed a Python interpreter in the runtime to call the compiler interface triton.compile of Triton. Embodiments of the present application provide necessary wrapping, and the triton jitC++ runtime provides a specialized signature to call triton.compile; then, the path of the editing result on the file system is returned to the triton jit C++ runtime. Therefore, embodiments of the present application can reuse the file cache managed by triton.compile.

[0089] For implementing the operator in C++ using Triton jit c++ runtime, the operator needs to provide an easy-to-use interface for specific multi-dimensional array library, and needs to use multi-dimensional array as the main input and output type. Therefore, when implementing the operator, the metadata access to the array, the construction of empty array and other interfaces provided by the multi-dimensional array library can be used, and different codes can be used according to the specific array library that needs to be adapted. For operator writing, through the triton jit C++ runtime provided by the embodiment of the present application, the developer can construct a TritonJITFunction instance on the C++ side, and compile and execute on demand through just-in-time compilation. Therefore, the embodiment of the present application provides a framework for implementing the operator using the function written in the triton language on the C++ side. In addition to calling TritonJITFunction, the part of the operator implementation can use the functional interfaces provided by the high-dimensional array library. For example, the tensor metadata query, metadata manipulation, and new empty tensor operations provided in the torch library. The performance can be higher than that of implementing equivalent logic on the python side.

[0090] The embodiment of the present application also provides a control device, the control device comprising a memory, a processor and a computer program stored on the memory and executable on the processor, and the processor executes the computer program to implement the cross-language function calling method described above.

[0091] The embodiment of the present application also provides a machine-readable storage medium, the machine-readable storage medium stores instructions, and the instructions make the machine execute the cross-language function calling method described above.

[0092] It should be noted that the control device and the machine-readable storage medium described above can implement the cross-language function calling method provided by the above embodiments, and the specific implementation manner can be referred to the description of the cross-language function calling method in the above embodiments, which will not be described here.

[0093] It can be understood that the circuit structure, name and parameter described in the above embodiments are only examples. Those skilled in the art can also easily combine and adjust the structural characteristics of the above multiple embodiments according to the use needs, and the concept of the present application should not be limited to the specific details of the above examples.

[0094] Although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that the technical solutions recorded in the foregoing embodiments can be modified, or some technical features can be replaced by equivalents; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. A cross-language function calling method based on just-in-time compilation, characterized in that: The cross-language function calling method includes: At runtime based on the first language side, the specialized signature is calculated based on the Triton just-in-time compiled function and the passed parameters; Based on the specialized signature and the hash table from the specialized signature to the CompiledKernel maintained in the cache of the Triton just-in-time compilation function, if the required CompiledKernel is found, the corresponding CompiledKernel is selected; Otherwise, trigger the Triton compiler to compile and place the compiled CompiledKernel into the cache of the Triton just-in-time compilation function; and Select the parameters for calling the required CompieldKernel from the incoming parameters of the Triton just-in-time compilation function, and call the corresponding CompiledKernel according to the required CompiledKernel and the selected CompiledKernel parameters.

2. The cross-language function calling method according to claim 1, characterized in that: Triggering the Triton compiler to compile includes: Calling the Triton compilation interface through the Python interpreter embedded in the runtime of the first language side; and Pass the specialization signature to the Triton compilation interface, compile using the Triton compiler, and package the compilation result into a CompiledKernel. The Triton compiler is equipped with a file cache mechanism, which specifies the path where the corresponding compilation results are stored in the file system according to the CompiledKernel that needs to be compiled. If the compilation result exists in the specified path, the path is returned. Otherwise, trigger the compilation first and save the compilation results to the specified path.

3. The cross-language function calling method according to claim 1, wherein: The cross-language function calling method further includes managing the instantiation of the Triton just-in-time compilation function through a multi-instance mode, including: Create a global static hash map to manage instantiated Triton just-in-time compilation function objects; Combining the file path of the Triton just-in-time compilation function with the function name to generate a unique identifier as the key of the hash map, and the corresponding Triton just-in-time compilation function as the value of the hash map; When a request is made to instantiate a Triton just-in-time compilation function, a hash map is queried in the global static hash map using the unique identifier. If the Triton just-in-time compilation function corresponding to the unique identifier is not in the queried hash map, a corresponding Triton just-in-time compilation function is created and stored in the queried hash map; and By reusing instances of Triton just-in-time compilation functions, the cache of the Triton just-in-time compilation functions is managed.

4. The cross-language function calling method according to claim 1, characterized in that: The cross-language function calling method further includes configuring the function parameter distribution rules of the Triton just-in-time compilation function based on the static signature rules, including: Assigning a corresponding enumeration value to each formal parameter of the Triton just-in-time compilation function, so as to use a series of enumeration values ​​as the static signature of the Triton just-in-time compilation function; When the Triton just-in-time compilation function object is constructed, a static signature object is generated by parsing its corresponding Triton function code; and When the Triton just-in-time compilation function object is called, the specialized signature is generated through the incoming parameters according to the static signature object, and the incoming parameters are distributed to select the parameters for calling CompiledKernel.

5. The cross-language function calling method according to claim 4, characterized in that: The enumeration values ​​include CONSTEXPR, SPECIALIZED and NON-CONSTEXPR, Among them, CONSTEXPR identifies compile-time constant parameters, SPECIALIZED identifies corresponding values ​​that affect non-constant parameters of specialized signatures, and NON-CONSTEXPR identifies corresponding values ​​that do not affect non-constant parameters of specialized signatures. The processing logic for the enumeration value is: For parameters marked as CONSTEXPR, the value of the corresponding actual parameter will be used as part of the specialization signature, and the value of the actual parameter will be embedded in the CompiledKernel generated by the compilation. For parameters marked as SPECIALIZED, at compile time, features are extracted from the corresponding actual parameters according to preset conditions as part of the specialization signature. For parameters marked as NON-CONSTEXPR, the values ​​of the actual parameters do not participate in the generation of the specialization signature, and the parameters of the compiled CompiledKernel include the parameters marked as NON-CONSTEXPR.

6. The cross-language function calling method according to claim 4, characterized in that: The cross-language function calling method further includes implementing polymorphic calling of the Triton just-in-time compilation function based on a variable parameter template, including: Defining the call operator of the Triton just-in-time compilation function as a variable parameter template function so that the Triton just-in-time compilation function object receives parameter packs of different numbers and types; During compilation, type deduction is performed on each parameter to determine the type of each parameter; and Based on the determined type and value of each parameter and the static signature of the Triton just-in-time compilation function, the specialized signature is determined, and the parameters for calling the CompiledKernel are selected.

7. The cross-language function calling method according to claim 4, characterized in that: The cross-language function calling method further includes determining the specialization signature based on a fold expression, and selecting parameters for calling the CompiledKernel, including: By folding the expression, a lambda expression is used to process the parameters, and each parameter is processed one by one according to the static signature rules.

8. The cross-language function calling method according to claim 5, characterized in that: The cross-language function calling method further includes: Call the device driver API to execute the required Kernel.

9. A control device, characterized in that: The control device includes: a memory, a processor, and a computer program stored in the memory and executable on the processor. The processor executes the computer program to implement the cross-language function calling method according to any one of claims 1 to 8.

10. A machine-readable storage medium, characterized in that The machine-readable storage medium stores instructions, which enable a machine to execute the cross-language function calling method according to any one of claims 1 to 8.

Images