Tensor kernel automatic lifting and cross-backend generation method and device

Abstract

The application discloses a kind of automatic promotion and cross backend generation method and device of tensor kernel, belong to GPU heterogeneous computing field.This method includes: obtaining target CUDA tensor kernel and converting into structured SIMT intermediate representation;According to thread block execution domain, build block context;For thread related branch and loop, generate corresponding predicate context and shared driving loop;Generate block intermediate representation, and through access image analysis, memory to register promotion, reduction and matrix multiplication normalization, SLP packing and layout allocation etc.Optimization, the block level calculation implied in low-level SIMT scheduling is restored to explicit block operation;The block level program after optimization is output to block backend generation target code.The application can reduce the artificial rewriting work required for the migration of CUDA tensor kernel to block level programming model, and improve the portability and optimization space of existing GPU kernel on different backends and hardware platforms.

Description

Technical Field

[0001] This invention relates to the fields of program analysis, compilation optimization, and GPU heterogeneous computing, and particularly to a method and apparatus for automatically promoting CUDA tensor kernels into block-level programs and outputting them to multiple block-level backends. Specifically, it relates to a method and apparatus for automatic promotion and cross-backend generation of tensor kernels. Background Technology

[0002] In deep learning systems and scientific computing programs, a large number of performance-critical paths rely on manually written CUDA tensor kernels. These kernels typically directly expose thread blocks, thread indices, shared memory, synchronization barriers, warp-level communication, and hardware-specific memory access patterns. While developers gain performance through these low-level mechanisms, they also create a significant gap between program logic and the intent of portable tensor computation. In recent years, block-level programming models such as Triton and TileLang have elevated the programmer's expression granularity from single threads to tensor blocks, enabling the backend to uniformly handle vectorized memory access, matrix instructions, asynchronous transport, and target hardware-specific scheduling. However, migrating existing CUDA kernels to block-level programs usually requires developers to manually read the original thread mappings, shared memory temporarys, and loop accumulation logic, and then rewrite equivalent block-level code, which is costly and prone to errors. Existing source code migration or redirection methods mostly retain the original CUDA thread-by-thread execution structure, primarily replacing interfaces, runtime, or target platform instructions, making it difficult to restore the block-level tensor operations implicit in the original program's thread-block collaborative computation.

[0003] Therefore, there is an urgent need for a technical solution that can automatically identify the local block computation semantics of thread blocks from the structured Single Instruction, Multiple Threads (SIMT) intermediate representation and convert them into code that can be further optimized and generated for the target platform by the block-level backend. Summary of the Invention

[0004] The purpose of this invention is to address the shortcomings of existing technologies by providing a method and apparatus for automatic tensor kernel enhancement and cross-backend generation.

[0005] To achieve the above objectives, the present invention adopts the following technical solution: an automatic tensor kernel promotion and cross-backend generation method, comprising the following steps:

[0006] Obtain the target CUDA tensor kernel and perform a structured transformation to obtain a structured SIMT intermediate representation; wherein, the structured SIMT intermediate representation includes a structured parallel region representing the execution domain of a thread block, thread index, buffer with storage space identifier and its memory access, synchronization operation, scalar operation, and a structured control flow containing control flow predicates and loops;

[0007] A block context is constructed based on the thread block execution domain represented by the structured parallel region, the thread index is mapped to block coordinates covering the thread block execution domain, and scalar expressions that depend on the thread index are promoted point by point to block values ​​in the block context.

[0008] A predicate context is constructed based on the control flow predicates in the structured control flow. In the block context, thread-related branch conditions are promoted to block predicates, and the predicate context is updated based on the block predicates.

[0009] The loops in the structured control flow are subjected to block lifting to obtain loop lifting results. For unified loops, the shared loop structure is maintained, and for thread-related loops, a shared driving loop and activity mask covering each block element are generated.

[0010] Based on the block context, predicate context, and loop promotion result, the scalar address expression is promoted to a block pointer, the scalar memory access is promoted to a block read / write, and an initial block intermediate representation is generated.

[0011] Perform block-level optimization on the block intermediate representation to obtain the optimized block intermediate representation;

