A TensorIR-based TPU automatic tensorization method

Abstract

The application belongs to the technical field of compilation optimization of artificial intelligence accelerators, and proposes a TPU automatic tensorization method based on TensorIR, which can automatically optimize and generate operation kernel code of a TPU tensor accelerator. The method analyzes an abstract syntax tree of the scheduled TensorIR through traversal, so as to automatically replace instructions of each hardware behavior. The application discards the cumbersome tensorization process of handwritten IR matching fragments in the existing TensorIR, more flexibly realizes tensorization of each neural network operator on a TPU, and can improve kernel operation performance by 2-3 times.

Description

Technical Field

[0001] This invention belongs to the field of compilation optimization technology for artificial intelligence accelerators, and particularly relates to an automatic tensor quantization method for TPU based on TensorIR, which can be used for automatic optimization and generation of kernel code on TPU tensor accelerators. Background Technology

[0002] To address the surge in tensor computation demands in the age of artificial intelligence, various hardware accelerators, such as TPUs (Tensor Processing Units), have been specifically developed for accelerating high-dimensional tensor computations. Their instruction set architecture supports a variety of neural network-level operators. On the other hand, to facilitate the transition of AI models from the design and development phase to the hardware deployment phase, neural network compilers, represented by TVMs (Tensor Virtual Machines), employ compilation optimization techniques similar to those in traditional compilers: utilizing intermediate representations (IRs) at layer and operator levels to progressively compile the neural network model into executable code on the TPU. Tensorization, in particular, transforms software-defined tensor operations into high-performance kernel code (hardware-side computational kernel functions) by calling hardware-specific instructions. This allows for flexible manipulation of various instruction behaviors of the hardware accelerator, and is one of the keys to enabling neural network compilers to adapt to various hardware accelerators and achieve high computational performance.

[0003] TVM is currently the most mature neural network compiler, which designs a TensorIR, an intermediate representation of the operator layer specifically for tensor quantization of hardware instructions. TensorIR decouples nested loops outside tensor operations from the internal computational body by utilizing an IR abstraction structure centered on blocks, thus facilitating the replacement of hardware instructions by the program IR to achieve tensor quantization. However, TensorIR primarily relies on hand-written IR fragments for matching and replacing hardware instructions. This process requires manually writing corresponding Block IR code fragments for various hardware behavioral instructions, such as initialization instructions, data read instructions, operation instructions, and data write-back instructions, and then matching these fragments with the original program IR to replace various behavioral blocks with instruction calls. Furthermore, it still requires manually writing a set of IR fragment combinations for each of the above instruction behaviors for the multiple neural network operators supported by the hardware accelerator. Therefore, this manual tensor quantization process is extremely cumbersome and difficult to scale in TPU scenarios.

[0004] Related patents, such as "A TVM-based NPU neural network model deployment method and device" (CN202210482702.2), propose a TVM-based model compilation optimization process for NPU-like hardware accelerators. However, NPUs do not provide neural network-level operator instructions (such as 2D convolution), and this method requires mapping other operators to matrix operation units on the NPU, thus it is not suitable for hardware like TPUs that support multiple neural network-level operators. "Extension method, device, apparatus, and readable storage medium for compiler hardware backend" (CN202310138812.1) proposes a general hardware backend extension method adapted to various artificial intelligence processors. However, this method mainly relies on the pre-set operator library on the hardware side for compiler-hardware operator adaptation and does not involve fine-grained support at the hardware instruction level, thus making it difficult to flexibly generate high-performance kernel operator code.

[0005] In summary, existing work rarely provides hardware tensor quantization methods for TPUs that support multiple neural network-level operators, while the existing manual tensor quantization process of TVM lacks flexibility and is not easily extended to multiple operators when facing TPU scenarios. Summary of the Invention

[0006] The purpose of this invention is to provide an automatic tensor quantization method for TPU tensor accelerators, based on and improved upon TensorIR. This method primarily involves automatically traversing the abstract syntax tree of the scheduled TensorIR and utilizing the decoupling characteristics of blocks to perform instruction replacement for various hardware behaviors. This invention eliminates the tedious process of manually writing IR fragment matching, thus enabling more flexible tensor quantization of various neural network operators on the TPU.

[0007] The technical solution of the present invention:

[0008] An automatic tensor quantization method for TPU based on TensorIR, the specific steps of which are as follows:

