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

Abstract

The invention provides a cross-language function calling method and device based on just-in-time compiling and a storage medium, and belongs to the technical field of computers. The method comprises the steps of calculating a specialized signature according to a Triton instant compilation function and an incoming parameter during operation based on a first language end; based on the specialized signature and a hash table which is maintained in a cache of the Triton instant compilation function and is from the specialized signature to the CompiledKernel, if the required CompiledKernel is queried, selecting the corresponding CompiledKernel, and if the required CompiledKernel is queried, selecting the corresponding CompiledKernel; otherwise, a Triton compiler is triggered to compile, and the compiled CompiledKernel is placed in the cache of the Triton instant compiling function; and selecting parameters used for calling the required CompieldKernel from the incoming parameters, and calling the corresponding CompieldKernel. The method comprises the following steps of: selecting parameters used for calling the required CompieldKernel from the incoming parameters, and calling the corresponding CompieldKernel. According to the embodiment of the invention, a basic framework required by the triton jit function can be used in the first programming language, so that a developer can develop an operator library realized based on the triton jit function, and a C + + interface is provided. According to the embodiment of the invention, on the basis of keeping a jit operation mechanism of the triton language, the overhead during operation is reduced.

Description

Technical Field

[0001] The present application relates to the field of computer technology, and in particular to a cross-language function calling method, device, and storage medium based on just-in-time compilation. Background Art

[0002] To meet the flexible demands of deep learning research and deployment, a number of Python-based embedded domain-specific languages (DSLs) or libraries have emerged, allowing users to write kernels for GPUs and other accelerators in a Python-like language. Triton Lang, one such language, only offers Python-based usage. Therefore, operators developed using Triton generally only provide Python interfaces. An operator typically consists of calling one or more kernel functions, supplemented by necessary parameter checking and preprocessing, starting configuration calculations, and wrapping them into a more user-friendly wrapper function. These functions typically accept Tensor objects provided by common high-dimensional array libraries as parameters (for example, array types in torch and jax). However, since the logic between the wrapper and the kernel call is executed in Python, the inefficiency of the Python interpreter can become a limiting factor in overall performance. Even if the kernel startup is asynchronous, control returns to the host thread after the kernel startup without waiting for the kernel execution to complete, but the wrapper takes too long or even exceeds the kernel execution time, it will still cause the GPU to be idle.

[0003] The main way to use the Triton language is just-in-time compilation: functions written in this language need to extract necessary information from the incoming parameters at runtime, construct specialized signatures, and have the compiler compile the required binary files just in time and then execute them at runtime. Since the Triton language is implemented as an embedded DSL for Python, the jit runtime (Runtime) is often also implemented on the Python side. The runtime is responsible for providing enough information to trigger the compiler to perform just-in-time compilation and execute the compiled results. The runtime may also be affected by the low efficiency of the Python interpreter. In addition, the compiler of the Triton language only provides interfaces on the Python side. Therefore, the problems with the relevant technical solutions include: the runtime of the Triton language depends on the Python interpreter and execution environment, which is not conducive to functions written in the Triton language providing interfaces in other languages; the low efficiency of the Python interpreter affects the overall execution efficiency. Summary of the Invention

[0004] The present application provides a cross-language function calling method, device and storage medium based on just-in-time compilation, which aims to solve the problem that the existing technical solutions rely on the Python interpreter and execution environment during the runtime of the Triton language, which is not conducive to providing interfaces of other languages for functions written in the 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: calculating a specialized signature based on the Triton just-in-time compilation function and incoming parameters during runtime based on a first language end; based on the specialized signature and a 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, selecting the corresponding CompiledKernel; otherwise, triggering Triton compiler compilation, 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 based on the required CompiledKernel and the selected CompiledKernel parameters.