[0012] The optimized block intermediate representation is output to at least one block-level backend, which generates the target platform code.

[0013] Furthermore, the structured transformation includes: converting CUDA meshes and thread blocks into structured parallel regions, the structured parallel regions carrying thread block execution domains determined by the dimensions of each thread block; converting shared memory into buffers with shared memory space identifiers; converting conditional branches and loops into structured control flow; and preserving the order constraints of synchronization operations in the structured SIMT intermediate representation.

[0014] Furthermore, the block context is derived from the thread block execution domain. In this context, the thread index no longer represents a scalar value of a single thread instance, but rather a coordinate block covering the entire thread block execution domain. Arithmetic expressions, comparison expressions, and address expressions involving the thread index are all evaluated point-by-point in this block context.

[0015] Furthermore, the predicate context includes Boolean blocks with the same shape as the current block context, used to represent the set of elements in the current block that are active; when thread-related conditions are encountered, the block predicate is calculated under the block context, and the current predicate context is updated to the conjunction of the current predicate and the block predicate or the inverse of the block predicate;

[0016] Specifically, for the thread-related branch with condition p, the block predicate P is calculated point by point in the block context; the predicate context of the then branch is updated to AND-OR P, and the predicate context of the else branch is updated to XOR P.

[0017] Furthermore, for unified loops whose loop boundaries do not depend on thread indices, their shared loop structure is maintained, and block promotion continues within the loop body. For thread-dependent loops whose loop boundaries depend on thread indices, a scalar loop upper bound covering each block element is calculated, and a shared driving loop is generated. In each iteration of the shared driving loop, the block value corresponding to the source loop inductive variable is reconstructed, and an activity mask is generated based on the source loop boundary. For block elements that have ended their source loop, their previous loop-carried values ​​are retained through element-by-element selection.

[0018] Furthermore, based on the block context, predicate context, and loop promotion result, scalar address expressions are promoted to block pointers, scalar memory accesses are promoted to block reads and writes, and an initial block intermediate representation is generated, including: for block operations in thread-related branches, using the current predicate context as a mask; for block operations in thread-related loops, using the conjunction result of the current predicate context and the corresponding active mask as a mask; and attaching the mask to block reads and writes, block atomic operations, or other block operations with side effects.

[0019] Furthermore, the initial block intermediate representation uses block value, block pointer, and block predicate as basic objects;

[0020] Here, the block value represents the element value calculated point by point in the block field; the block pointer represents a set of memory locations determined by the base address and the block offset; and the block predicate represents a Boolean block obtained by evaluating thread-related conditions point by point.

[0021] The initial block intermediate representation includes coordinate generation, dimension expansion, broadcasting, pointer construction, block loading, block storage, block atomic operations, and structured loop operations. For side effect operations in thread-related branches, the current predicate context is used as a mask. For side effect operations in thread-related loops, the conjunction result of the current predicate context and the corresponding active mask is used as a mask so that inactive block elements do not produce corresponding memory effects.

[0022] Furthermore, the block-level optimization includes one or more of the following: access image analysis, memory-to-register promotion, reduction normalization, matrix multiplication normalization, superword-level parallelism (SLP) packing, and layout allocation.

[0023] Furthermore, the access image analysis includes: for each block read / write record base object, block shape, and affine mapping from block coordinates and signed operands to linear storage offsets, and determining the necessary aliasing, no aliasing, or compatible rearrangement relationship between block reads and writes based on the access image;

[0024] The memory-to-register promotion includes: establishing symbolic storage states for block storage and block loading that are proven to be the same logical access set by the access image; if the elements read by subsequent loading are all covered by the recorded storage, then the subsequent loading is replaced with the corresponding block SSA value; if there are overlapping writes that cannot be proven to be equivalent, then the original memory operation is retained.

[0025] The reduction normalization includes: identifying loops consisting of loops carrying accumulators, reduction identities, and associative merging operations, and rewriting them as block reduction operations; the matrix multiplication normalization includes identifying loops that update accumulators by multiplying or fused multiplying, and rewriting them as block matrix multiplication operations when operands access the image to form compatible M×K and K×N views.