[0009] Step 1: Describe the operator logic. TE (Tensor Expression) is an abstract language in the TVM compiler specifically used to describe the tensor operation logic of the software layer. The TE module is used to describe the initial tensor operation expression of a certain neural network operator: declare the dimension information and shape size of each tensor data, use lambda expressions to describe the operation logic from input tensor to output tensor, and generate the initial TensorIR program.

[0010] Step 2: Scheduling TensorIR. Scheduling is the program scheduling module in TVM applied to the low-level operator layer IR. It generates the corresponding schedule for the initial TensorIR program in Step 1 through the TVM's program scheduling interface, and calls various scheduling primitive functions to perform program transformations on the TensorIR. The main steps are as follows:

[0011] First, the read cache and write cache scheduling primitives are invoked to create Blocks corresponding to data read and write operations. Second, the loop splitting and loop rearrangement primitives are invoked to perform two-level loop splitting and rearrangement of nested loops to adapt to the two-level storage structure on the TPU. Then, the specified computation primitive is invoked to specify the loop dimension for the read and write Blocks. Finally, the dereduction primitive is invoked to decompose the initial computation Block into computation initialization Blocks and computation update Blocks. After applying a series of scheduling primitives, TensorIR is transformed into an abstract form that more closely resembles the computation kernel on the TPU.

[0012] Step 3: Automatic TensorQuantization. TensorIR is an intermediate representation designed based on Python syntax. Using Python's AST (Abstract Syntax Tree) module, the TensorIR scheduled in Step 2 is traversed through its abstract syntax tree, and instructions for each line of blocks are automatically replaced to generate the tensorquantization program. The specific process is as follows:

[0013] (a): Use Python AST parsing to traverse the abstract syntax tree of the scheduled TensorIR and obtain the AST nodes corresponding to each behavior Block (initialization, data reading, data writing, data computation).

[0014] (b): Traverse the nested loop nodes outside the Block, collect the iteration variables of each node in the corresponding dimension and store them as an iteration variable pool.

[0015] (c): Iterate through and extract the statements that bind the loop iteration variables inside and outside each Block, remove the independent variables that do not access the external variable pool in the binding statements corresponding to each dimension within the Block, and retain the remaining binding sub-statements.

[0016] (d): Finally, replace each Block node in the abstract syntax tree with a virtual instruction call node. Note that the instruction call nodes replaced in this step are virtual function call nodes, which are still abstract representations of various hardware behaviors. They can be replaced with actual instruction calls according to the instruction requirements of different neural network operators.

[0017] Step 4: Preset Instruction Format. Different neural network operators on the TPU require different computation and data transfer instructions. It is necessary to preset the required TPU-side data initialization, data transfer, and data computation instruction formats for the neural network operators, including the names of the instruction functions and the specific format of the function parameters.

[0018] Step 5: Kernel Code Generation. A corresponding C code generation module is added to the TPU. During code generation by traversing the ASTs of the TensorIR, the virtual instructions in the quantization program IR from Step 3 are converted one-to-one into real function calls according to the instruction function call format from Step 4. This includes replacing function names, adjusting function parameter formats, etc., and fetching data from the TPU memory based on the binding sub-statements retained in Step 3; finally, the C code kernel program on the TPU is generated.

[0019] The beneficial effects of this invention are as follows: The method of this invention can automatically replace hardware instructions and generate kernel code for tensor operation programs defined in the software layer, reducing the need for additional steps of manually writing IR matching fragments, making the tensor quantization process simpler and more easily scalable. This invention can support the automatic optimization and generation of operation kernels for neural network-level operators such as matrix multiplication, 2D convolution, dilated convolution, and depthwise separable convolution on TPU, and can improve kernel operation performance by 2 to 3 times. Attached Figure Description

[0020] Figure 1 This is a flowchart illustrating the TPU automatic tensor quantization method based on TensorIR of this invention. Detailed Implementation

[0021] The method of the present invention will be described in detail below with reference to the accompanying drawings, technical solutions, and embodiments.

[0022] This invention uses a TPU accelerator as an experimental example for illustration. The corresponding offline compilation platform is deployed on a Linux server, and the specific server configuration is shown in Table 1. According to the experimental requirements, TVM version 0.11.1, Python version 3.9.7, and SDK version 2.7.0 are installed on the server.

[0023] Table 1. Linux Server Configuration Information

[0024]