[0006] In an optional embodiment, the triggering of the Triton compiler compilation includes: calling the Triton compilation interface through the Python interpreter embedded in the runtime of the first language end; and passing the specialized signature to the Triton compilation interface, compiling using the Triton compiler, and packaging the compilation result into a CompiledKernel. Among them, the Triton compiler is configured with a file caching mechanism. According to the CompiledKernel compiled as needed, the path where the corresponding compilation result is stored in the file system 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 also includes managing the instantiation of the Triton just-in-time compilation function through a multi-instance mode, including: creating a global static hash mapping table for managing the instantiated Triton just-in-time compilation function objects; combining the file path where the Triton just-in-time compilation function is located with the function name to generate a unique identifier as the key of the hash mapping, and the corresponding Triton just-in-time compilation function as the value of the hash mapping; when requesting to instantiate the Triton just-in-time compilation function, 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, then creating a corresponding Triton just-in-time compilation function and storing it in the query hash mapping; and managing the cache of the Triton just-in-time compilation function by reusing instances of the Triton just-in-time compilation function.

[0008] In an optional embodiment, the cross-language function calling method also includes configuring function parameter distribution rules of the Triton just-in-time compilation function based on 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 parameters for calling CompiledKernel.

[0009] In an optional embodiment, the enumeration values include CONSTEXPR, SPECIALIZED and NON-CONSTEXPR. Among them, CONSTEXPR identifies a compile-time constant parameter, SPECIALIZED identifies that the corresponding value will affect the non-constant parameter of the specialization signature, and NON-CONSTEXPR identifies that the corresponding value will not affect the non-constant parameter of the specialization signature. The processing logic for the enumeration values is that for the formal parameters marked as CONSTEXPR, the values of the corresponding actual parameters will be used as part of the specialization signature, and the values of the actual parameters will be embedded in the CompiledKernel generated by the compilation. For the formal parameters marked as SPECIALIZED, during compilation, features are extracted from the corresponding actual parameters according to preset conditions as part of the specialization signature. For the formal parameters marked as NON-CONSTEXPR, the values of the actual parameters do not participate in the generation of the specialization signature, and the formal parameters of the CompiledKernel generated by the compilation include the formal parameters marked as NON-CONSTEXPR.

[0010] In an optional embodiment, the cross-language function calling method also includes implementing polymorphic calling of the Triton just-in-time compilation function based on a variable parameter template, including: defining the calling 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 to support calls to different Triton functions, as well as calls to the same Triton function using different parameter combinations; at compile time, performing type deduction on each parameter to determine the type of each parameter; and determining the specialized signature based on the determined type, value and static signature of each parameter of the Triton just-in-time compilation function, and selecting parameters for calling the CompiledKernel.

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

[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, comprising: a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein 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, on which instructions are stored, and the instructions enable a machine to execute the above-mentioned cross-language function calling method.

[0015] Through the above technical solution, the cross-language function calling method provided by the embodiment of the present application, when running based on the first language end, calculates the specialized signature according to the Triton just-in-time compilation function and the incoming 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 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 compilation function; and the parameters for calling the required CompieldKernel are selected from the incoming parameters of the Triton just-in-time compilation function, and the corresponding CompiledKernel is called according to the required CompiledKernel and the selected CompiledKernel parameters. The embodiment of the present application can use the basic framework required for triton jit function in the first programming language (for example, C++ programming language), so that developers can develop an operator library based on triton jit function implementation, and provide a C++ interface. The embodiment of the present application reduces its runtime overhead while retaining the jit operating mechanism of the triton language.

[0016] Other features and advantages of the present application will be described in the following description, and in part will become apparent from the description, or will be understood by practicing the present application. The purpose and other advantages of the present application can be achieved and obtained through the structures and processes indicated in the description and the drawings. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] In order to more clearly illustrate the technical solutions in the embodiments of the present application or related technologies, the following is a brief introduction to the drawings required for use in the embodiments or related technical descriptions. It is obvious that the drawings described below are certain embodiments of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without any creative work.