[0026] The SLP packaging includes: identifying multiple chunk computations that access adjacent elements and are isomorphic to the computation tree, merging them into a single chunk value with a newly added secondary dimension, and maintaining consistency in the corresponding read / write mask and memory attributes;

[0027] The layout allocation includes: establishing a candidate layout set for block values, operand usage points, and memory accesses; ensuring the consistency of representation among block pointers, block values, and block predicates through a constraint model; and selecting a target layout based on continuous access, unit step access, and layout conversion overhead.

[0028] Furthermore, the block-level backend includes a Triton backend, a TileLang backend, or other backends that support block-level tensor operations, and the target platform code includes code for GPUs or other parallel accelerators.

[0029] To achieve the above objectives, the present invention also provides a tensor kernel block-level program automatic promotion and cross-backend generation apparatus, including one or more processors for implementing the above method.

[0030] To achieve the above objectives, the present invention also provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the above-described method.

[0031] The beneficial effects of this invention are: it can recover the local block-level computational semantics of thread blocks from existing CUDA tensor kernels, reducing the understanding and rewriting work required for manual migration to a block-level programming model; by unifying thread indexes, control predicates, loops, and memory access into block representations, it can provide a higher-level program structure for backend optimization while maintaining the original computational semantics; and through a multi-backend generation mechanism, it can improve the portability of existing GPU kernels on different block-level backends and different target platforms. Attached Figure Description

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

[0033] Figure 1 This is a flowchart illustrating the overall process framework of the method of this invention;

[0034] Figure 2 This is a schematic diagram of the block context and predicate mask of the present invention;

[0035] Figure 3 This is a schematic diagram of the structure of the device of the present invention. Detailed Implementation

[0036] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of the present invention, and not all of them. Other embodiments obtained by those skilled in the art based on the embodiments of the present invention without creative effort are all within the protection scope of the present invention.

[0037] like Figure 1 As shown, this embodiment of the invention provides a method for automatic tensor kernel promotion and cross-backend generation, including the following steps:

[0038] Step 1: Input the target CUDA tensor kernel, identify its thread block structure, thread index, memory space, synchronization operations, scalar operations and memory access, and convert it into a structured SIMT intermediate representation.

[0039] The structured SIMT intermediate representation includes a structured parallel region representing the execution domain of a thread block, a thread index, a buffer with a storage space identifier and its memory access, synchronization operations, scalar operations, and a structured control flow containing control flow predicates and loops.

[0040] This intermediate representation uses structured parallel regions to represent CUDA grids and thread blocks, explicit buffers to represent shared memory, and synchronization operations to represent inter-thread order constraints.

[0041] Step 2: Construct a block context based on the thread block execution domain characterized by the structured parallel region.

[0042] In this application, the execution domain of a thread block refers to the set of thread index values ​​corresponding to each thread instance within a CUDA thread block. For a thread block with a d-dimensional thread structure, the execution domain of the thread block is the Cartesian product of the ranges of thread index values ​​in each dimension: D = [0, S0) × ... × [0, Sd-1). Each thread index corresponds to a coordinate block. Scalar expressions in the source program that depend on thread indices are promoted point-by-point to block values ​​in this context; scalar values ​​that do not depend on thread indices are broadcast to the block shape when needed.

[0043] Step 3: Construct the predicate context based on the control flow predicates in the structured control flow.

[0044] For thread-related conditions, the block predicate is calculated in the block context, and the current active predicate is combined with the block predicate.

[0045] For memory read / write or atomic operations with side effects, instead of generating new low-level branches, a mask is added to the block read / write operations to maintain the semantics that only some threads in the source program execute the operation.

[0046] Step 4: Perform block lifting on the loops in the structured control flow to obtain the loop lifting results.

[0047] For a unified loop whose loop boundary is independent of thread coordinates, retain the loop and continue to perform block promotion within the loop body; for a loop whose loop boundary depends on thread coordinates, calculate the shared upper bound of the driving loop covering all block elements, and reconstruct the block value of the source loop inductive variable and the corresponding activity mask in each driving iteration.

[0048] Step 5: Generate the initial block intermediate representation based on the block context, predicate context and loop promotion result.