[0025] like Figure 1As shown, the automatic tensor quantization method for TPU based on TensorIR of the present invention proceeds as follows: First, the tensor operation program of the software layer is defined using the TE module of TVM. Then, the initial TensorIR is scheduled and transformed. Next, the TensorIR abstract syntax tree is traversed and instructions are automatically replaced for each line of Block. Finally, kernel code executable on the TPU is generated according to the specific instruction call requirements of the operator. Specific implementation details are as follows:

[0026] Step 1: Describe the operator logic. Use the TE (Tensor Expression) module of the TVM compiler to describe the basic tensor operation expression of a neural network operator.

[0027] Specifically, it utilizes TE to declare the shape, size, and reduction dimension information of each input and output tensor in tensor operations, as well as to describe the operational logic mapping from the input tensor to the output tensor. The following formula, using the lambda expression of matrix multiplication as an example, describes the operational mapping formula from input matrices A[m,k] and B[k,n] to output matrix C[m,n], and specifies that the reduction dimension k is the accumulation dimension.

[0028] C=te.compute(lambdam,n:te.sum(A[m,k]*B[k,n]),axis=k)

[0029] TE will automatically generate a tensor operation program based on the operation map. The operation program described in this step is the initial TensorIR program without any scheduling transformations.

[0030] Step 2: Schedule TensorIR. Generate a corresponding schedule for the initial TensorIR program from Step 1 using the TVM program scheduling interface. Obtain the initial operation block and various dimensions in the tensor operation from the schedule, and call various scheduling primitives to perform program transformations on the TensorIR.

[0031] Specifically: First, the read cache and write cache scheduling primitives are called to create a read block for the input tensor and a write block for the output tensor of the initial operation block. Second, the loop splitting primitive is called to split each loop dimension in the tensor operation. The original dimension axis[i] is split into an outer dimension axis_outer[i] and an inner dimension axis_inner[i], thus adapting to the two-level storage structure of global memory and local memory on the TPU. Next, the loop rearrangement primitive is called to rearrange the loop order after splitting, placing all outer dimensions in the outer layer and all inner dimensions in the inner layer. For example, the nested loop order of the three dimensions m, n, and k in matrix multiplication after splitting and rearranging is [m_o, n_o, k_o, m_i, n_i, k_i]. Then, the specified computation primitive is called to specify the loop dimension for the read block and the write block. The read block of each input tensor is specified to be computed in the outer reduction dimension, and the write block of the output tensor is specified to be computed in the outer spatial dimension. Finally, the dereduction primitives are invoked to decompose the initial computation block into a computation initialization block and a computation update block. After applying a series of scheduling primitives, TensorIR is transformed into an abstraction that more closely resembles the computation kernel on the TPU.

[0032] Step 3: Automatic Tensorification. Using the Python AST module, the TensorIR scheduled in Step 2 is traversed, its abstract syntax tree is automatically replaced with instructions for each line of blocks, generating the tensorification program. The specific algorithm flow is as follows:

[0033] Process a: Use the Python AST module to parse and traverse the abstract syntax tree of the scheduled TensorIR to obtain the AST nodes corresponding to each Block (initialization, data read, data write, data computation).

[0034] Process b: Traverse the nested for loop nodes outside the Block, collect the iteration variables of the corresponding dimension of each for node and store them as an iteration variable pool, such as the outer iteration variable pool of matrix multiplication is [m_0,k_0,n_0].

[0035] Process c: Iterate through and extract the statements that bind loop iteration variables within and outside each Block. Remove independent variables from the binding statements for each dimension within the Block that do not access the external variable pool, and retain the remaining binding sub-statements. The following formula uses the Block for reading matrix A[m,k] in matrix multiplication as an example. The left side of the arrow represents the original binding statements within the Block, and the right side represents the binding sub-statements that need to be retained. Here, ax0 and ax1 are independent variables within the Block and need to be removed; sub-statements m_0*256 and k_0*256 access the external iteration variable pool and need to be retained.

[0036]

[0037] Process d: Finally, the NodeTransformer class of Python AST is used to replace each Block node in the abstract syntax tree with a virtual instruction call node. The initialization, read, compute, and write Block nodes will be replaced with function call nodes such as Init, Load1, Load2, Compute, and Store, respectively. Note that the instruction call nodes replaced in this step are virtual function call nodes, which are still abstract representations of various hardware behaviors. They can be replaced with actual instruction calls according to the instruction requirements of different neural network operators.