[0018] Figure 1 This is a flow chart of a cross-language function calling method based on just-in-time compilation provided in an embodiment of the present application.

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

[0020] Figure 3 This is a schematic diagram of the process of obtaining the Triton just-in-time compilation function instance provided in the embodiments of this application.

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

[0022] In order to make the purpose, technical solutions and advantages of the embodiments of this application more clear, the technical solutions in the embodiments of this application will be clearly and completely described below in conjunction with the drawings in the embodiments of this application. Obviously, the described embodiments are part of the embodiments of this application, not all of the embodiments. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of this application.

[0023] As mentioned above, since the Triton language is implemented as an embedded DSL of Python, its just-in-time compilation (Just-in-time Compilation) runtime is also implemented on the Python side. The runtime is responsible for providing sufficient information to trigger the compiler to perform just-in-time compilation and execute the compiled results. The runtime may also be affected by the low efficiency of the Python interpreter. In addition, the compiler of the Triton language also only provides an interface on the Python side. Therefore, the problems existing in the relevant technical solutions include: the runtime of the Triton language depends on the Python interpreter and execution environment, which is not conducive to the functions written in the Triton language to provide interfaces in other languages; the low efficiency of the Python interpreter affects the overall execution efficiency.

[0024] In this regard, an embodiment of the present application provides a cross-language function calling method based on just-in-time compilation, which provides a solution for using functions written in an external language (for example, Triton Lang) through just-in-time compilation in a first language, that is, a low-level language (for example, C++ language), so that developers can develop an operator library based on the Triton just-in-time compilation function (TritonJITFunction), and provide a low-level language interface, and reduce its runtime overhead while retaining the JIT operating mechanism of the Triton language.

[0025] Figure 1 This is a flow chart of a cross-language function call method based on just-in-time compilation provided by an embodiment of the present application. The method may include the following steps:

[0026] Step S110: During the runtime based on the first language end, the specialized signature is calculated according to the Triton just-in-time compilation function and the incoming parameters.

[0027] The Triton just-in-time compilation function is a type defined at runtime. It is an encapsulation of a function written in the Triton language and can be called by specifying the file path and function name of the corresponding function. CompiledKernel is an encapsulation of the compilation result of a function written in the Triton language and is constructed by specifying the file path and function name of the compilation result. CompiledKernel is a type defined at runtime and can be called.

[0028] In the embodiment of the present application, the torch operator is implemented on the first language side (e.g., C++ side) using the Triton JIT C++ runtime. Figure 2 For example, at runtime, a specialization signature is calculated based on the Triton JITFunction, incoming parameters, runtime environment, compilation target, etc., and the corresponding CompiledKernel can be selected based on the specialization signature. When the Triton JIT C++ runtime is running, it can be implemented as a C++ library, and its core configuration types are TritonJITFunction and CompiledKernel.

[0029] In an embodiment of the present application, the cross-language function calling method may also include managing the instantiation of the Triton just-in-time compilation function through a multi-instance mode, including: creating a global static hash mapping table for managing the instantiated Triton just-in-time compilation function objects: combining the file path where the Triton just-in-time compilation function is located with the function name to generate a unique identifier as the key of the hash mapping, and the corresponding Triton just-in-time compilation function as the value of the editing result hash mapping; when requesting to instantiate the Triton just-in-time compilation function, 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, then creating the corresponding Triton just-in-time compilation function and storing it in the query hash mapping; and managing the cache of the Triton just-in-time compilation function by reusing instances of the Triton just-in-time compilation function.

[0030] Please refer to Figure 3For example, regarding the construction of the TritonJITFunction instance, on the C++ side, provide the file path and function name of the function implemented in the triton language to obtain the corresponding TritonJITFunction instance. TritonJITFunction uses the multiton mode to manage its instances. All TritonJITFunctions exist in a global hash table with the concatenation of the file path and function name of the function as the key and TritonJITFunction as the value. When providing the file path and function name of the function implemented in the triton language to obtain the TritonJITFunction instance, first concatenate the path and function name as the key and query the hash table. If the key does not exist in the hash table, first create a TritonJITFunction instance and store it in the hash table. Otherwise, directly return a reference to the queried TritonJITFunction instance.