[0049] The initial block intermediate representation includes three basic objects: block value, block pointer, and block predicate, and includes operations such as coordinate generation, dimension expansion, broadcasting, pointer construction, block loading, block storage, block atomic operations, and structured loops.

[0050] In this context, the initial block represents the tensor block computation of the thread block without explicitly preserving the specific thread scheduling of the source CUDA program.

[0051] Step 6: Perform block-level optimization on the initial block intermediate representation to obtain an optimized block intermediate representation.

[0052] Wherein, the block-level optimization comprises access pattern analysis, memory-to-register promotion, reduction normalization, matrix multiplication normalization, SLP packing, and layout allocation.

[0053] Through the above optimizations, the block computation implicitly expressed by shared memory and scalar loops is restored to high-level block operations that can be directly recognized by the back end.

[0054] Step 7: Output the optimized block intermediate representation to the block-level back end.

[0055] In this embodiment, the block-level back end may be Triton, TileLang, or other compilation back ends that support block-level tensor operations. Different back ends can select corresponding code generation strategies according to the target platform to generate target code on a GPU or other parallel accelerators.

[0056] As a preferred embodiment, the block context in step 2 is as Figure 2 shown. Each thread instance in the thread block corresponds to a position in the block domain, and the thread indices tx and ty are interpreted as two coordinate blocks T0 and T1 respectively.

[0057] If the address expression in the source program is base+tx or row×stride+ty, a corresponding block pointer is obtained after promotion; if the operation expression in the source program contains both a thread-related value and a uniform scalar value, the uniform scalar value is broadcast and then participates in element-wise operations.

[0058] As a preferred embodiment, the predicate context in step 3 is also as Figure 2 shown. When a condition such as if(tx<n) exists in the source program, a boolean block consistent with the block shape is generated. Block elements that satisfy the condition are active elements, and elements that do not satisfy the condition are inactive elements. Block loading, block storing, and loop-carried value updating are all performed under this active mask.

[0059] As a preferred embodiment, the thread-related loop in step 4 may be a thread-strided loop. For a loop in the source program that uses a thread index as an initial value and a thread block size as a stride, the thread index is promoted to a coordinate block, and a shared driving loop is used to represent the iteration set of all threads.

[0060] In each driving iteration, the current block induction variable and the active mask are calculated, so that block elements that have not yet reached the upper bound of the source loop continue to execute the loop body, and block elements that have finished maintain their previous states.

[0061] As a preferred embodiment, the access image analysis in step 6 includes: recording the base object, block shape, and affine address mapping for each block loading or block storage.

[0062] When two block read / write address mappings correspond to the same set of logical elements, or to views that can be proven to be compatible such as transpose, projection, or broadcast, block values ​​can be forwarded or temporary memory states can be promoted while preserving semantics. If access equivalence or non-overlap cannot be proven, the original memory operations are conservatively preserved.

[0063] In a preferred embodiment, the reduction normalization includes: detecting whether the loop-carrying accumulator is updated by associative operations such as addition, multiplication, maximum or minimum value, and determining whether the initial value is the corresponding identity. When the condition is met, the loop is rewritten as an explicit block reduction operation.

[0064] The matrix multiplication normalization includes: detecting the multiply-accumulate pattern and determining whether an M×K and K×N matrix view is formed by combining the operand access image; when the condition is met, the accumulation loop is rewritten as an explicit block matrix multiplication operation.

[0065] In a preferred embodiment, the SLP packaging is used to handle the scenario where a single thread computes multiple adjacent logical elements. This includes: identifying multiple memory-access-adjacent, mask-consistent, and computation-tree-isomorphic block computations, and merging them into a block value with a newly added secondary dimension. This packaging result is then passed to the layout allocation stage to choose whether to flatten, rearrange, or maintain the original state, so that the backend can generate more suitable access and computation code.

[0066] As a preferred embodiment, the layout allocation includes: establishing a finite set of candidate layouts for each block value, operand usage point, and memory access, and establishing layout consistency constraints. For example, pointers, values, and masks for block storage operations need to use compatible layouts. A cost function is used to favor sequential or unit-step accesses and penalize additional rearrangements, broadcasts, and flattening transformations, thereby selecting a layout combination with lower overall cost.