[0038] Step 4: Preset Instruction Format. Preset the call parameter format for the instruction functions required by the neural network operator for TPU-side data initialization, data transfer, data operation, etc.

[0039] Specifically, this involves: pre-setting the TPU-side instruction set required by the neural network operator, such as data initialization instructions, data transfer instructions, and data computation instructions. It declares the function name of each instruction and the specific parameter format during actual invocation. Table 2 lists the names and parameters of each required instruction function, using matrix multiplication as an example. For instance, the matrix transfer instruction needs to specify the source address in global memory, the target address in local memory, and the row and column information of the matrix. The actual instructions of each operator are then mapped and stored in dictionary form with the virtual instructions from step 3: {IR_func:[TPU_func,params]}.

[0040] Table 2. TPU Matrix Multiplication Operator Instruction Set

[0041]

[0042]

[0043] Step 5: Kernel Code Generation. Add a C code generation module and, based on the instruction function call format in Step 4, convert the virtual instructions in the quantization program IR from Step 3 into real function calls one-to-one, thereby generating the C code kernel program on the TPU.

[0044] Specifically, a corresponding C code generation module is added to the TPU. When generating code by traversing the AST of TensorIR, the virtual instructions in the tensor quantization program IR in step 3 are converted into real function calls one-to-one according to the instruction function call format in step 4. This includes replacing function names, passing specific function parameters, etc., and performing data addressing on the TPU's global memory based on the dimension binding sub-statements retained in each Block in step 3. The addressing information of the tensor data in one iteration is equal to its initial base address plus the offset calculated by the following formula. Where [dim(1), dim(2), ... dim(n)] represents n-dimensional tensor data, O represents the addressing offset, S represents the stride of a certain dimension in the tensor data, and I represents the binding sub-statement of a certain dimension retained in step 3. The addressing offset O of the tensor data in one iteration is equal to the sum of the products of S and I of each dimension, and the specific calculation formula is as follows:

[0045]

[0046] The memory addressing information will be used as the parameter information for global memory data access in each instruction, and will eventually generate the C code kernel program on the TPU.

Claims

1. A TPU automatic tensor quantization method based on TensorIR, characterized in that, The specific steps are as follows: Step 1: Describe the operator logic; Use the TE module to describe the initial tensor operation expression of a neural network operator: declare the dimension information and shape size of each tensor data, and use lambda expressions to describe the operation logic from input tensor to output tensor, generating the initial TensorIR program; Step 2: Schedule TensorIR; The TVM program scheduling interface generates the corresponding schedule for the initial TensorIR program from Step 1, and various scheduling primitive functions are called to perform program transformations on the TensorIR, specifically as follows: First, the read cache and write cache scheduling primitives are called to create blocks corresponding to data read and write operations; second, the loop splitting and loop rearrangement primitives are called to perform two-level loop splitting and rearrangement of nested loops to adapt to the two-level storage structure on the TPU; then, the specified computation primitive is called to specify the loop dimension for the read and write blocks; finally, the dereduction primitive is called to decompose the initial operation block into an operation initialization block and an operation update block; Step 3: Automatic Tensorification; Using Python's AST module, the TensorIR scheduled in Step 2 is traversed, and its abstract syntax tree is automatically replaced with instructions for each line of Blocks to generate the tensorification program; the specific process is as follows: (a): Use Python AST to parse and traverse the abstract syntax tree of the scheduled TensorIR to obtain the AST nodes corresponding to each behavior Block; (b): Traverse the nested loop nodes outside the Block, collect the iteration variables of each node in the corresponding dimension and store them in the iteration variable pool; (c): Iterate through and extract the statements that bind the loop iteration variables inside and outside each Block, remove the independent variables that do not access the external variable pool in the binding statements corresponding to each dimension within the Block, and retain the remaining binding sub-statements. (d): Finally, replace each Block node in the abstract syntax tree with a virtual instruction call node; Step 4: Preset instruction format; Preset the TPU instruction format required for data initialization, data transfer, and data operation for the neural network operator, including the name of the instruction function and the specific format of function parameter passing; Step 5: Kernel code generation; Add a corresponding C code generation module to the TPU. When generating code by traversing the AST of the TensorIR, convert the virtual instructions in the quantization program IR in Step 3 into real function calls one-to-one according to the instruction function call format in Step 4. This includes replacing function names and adjusting function parameter formats, and fetching data on the TPU memory according to the binding sub-statements retained in Step 3. Finally, generate the C code kernel program on the TPU.

Images