[0031] In an embodiment of the present application, for the call of the TritonJITFunction instance, the incoming parameters include a fixed part and a variable part. The fixed part includes stream, the started grid (which may include x, y, z dimensions), compilation options num_warps, num_stages, etc., which can be determined by the backend. The variable part is defined by the function implemented in the corresponding triton language. Because the calling method of TritonJITFunction is implemented as a variable parameter template (which will be introduced in detail later, so I will not go into details here), variable parameters can be supported.

[0032] Step S120: 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.

[0033] In the embodiment of the present application, the call of Triton's instant compilation function is implemented by specializing and calling CompiledKernel. Figure 2For example, the information required to specialize CompiledKernel can come from the incoming parameters, the runtime environment, the compilation target, etc. The information extracted from the source and used to determine the CompiledKernel is referred to as the specialization signature. The TritonJITFunction object is preferably configured to maintain a hash table. Each TritonJITFunction instance maintains a corresponding cache, which is implemented as a hash table, where the key is the specialization signature and the value is the CompiledKernel, which is used to quickly find the required CompiledKernel based on the specialization signature.

[0034] Step S130: Otherwise, trigger the Triton compiler to compile, and put the compiled CompiledKernel into the cache of the Triton just-in-time compilation function.

[0035] In a preferred embodiment of the present application, triggering the Triton compiler to compile may include: calling the Triton compilation interface through the Python interpreter embedded in the runtime of the first language end; and passing the specialized signature to the Triton compilation interface, compiling using the Triton compiler, and packaging the compilation result into a CompiledKernel. Among them, the Triton compiler is configured with a file caching mechanism. According to the CompiledKernel to be compiled, the path where the corresponding compilation result is stored in the file system 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.

[0036] Please refer to Figure 2 For example, in an embodiment of the present application, the Runtime is implemented on the C++ side, and the Compiler part can reuse the Triton compiler. The embodiment of the present application can configure a separate compilation script to call the Triton compilation interface (triton.compile). For example, when the required CompiledKernel is not in the cache of the Triton just-in-time compilation 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 its editing result can be used to construct a CompiledKernel instance according to the path and store it in the cache of the Triton just-in-time compilation function.

[0037] Step S140: Select parameters for calling the required CompiledKernel from the input parameters of the Triton just-in-time compilation function, and call the corresponding CompiledKernel according to the required CompiledKernel and the selected CompiledKernel parameters.

[0038] Please refer to Figure 2 , select the parameters required for calling the selected CompiledKernel from the incoming parameters when calling the Triton just-in-time compilation function, and call the corresponding CompiledKernel to implement the execution of Triton just-in-time compilation on the C++ side.

[0039] An embodiment of the present application provides a cross-language calling method based on just-in-time compilation, which can be a Runtime part and a Compiler part. Among them, the core types provided by Runtime are TritonJITFunction and CompiledKernel. TritonJITFunction (Triton just-in-time compilation function) corresponds to a function written in the triton language. Each TritonJITFunction instance maintains a corresponding cache, which is implemented as a hash mapping table from specialized signatures to CompiledKernel; TritonJITFunction is responsible for calculating specialized signatures, CompiledKernel selection, selecting parameters for calling Kernel, and Kernel calls. CompiledKernel can be configured to encapsulate the compilation results (for example, cubin binary code on the Nvidia GPU backend), and is responsible for calling the corresponding binary code through the driver API.