[0067] As an experimental embodiment, the present invention performs the following specific experimental verification based on the above method: First, 145 CUDA tensor kernel implementations are collected as candidate inputs, and the implementations cover types such as element-wise computation, reduction, layout transformation, convolution, attention, quantization and matrix multiplication.

[0068] The optimized block intermediate representation is output to the Triton backend and the TileLang backend respectively, and runs on NVIDIA A100, NVIDIA H100 and AMD MI300X platforms.

[0069] Performance measurement involved 10 warm-ups, 50 kernel iterations per sample, and was repeated 3 times overall, with the median taken. Compilation time was not included in kernel execution time.

[0070] To evaluate the achievable performance upper bounds of the two backends, Triton and TileLang versions were compiled separately for each verified implementation, and the effective results with shorter execution times were selected offline.

[0071] Compared to the source CUDA program, the geometric mean speedup on the A100 and H100 platforms is 1.41x and 1.47x, respectively; compared to the source program converted by HIPIFY, the geometric mean speedup on the MI300X platform is 1.87x.

[0072] Corresponding to the foregoing method embodiments, the present invention also provides an embodiment of a tensor kernel-based block-level program automatic promotion and cross-backend generation apparatus. For example... Figure 3 As shown, the device includes an input acquisition module, a structured transformation module, a context construction module, a semantic enhancement module, an access analysis module, a normalization optimization module, a layout and packaging module, and a backend generation module. These modules can be implemented in software, hardware, or a combination of both.

[0073] Specifically, the input acquisition module is used to receive the target CUDA tensor kernel and its compilation parameters;

[0074] The structured transformation module is used to generate a structured SIMT intermediate representation;

[0075] The context building module is used to build block contexts and predicate contexts;

[0076] The semantic enhancement module is used to generate block intermediate representations;

[0077] The access analysis module is used to construct an access image and determine read-write relationships;

[0078] The normalization optimization module is used to perform memory-to-register promotion, reduction normalization, and matrix multiplication normalization;

[0079] The layout and packing module is used to perform SLP packing and layout allocation;

[0080] The backend generation module is used to output to the block-level backend and generate target code.

[0081] Embodiments of the device of the present invention can be applied to any device with data processing capabilities, such as a computer, server, compiler server, development workstation, or embedded development device. Taking software implementation as an example, as a logical device, it is formed by the processor of the device loading the corresponding computer program instructions from non-volatile memory into memory for execution.

[0082] This invention also provides a computer-readable storage medium storing a program thereon, which, when executed by a processor, implements the methods described in the above embodiments. The computer-readable storage medium may be a hard disk, a solid-state drive, a smart memory card, an SD card, a flash memory card, a read-only memory, a random access memory, or other media capable of storing program instructions.

[0083] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention. Any modifications, equivalent substitutions, or improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

[0084] The above embodiments are only used to illustrate the design concept and features of the present invention, and are intended to enable those skilled in the art to understand the content of the present invention and implement it accordingly. The scope of protection of the present invention is not limited to the above embodiments. Therefore, all equivalent changes or modifications made based on the principles and design ideas disclosed in the present invention are within the scope of protection of the present invention.

Claims

1. A method for automatic promotion and cross-backend generation of tensor kernels, characterized in that, include: Obtain the target CUDA tensor kernel and perform a structured transformation to obtain a structured SIMT intermediate representation; wherein, the structured SIMT intermediate representation includes a structured parallel region representing the execution domain of a thread block, thread index, buffer with storage space identifier and its memory access, synchronization operation, scalar operation, and a structured control flow containing control flow predicates and loops; A block context is constructed based on the thread block execution domain, the thread index is mapped to block coordinates covering the thread block execution domain, and scalar expressions that depend on the thread index are promoted point by point to block values ​​in the block context. A predicate context is constructed based on the control flow predicate, and thread-related branch conditions are promoted to block predicates in the block context, and the predicate context is updated based on the block predicates; Block promotion is performed on the loop to obtain the loop promotion result. For the unified loop, the shared loop structure is maintained, and for the thread-related loop, a shared driving loop and activity mask covering each block element are generated. Based on the block context, predicate context, and loop promotion result, the scalar address expression is promoted to a block pointer, the scalar memory access is promoted to a block read / write, and an initial block intermediate representation is generated. Perform block-level optimization on the initial block intermediate representation to obtain the optimized block intermediate representation; The optimized block intermediate representation is output to at least one block-level backend, which generates the target platform code.