[0040] Wherein, the Compiler part is configured to be realized by calling a separate compilation script through the Python interpreter embedded in the runtime. The compilation script can reuse the triton.compile compilation interface of the python end, and the Runtime provides function body (fn), signature (signature), constant parameters (constants), parameter specialization information (spec), compilation target (target) and compilation options (compile_options) and other information; Compiled into binary code and return path to Runtime. Runtime constructs 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 the C++ programming language, so that developers can develop operator libraries based on the functions written in triton, and provide C++ interface. The embodiment of the present application can reduce its runtime overhead on the basis of retaining the jit operation mechanism of the triton language.

[0041] As mentioned above, see Figure 3, the embodiment of the present application can manage the CompiledKernel cache through two layers of instance reuse (TritonJITFunction instance reuse and CompiledKernel instance reuse within TritonJITFunction) to avoid frequent triggering of compilation.

[0042] In the embodiment of the present application, since the same function written by Triton may be used in multiple places in the C++ function, if a separate TritonJITFunction instance is created in each place, it cannot be guaranteed that the cache it manages can be reused. In order for the cache to be reused, the embodiment of the present application is configured to ensure that the same TritonJITFunction instance is used in multiple uses. To this end, the embodiment of the present application is based on the multi-instance mode (Multiton mode) to manage and create TritonJITFunction instances. Please refer to Figure 3 For example, in this embodiment of the application, a global static hash map is configured to manage all instances created by this type. The hash map uses the file path and function name of the function as the key. As long as the file path and function name passed in are the same, a reference to the same TritonJITFunction instance will be obtained. The Multiton mode can facilitate the management of instance creation.

[0043] In an optional implementation of the present application, other ways to achieve the above purpose may also include creating a corresponding global static object or creating a corresponding getter function for each TritonJITFunction, wherein a local static object is returned.

[0044] In a preferred embodiment of the present application, the cross-language function calling method may also include 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, a 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 the CompiledKernel.

[0045] The preferred enumeration types of the embodiment of the present application may include CONSTEXPR, SPECIALIZED and NON-CONSTEXPR. Among them, CONSTEXPR identifies a compile-time constant parameter, SPECIALIZED identifies a non-constant parameter whose corresponding value will affect the specialized signature, and NON-CONSTEXPR identifies a non-constant parameter whose corresponding value will not affect the specialized signature. Preferably, the processing logic for the enumeration value is: for the formal parameter marked as CONSTEXPR, the value of the corresponding actual parameter will be used as part of the specialized signature, and the value of the actual parameter will be embedded in the CompiledKernel generated by compilation. For the formal parameter marked as SPECIALIZED, at compile time, features are extracted from the corresponding actual parameter according to preset conditions as part of the specialized signature. For the formal parameter marked as NON-CONSTEXPR, the value of the actual parameter does not participate in the generation of the specialized signature, and the formal parameters of the CompiledKernel generated by compilation include the formal parameter marked as NON-CONSTEXPR.

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

[0047] Table 1 Enumeration values and their corresponding processing logic

[0048]

[0049]

[0050] When constructing a Triton just-in-time compilation function, it is necessary to analyze the corresponding Triton function code. For each formal parameter, if its type is marked as tl.constexpr, the corresponding enumeration value is CONSTEXPR; otherwise, depending on whether the parameter is marked as do_not_specialize in the function decorator, if it is marked as do_not_specialize, the corresponding enumeration value is NON-CONSTEXPR, otherwise, the corresponding enumeration value is SPECIALIZED.

[0051] For example, the Triton function signature and function decorator can be expressed 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 specified as do_not_specialize, its corresponding enumeration value is NON-CONSTEXPR, and the remaining parameters are SPECIALIZED.

[0055] In a preferred embodiment of the present application, the cross-language function call method may also include implementing a polymorphic call of the edited result Triton just-in-time compilation function based on a variable parameter template, which may include: defining the call operator of the Triton just-in-time compilation function as a variable parameter template function, so that the operator receives parameter packs of different numbers and types. Among them, for different Triton just-in-time compilation functions, different numbers and types of parameter packs can be accepted; and different types of parameter packs can be accepted for the same Triton just-in-time compilation function; at compile time, 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, a specialized signature is determined, and parameters for calling the edited result CompiledKernel are selected.

[0056] The variable means that different functions can have different parameter combinations, or the same function can also pass different types of parameter combinations. The Triton function is the call operator or call method of the Triton just-in-time compiled function, that is, operator().

[0057] In an embodiment of the present application, since the number of parameters required for calling different Triton functions is not the same, even if the number of parameters required for two triton jit functions is the same, their parameter distribution logic may also be different. Even if it is the same Triton function, a combination of different types of parameters can be passed in. Therefore, it is necessary to forward parameters according to a specific logic and call CompiledKernel. This logic can be implemented by templates or runtime code generation. The preferred embodiment of the present application uses C++ variable parameter templates and expression folding to implement the operator method of TritonJITFunction to support polymorphism. Specifically, the calling method of TritonJITFunction is implemented as a variable parameter template (Variadic Template), and its incoming parameters include a fixed part and a variable part. The fixed part includes stream, the started grid (which may 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 in the corresponding triton language. Some parameters determined by the function implemented by the triton language are declared as a parameter pack (Parameter Pack).

[0058] In a preferred embodiment of the present application, the cross-language function calling method may also include determining the specialized signature of the editing result based on the folding expression, and selecting the parameters for calling the CompiledKernel, which may include: using the lambda expression for processing parameters through the folding expression, and processing each parameter item by item according to the static signature rules of the editing result.

[0059] For example, by expanding the parameter pack through a folding expression, using lambda expressions, and following static signature rules, each parameter is processed item by item. This item-by-item processing can include extracting parameter types, extracting features, and parameter distribution (determining whether the parameter is used to call the CompiledKernel).

[0060] Please refer to Figure 4 Continuing with the above example, in the code logic, define a lambda expression for processing parameters, and then use fold expressions to process each parameter in the parameter pack one by one. This lambda can be configured to: construct a specialized signature; collect parameters for invocation. The configuration logic is as follows:

[0061] 1) Construct a specialized signature: First, initialize the Signatures string vector to collect the corresponding part of each parameter of tritonjitfunction in the specialized 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 a preferred embodiment of the present application, the cross-language function calling method may further include: calling a driver API of the device to execute the required Kernel.

[0079] Continuing with the above example, after using lambda to process the parameters in the parameter pack one by one, the values of the three vectors are prepared. Specialized signatures are concatenated using signatures, and the CompiledKernel is retrieved from the cache. If the required CompiledKernel is not in the cache, compilation is triggered. The device driver API (for example, cuLaunchKernel on NVIDIA GPUs) is then called using Pointers to Kernel Args to invoke the compiled CompiledKernel.

[0080] In a preferred embodiment of the present application, using Triton JIT C++ Runtime to implement a torch operator may include: writing the triton jit function code in the triton language; defining a function for processing tensors in C++. Defining the function for processing tensors in C++ may include: configuring the wrapper part using atenAPI, relying on libtorch; and using TritonJITFunction to call the triton jit function.