2. The method according to claim 1, characterized in that, The structured transformation includes: converting CUDA meshes and thread blocks into structured parallel regions carrying thread block execution domains determined by the dimensions of each thread block; converting shared memory into buffers with storage space identifiers; converting conditional branches and loops into structured control flow; and preserving the order constraints of synchronization operations in the structured SIMT intermediate representation.

3. The method according to claim 1, characterized in that, The predicate context includes Boolean blocks with the same shape as the current block context; when thread-related conditions are encountered, the block predicate is calculated under the block context, and the current predicate context is updated to the conjunction of the current predicate and the block predicate or the inverse of the block predicate.

4. The method according to claim 1, characterized in that, The process of generating a shared driving loop and an activity mask covering each block element for thread-related loops includes: for thread-related loops whose loop boundaries depend on thread indices, calculating the scalar loop upper bound covering each block element and generating a shared driving loop; in each iteration of the shared driving loop, reconstructing the block value corresponding to the source loop inductive variable and generating an activity mask based on the source loop boundary; and for block elements that have ended their source loop, retaining their previous loop carry value by selecting each element one by one.

5. The method according to claim 1, characterized in that, Based on the block context, predicate context, and loop promotion result, scalar address expressions are promoted to block pointers, scalar memory accesses are promoted to block reads and writes, and an initial block intermediate representation is generated, including: for block operations in thread-related branches, using the current predicate context as a mask; for block operations in thread-related loops, using the conjunction result of the current predicate context and the corresponding active mask as a mask; and attaching the mask to block reads and writes, block atomic operations, or other block operations with side effects.

6. The method according to claim 1, characterized in that, The block-level optimization includes one or more of the following: access image analysis, memory-to-register promotion, reduction normalization, matrix multiplication normalization, SLP packing, and layout allocation.

7. The method according to claim 6, characterized in that, The access image analysis includes: for each block read / write record base object, block shape, and affine mapping from block coordinates and signed operands to linear storage offsets, and determining the necessary aliasing, no aliasing, or compatible rearrangement relationship between block reads and writes based on the access image; The memory-to-register promotion includes: establishing symbolic storage states for block storage and block loading that are proven to be the same logical access set by the access image; if the elements read by subsequent loading are all covered by the recorded storage, then the subsequent loading is replaced with the corresponding block SSA value; if there are overlapping writes that cannot be proven to be equivalent, then the original memory operation is retained. The reduction normalization includes: identifying loops consisting of loops carrying accumulators, reduction identities, and associative merging operations, and rewriting them as block reduction operations; the matrix multiplication normalization includes identifying loops that update accumulators by multiplying or fused multiplying, and rewriting them as block matrix multiplication operations when operands access the image to form compatible M×K and K×N views. The SLP packaging includes: identifying multiple chunk computations that access adjacent elements and are isomorphic to the computation tree, merging them into a single chunk value with a newly added secondary dimension, and maintaining consistency in the corresponding read / write mask and memory attributes; The layout allocation includes: establishing a candidate layout set for block values, operand usage points, and memory accesses; ensuring the consistency of representation among block pointers, block values, and block predicates through a constraint model; and selecting a target layout based on continuous access, unit step access, and layout conversion overhead.

8. The method according to claim 1, characterized in that, The block-level backend includes the Triton backend, the TileLang backend, or other backends that support block-level tensor operations, and the target platform code includes code for GPUs or other parallel accelerators.

9. A device for automatic promotion and cross-backend generation of block-level programs using tensor kernels, characterized in that, It includes one or more processors for implementing the method according to any one of claims 1-8.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the method described in any one of claims 1-8.

Images