[0081] In the preferred embodiment of this application, the implemented torch operator also provides a python end use entry, please refer to Figure 4 , you can define a custom TorchLibrary m through the macros provided by torch (the usage is TORCH LIBRARY(library_name, m), then define a custom operator through m.def(op_name, ...), and register the operator implementation to the torch operator table through m.impl(op_name, ...); it can be called through torch.ops.library_name.op_name, where library_name is the library name and op_name is the operator name.

[0082] Accordingly, the cross-language function calling method provided by the embodiment of the present application calculates the specialized signature based on the Triton just-in-time compilation function and the incoming parameters during the runtime based on the first language end; 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 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 compilation function; and the parameters for calling the required CompieldKernel are selected from the incoming parameters of the Triton just-in-time compilation function, and the corresponding CompiledKernel is called according to the required CompiledKernel and the selected CompiledKernel parameters. The embodiment of the present application can use the basic framework required for the triton jit function in the first programming language (for example, the C++ programming language) so that developers can develop an operator library based on the triton jit function implementation and provide a C++ interface. While retaining the Triton language's just-in-time compilation and execution mechanism, the present embodiment reduces the runtime overhead of TritonJITFunction. Furthermore, since the runtime is implemented in C++, the rest of the operator (e.g., metadata calculation, parameter preprocessing, output allocation, etc.) can be developed in C++, reducing the overall runtime overhead of the operator.

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

[0084] 1) Calculate the specialization signature and obtain the CompiledKernel (extract information from the specific incoming parameters, such as the data type of Tensor type parameters, the type of non-Tensor parameters, the specialization information of some integer parameters and Tensor data pointer parameters, and the values of compiler constant parameters, to specialize 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 the kernel, while some parameters are passed to the compiled kernel. Parameters with the CONSTEXPR enumeration value in the function's static signature are compile-time constant parameters used to specialize a specific CompiledKernel, rather than CompiledKernel runtime parameters. Parameters with the NON-CONSTEXPR enumeration value in the function's static signature indicate that the parameter value is not used to specialize the CompiledKernel. Parameters with the SPECIALIZED enumeration value in the function's static signature are used to specialize the CompiledKernel.

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

[0087] Since the compiler and runtime are separable in the just-in-time compilation mechanism, the embodiment of the present application does not require major modifications to the Triton compiler. The compilation interface (triton.compile) of the Triton language itself can be reused. It only needs to add necessary auxiliary functions so that it can be triggered from the C++ runtime and the editing results can be handed over to the runtime for use.

[0088] For calling triton.compile, the embodiment of the present application embeds a Python interpreter in the runtime and calls Triton's compiler interface triton.compile. The embodiment of the present application provides the necessary packaging, with the triton jitC++ runtime providing a specialized signature to call triton.compile; then, the path of the edit result on the file system is returned to the triton jit C++ runtime. Therefore, the embodiment of the present application can reuse the file cache managed by triton.compile.

[0089] For implementing operators in C++ using the Triton jit c++ runtime, the operator needs to provide an easy-to-use interface for a specific multidimensional array library, and needs to use multidimensional arrays as the main input and output types. Therefore, when implementing the operator, you can use the interfaces provided by the multidimensional array library to access the array metadata, construct empty arrays, etc., and use different codes 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, developers can construct TritonJITFunction instances on the C++ side, compile and execute them on demand through just-in-time compilation. Therefore, the embodiment of the present application provides a framework for implementing operators using functions written in the triton language on the C++ side. In addition to calling TritonJITFunction, the part of the operator implementation other than calling TritonJITFunction can use the functional interface provided by the high-dimensional array library. For example, operations such as tensor metadata query, metadata manipulation, and new empty tensor provided in the torch library. It can achieve higher performance than implementing equivalent logic on the python side.

[0090] An embodiment of the present application further provides a control device, which 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 above-mentioned cross-language function calling method.

[0091] An embodiment of the present application further provides a machine-readable storage medium having stored thereon instructions, which enable a machine to execute the above-mentioned cross-language function calling method.

[0092] It should be noted that the above-mentioned control device and machine-readable storage medium can implement the cross-language function calling method provided in the above-mentioned embodiment. The specific implementation method can be found in the description of the cross-language function calling method in the above-mentioned embodiment, which will not be repeated here.

[0093] It is understood that the circuit structures, names, and parameters described in the above embodiments are merely examples. Those skilled in the art may also readily conceive of combinations and adjustments to the structural features of the above embodiments as needed, and should not limit the concept of this application to the specific details of the above examples.

[0094] Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein; and these modifications or replacements do not cause the essence of the corresponding technical solutions to 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