Calculation method of Conv2D operator on TPU chip

Abstract

The invention relates to the technical field of deep learning reasoning, in particular to a calculation method of a Conv2D operator on a TPU chip. The method comprises the following steps: decomposing a size convolution kernel to be processed into 1 * 1 sub-kernels; dynamically determining an optimal block size and a segmentation strategy based on the VMEM capacity, the convolution parameter and the input size on the TPU chip; on the basis of convolution parameter features, different calculation scenes are adapted through a hierarchical multi-branch optimization strategy, and a calculation adaptation path of the 1 * 1 sub-kernel is determined; according to the optimal block size and the adaptive path, matrix multiplication is executed, and a calculation result is obtained; and accumulating all calculation results to obtain the final output of the Conv2D operator. According to the method, the problems that memory redundancy expansion is caused by explicit im2col, and overflow is easily caused by dependence on a transfer space can be solved; serial execution of data transmission and calculation; and the scene adaptability of a fixed partitioning strategy is weak.

Description

Technical Field

[0001] This invention relates to the field of deep learning inference technology, and in particular to a method for calculating Conv2D operators on a TPU chip. Background Technology

[0002] Conv2D is a core computational operator of Convolutional Neural Networks (CNNs). It accounts for a very high proportion of computation in mainstream models such as Residual Networks (ResNet) and YOLO. Its implementation efficiency directly determines the end-to-end inference latency and the overall system performance, making it a key target for hardware platform adaptation and optimization.

[0003] Existing convolution implementations include explicit im2col+GEMM (matrix multiplication after data recombination), direct convolution (multi-layer loop computation), and the Winograd algorithm (reducing the number of multiplications). Solutions on graphics processing units (GPUs) or central processing units (CPUs) are difficult to directly port to tensor processing units (TPUs) due to differences in hardware architecture. Existing TPU solutions are not fully adapted to its dual-core, hierarchical memory, and dedicated instruction set characteristics.

[0004] Existing technologies have significant drawbacks: explicit im2col leads to memory redundancy and reliance on intermediate space, which can easily cause overflow; data transmission and computation are executed serially, and the hardware parallelism is not fully utilized; fixed block partitioning strategies cannot adapt to diverse convolution parameters and input sizes, resulting in weak scenario adaptability. Summary of the Invention

[0005] This invention provides a method for calculating Conv2D operators on a TPU chip to solve the problems in the prior art, such as memory redundancy expansion caused by explicit im2col, easy overflow due to reliance on intermediate space; serial execution of data transmission and computation, which does not fully utilize hardware parallelism; and the inability of fixed block strategy to adapt to diverse convolution parameters and input sizes, resulting in weak scene adaptability.

[0006] In a first aspect, embodiments of the present invention provide a method for calculating Conv2D operators on a TPU chip, including: To be processed Size convolution kernel decomposed into Each 1×1 sub-core; Indicates the height of the convolution kernel. Indicates the width of the convolution kernel; Based on the VMEM capacity, convolution parameters, and input size on the TPU chip, the optimal block size and segmentation strategy are dynamically determined to provide memory adaptation support for the computation of the 1×1 sub-core; Based on the features of convolution parameters, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios and determine the computing adaptation path of the 1×1 sub-kernel. According to the optimal block size and the adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through the 1×1 sub-core to obtain the calculation result. The final output of the Conv2D operator is obtained by summing all the calculation results.

[0007] One possible implementation also includes: By utilizing the hierarchical memory architecture and double buffering mechanism of the TPU chip, parallel execution of computation and data transmission is achieved during the matrix multiplication operation.

[0008] In one possible implementation, the decomposition process and matrix multiplication are performed through a four-layer loop structure: output channel segmentation, input channel segmentation, convolution kernel position traversal, and spatial dimension sliding processing. The output channel is segmented based on the hardware architecture of the TPU, the base number of the output channel segments is determined, and the output channels are segmented according to the base number of the output channel segments; each time, the weight data corresponding to the output channel segment is loaded into the VMEM to adapt to the VMEM capacity constraints and the channel alignment access requirements of the TPU. The input channel is segmented based on the hardware architecture of the TPU, the base number of the input channel segments is determined, and the input channels are segmented according to the base number of the input channel segments; each time, the original input data corresponding to the input channel segment is loaded into the VMEM to match the parallel processing capability of the computing units in the TPU; The convolution kernel position traversal is to traverse the... The spatial location of each 1×1 sub-kernel is used to transform the single convolution operation of the original large convolution kernel into batch computation of each 1×1 sub-kernel; The spatial dimension sliding process involves sliding the 1×1 sub-kernel in the height and width directions of the input feature map, and achieving continuous access to the original input data through offset calculation related to the position of the convolution kernel.

[0009] In one possible implementation, dynamically determining the optimal block size and segmentation strategy based on the VMEM capacity, convolution parameters, and input size on the TPU chip includes: Establish a memory requirement model, and based on the memory requirement model, calculate the total memory usage of the input sub-blocks, output sub-blocks, weight data, and bias data required to process the target output block; Based on the convolution parameters and input size, the output feature map size is calculated, and based on the output feature map size, a segmentation strategy is determined; Based on the segmentation strategy and the memory requirement model, the maximum feasible block size is determined by iterative calculation within the VMEM capacity constraint using a binary search algorithm. The maximum feasible block size and number of channels are adjusted for hardware alignment to adapt to the access requirements of the TPU.

[0010] In one possible implementation, determining the segmentation strategy based on the output feature map size includes: When the height of the output feature map is greater than or equal to its width, a row segmentation strategy is adopted; When the height of the output feature map is less than its width, a column splitting strategy is adopted.

[0011] In one possible implementation, based on the partitioning strategy and the memory requirement model, the maximum feasible block size is determined iteratively within the VMEM capacity constraint using a binary search algorithm, including: Based on the row segmentation strategy or the column segmentation strategy, determine the initial search interval for the binary search algorithm; Based on the memory requirement model, the current total memory usage corresponding to the middle block size of the current search interval is calculated iteratively. After comparing it with the VMEM capacity, the current search interval is shrunk until the maximum feasible block size is determined.

[0012] In one possible implementation, the hardware alignment adjustment includes: If both the number of input channels and the number of output channels are greater than or equal to 128, then the number of input channels and the number of output channels are rounded up to a multiple of 128, and the block size is aligned to a multiple of 128 through bit operations. If both the number of input channels and the number of output channels are less than 128, the maximum divisible size is calculated using a preset block function.

[0013] In one possible implementation, determining the computational adaptation path of the 1×1 sub-kernel based on convolutional parameter features and using a hierarchical multi-branch optimization strategy to adapt to different computational scenarios includes: When the kernel size is 1×1, the 1×1 convolution branch is enabled to compute the 1×1 sub-kernel; When the convolution kernel size is greater than 1×1 and the number of input channels and the number of output channels are both less than or equal to 128, the small channel branch is enabled to calculate the 1×1 sub-kernel, and two adjacent 1×1 sub-kernels are processed at one time. When the convolution kernel size is greater than 1×1 and the number of input channels and / or output channels is greater than 128, the general branch is enabled to compute the 1×1 sub-kernel.

[0014] In one possible implementation, according to the optimal block size and the adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through the 1×1 sub-kernel to obtain the calculation result, including: Based on the current spatial position of the 1×1 sub-kernel, calculate the offset related to the position of the convolution kernel to determine the sliding access address of the input data; According to the sliding access address corresponding to the offset, a single DMA transfer is performed. Large-scale continuous input data is used to move data from the TPU's HBM to the VMEM; Indicates window size. Indicates the number of input channels; The dual PGX matrix multiplication unit of the TPU is activated, and matrix multiplication operations are performed in parallel based on the input data in the VMEM and the weight data of the 1×1 sub-core. During matrix multiplication, the intermediate results of each parallel computation are accumulated in real time to obtain the final result.

[0015] In one possible implementation, the hierarchical memory architecture includes HBM, VMEM, and CMEM; the double buffering mechanism is to divide the current computation buffer and the preload buffer in VMEM or CMEM, asynchronously preload the next block of data when computing the current block of data, and switch the buffer through pointer swapping; for 1×1 convolution scenarios, the output buffer is deployed in CMEM.

[0016] This invention provides a method for calculating Conv2D operators on a TPU chip, by processing the... Size convolution kernel decomposed into Each 1×1 sub-core performs calculations directly on the original input data via a sliding window, avoiding redundant expansion of the explicit im2col and HBM transfer space, thus significantly reducing memory usage. Based on the VMEM capacity, convolution parameters, and input size on the TPU chip, the optimal block size and segmentation strategy are dynamically determined, providing memory adaptation support for the 1×1 sub-core calculation, ensuring no memory overflow and high computational efficiency. Based on the convolution parameter features, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios, determining the computation adaptation path of the 1×1 sub-core, enabling multi-dimensional deep adaptation to TPU hardware characteristics and significantly improving hardware utilization. According to the optimal block size and adaptation path, a sliding window mechanism is used to perform matrix multiplication operations on the original input data through the 1×1 sub-core to obtain the calculation results. This fully utilizes the parallel processing capability of the TPU's dual PGX matrix multiplication units, enabling it to concurrently process multiple weight blocks and matrix multiplication core calculations, efficiently converting hardware idle periods into continuous computing power. All calculation results are accumulated to obtain the final output of the Conv2D operator. Attached Figure Description

[0017] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the 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.

[0018] Figure 1 This is a flowchart illustrating the implementation of the Conv2D operator calculation method on the TPU chip provided in this embodiment of the invention. Figure 2 This is a flowchart illustrating the implementation of dynamically determining the optimal block size and segmentation strategy provided in this embodiment of the invention. Figure 3 This is a flowchart illustrating the implementation of a Conv2D operator calculation method on a TPU chip according to another embodiment of the present invention. Figure 4 This is a flowchart illustrating the implementation of matrix multiplication using a sliding window mechanism, as provided in an embodiment of the present invention. Detailed Implementation

[0019] In the following description, specific details such as particular system architectures and techniques are set forth for illustrative purposes and not for limitation, in order to provide a thorough understanding of the embodiments of the invention. However, those skilled in the art will understand that the invention can be implemented in other embodiments without these specific details. In other instances, detailed descriptions of well-known systems, apparatuses, circuits, and methods are omitted so as not to obscure the description of the invention with unnecessary detail.

[0020] To make the objectives, technical solutions, and advantages of the present invention clearer, specific embodiments will be described below in conjunction with the accompanying drawings.

[0021] Figure 1 This is a flowchart illustrating the implementation of a Conv2D operator calculation method on a TPU chip, as provided in an embodiment of the present invention.

[0022] First, let's introduce the hardware characteristics of the TPU: ① Core computing unit architecture.

[0023] The Xuanyuan Sword (XYS) core and its computing units: The TPU chip integrates two completely independent XYS cores, supporting multi-core parallel computing and serving as the core carrier of control and computation flow. Each XYS core includes a Scalar processor responsible for scalar operations, Direct Memory Access (DMA) control, and privilege register management; a vector processor equipped with 32 vector data registers, capable of processing 8×128 elements in parallel per instruction; and dual Pangu Ave (PGX) matrix multiplication units. Each PGX unit is dedicated matrix multiplication hardware, corresponding to an independent matrix result FIFO (MRF) (matrix multiplication result output unit). A single PGX supports matrix multiplication operations on sub-blocks of up to 8×128×128 elements. The dual PGX units support parallel execution of matrix multiplication operations, capable of processing two independent computational subtasks simultaneously, and forming a high-efficiency pipeline with asynchronous DMA transfer, thereby maximizing hardware parallelism and significantly improving overall computational throughput.

[0024] ② Hierarchical memory system: High Bandwidth Memory (HBM): Serves as a global data storage medium with a capacity of 64GB. It is used to store complete input tensors, weight matrices, and final output results, and interacts with VMEM via DMA.

[0025] Vector Memory (VMEM): Each XYS core is equipped with an independent 16MB capacity, which is used to cache input sub-blocks, intermediate results and weight blocks for vector operations.

[0026] Common Memory (CMEM): 32MB in size, it is a shared data cache space for both XYS cores. Its main function is to store reusable data (such as weights and constants) and provide an efficient data exchange channel between different computing cores or computing stages, thereby significantly reducing the interaction frequency and latency between the cores and HBM.

[0027] Memory access constraints: Vector instructions require data to be aligned to 512-byte boundaries. Unaligned access will cause instruction splitting and reduce transmission efficiency.

[0028] ③ Data transmission and command characteristics: DMA Asynchronous Transfer: Supports asynchronous transfer between HBM and on-chip memory, can be executed in parallel with computation, and supports batch transfer. Large-scale continuous data, adapted to convolutional sliding input requirements. Indicates window size. Indicates the number of input channels.

[0029] Dedicated instruction set: Includes register-level transpose instructions (m_transpose), vector loading / storage instructions (load8_128_stride_ldmk, etc.), and block transpose instructions (m_transpose_start / mid / end), adapted to convolution core operations.

[0030] Instruction-level parallelism: The VLIW architecture supports packaging multiple independent operations into a very long instruction word, and executing multiple operations in parallel within a single clock cycle.

[0031] Existing TPU-based convolutional schemes are as follows: The traditional im2col display solution is as follows: im2col_DMA_withPad(input_hbm, hbm_tmp, ...); / / HBM→HBM transfer matmul_all_f32(hbm_tmp, weight_hbm, out_hbm); / / GEMM calculation Traditional display im2col solutions do not utilize the dual XYS cores and dual PGX of the TPU, resulting in low hardware utilization. They rely on the HBM transfer space (hbm_tmp), leading to data redundancy. Secondary DMA transfers increase latency, and they are not adapted to the hierarchical memory characteristics of VMEM / CMEM, resulting in a high risk of network congestion.

[0032] Data reorganization is completed during the loading process from global memory to shared memory without the need for explicit im2col. However, the GPU memory model and instruction set differ significantly from those of the TPU, making direct porting impossible. Furthermore, it is not adapted to the TPU's dual XYS cores, dual PGX units, and CMEM features, and the TPU's 512-byte memory alignment requirement is not considered. This would result in significant unaligned access losses after porting, failing to leverage the TPU's hardware advantages.

[0033] The output buffer is deployed in VMEM, failing to fully utilize the low latency and high capacity hardware features of CMEM; and a general block partitioning strategy is used for... Divide into blocks, , This represents the height of the output feature map after the convolution operation. This represents the width of the output feature map after the convolution operation. This indicates the effective element size of the output feature map in a single block processing operation, and the current block size (referring to...). or Smaller size leads to The shortcomings include limited matmul computation efficiency. In scenarios with stride=1 and padding=0, batch DMA loading optimization was not implemented, failing to fully utilize the TPU's DMA batch transfer capabilities, resulting in an excessively high DMA transfer ratio. Furthermore, the parallel processing capabilities of the dual PGX matrix multiplication units were not utilized, hindering their parallel computing advantages in small-to-medium channel scenarios and further limiting computational efficiency. Additionally, the H2v Ping-pong mechanism was not adapted, leading to serial data transmission and computation, resulting in low overall hardware utilization.

[0034] This invention addresses the aforementioned problems and achieves efficient computation of large-size convolutions by employing a Conv2D operator computation method on a TPU chip, utilizing the Implicit GEMM algorithm, and combining TPU dual XYS cores, dual PGX units, weighted block division, sliding input, CMEM transfer, and other technologies.

[0035] The calculation method of the Conv2D operator on the TPU chip is described in detail below: Step 101, process the... Size convolution kernel decomposed into Each 1×1 subnucleus.

[0036] Indicates the height of the convolution kernel. This indicates the width of the convolution kernel.

[0037] This step uses weighted block partitioning technology to divide the material to be processed. Size convolution kernel is decomposed into at the algorithm level as Each independent 1×1 sub-core lays the foundation for subsequent implicit data recombination and efficient computation.

[0038] This decomposition process does not physically split and store the convolutional kernels. Instead, it defines the spatial correspondence between each 1×1 sub-kernel and the original large convolutional kernel through algorithmic logic. That is, each 1×1 sub-kernel corresponds to a different location in the original kernel. The weight parameters of the convolutional kernel at different spatial positions. After decomposition, in subsequent calculations, there is no need to explicitly perform im2col reorganization on the input data. Instead, matrix multiplication operations of each 1×1 sub-kernel are directly executed on the original input data through a sliding window mechanism. Finally, by accumulating the calculation results of all 1×1 sub-kernels, an output completely equivalent to the original convolutional operation is achieved.

[0039] From the perspective of mathematical equivalence, in the traditional explicit GEMM scheme, the input data needs to be reorganized into matrix, and then multiplied by the weight matrix; after this step of decomposition, Implicit GEMM only needs to loop times, and each time a 1×1 window slides on the original input to perform matrix multiplication, and the intermediate results are accumulated in CMEM (cache memory). Finally, the accumulated sum is exactly the same as the original convolutional output.

[0040] This decomposition process and the execution process of the matrix multiplication operation in step 103 are implemented through a four-layer loop structure of output channel block division, input channel block division, convolutional kernel position traversal, and spatial dimension sliding processing; the specific code logic is as follows: plaintext for(int ih_weight = 0; ih_weight<C_out_low128; ih_weight += 128) { / / Output channel block division for(int iC_in = 0; iC_in<C_in_low256; iC_in += 256) { / / Input channel block division for(int ik = 0; ik<n_kernel; ik++) { / / Convolutional kernel position traversal (n_kernel = k_h× k_w) / / Sliding processing of height and width in the spatial dimension slide_left_aw_matmul_f32(...); / / Sliding window matrix multiplication } } } For output channel block division, based on the hardware architecture of the TPU, determine the output channel block division base number, and divide the output channels according to the output channel block division base number; each time, load the weight data corresponding to the output channel block into VMEM to adapt to the capacity constraint of VMEM and the channel alignment access requirements of the TPU; The TPU's VMEM is only 16MB, and it requires the number of channels to be aligned to multiples of 128 (ALIGN128). If the number of output channels ( If the weights are large (e.g., 512, 1024), directly loading the weights of all output channels would exceed the VMEM capacity. Therefore, this loop divides the output channels into blocks of 128, loading only the weight blocks corresponding to 128 output channels into the VMEM each time. This avoids the weight data occupying too much memory, while also meeting the channel alignment access requirements of the TPU hardware, laying the foundation for the parallel execution of subsequent matrix multiplication.

[0041] Input channel segmentation is based on the TPU hardware architecture. The input channel segmentation base number is determined and the input channels are segmented according to the input channel segmentation base number. The original input data of the corresponding input channel segment is loaded into VMEM each time to match the parallel processing capability of the computing units in the TPU. Number of input channels This could also far exceed the VMEM caching capacity, and the optimal input channel size for a single processing cycle of the TPU's dual PGX matrix multiplication unit is 256. Therefore, the loop divides the input channels into blocks of 256, loading only the raw data from 256 input channels into the VMEM each time. This ensures that the amount of data processed in a single computation matches the processing capacity of the TPU computing unit, avoiding computational blockage due to excessive data volume, while simultaneously improving the parallel efficiency of data loading and computation.

[0042] The convolution kernel position traversal is a traversal. The spatial locations of each 1×1 sub-kernel are used to transform the single convolution operation of the original large convolution kernel into batch computation of each 1×1 sub-kernel; this avoids explicit im2col data recombination and ensures that the computation results of each sub-kernel can be merged into the equivalent result of the original convolution in the subsequent accumulation stage. Spatial dimension sliding processing involves sliding a 1×1 sub-kernel in the height and width directions of the input feature map, and achieving continuous access to the original input data through offset calculation related to the position of the convolution kernel.

[0043] Spatial dimension sliding is the "implicit reorganization core" of Implicit GEMM, responsible for sliding a 1×1 sub-kernel in the height and width directions of the input feature map: By calculating offsets related to the convolution kernel position, continuous sliding access to the original input data is achieved, enabling single DMA transfers. Large-scale continuous data avoids DMA efficiency degradation caused by non-continuous access; During the sliding process, the 1×1 sub-core performs matrix multiplication with the current input sub-block (relying on the parallel computation of dual PGX), and the intermediate results are accumulated in real time to the accumulator in CMEM, without the need for additional storage of intermediate data, thus maximizing data reuse rate; Finally, by sliding through the entire input feature map, all spatial locations are covered, and all calculations of the current 1×1 sub-kernel under the current input / output channel block are completed.

[0044] Through this four-layer loop structure, each 1×1 sub-kernel can complete the calculation independently and be merged in the accumulation stage, which not only ensures mathematical equivalence but also maximizes data reuse. During the convolution kernel traversal, the loaded input data can be reused by multiple 1×1 sub-kernels, avoiding the problems caused by explicit im2col transformation. The doubled input data expansion significantly reduces the storage overhead of VMEM (Vector Memory), requiring only the memory space of "current 1×1 weight + current input / output block".

[0045] Step 102: Based on the VMEM capacity, convolution parameters and input size on the TPU chip, dynamically determine the optimal block size and segmentation strategy to provide memory adaptation support for the computation of 1×1 sub-cores.

[0046] The dynamic block partitioning strategy described in this step is a key step in adapting to the VMEM capacity constraints of the TPU chip and ensuring the efficient execution of the Conv2D operator. Its core is to calculate the optimal block size and partitioning strategy in real time based on the VMEM capacity, convolution parameters and input size of the TPU, so as to avoid memory overflow and maximize computational efficiency.

[0047] In one embodiment, such as Figure 2 As shown, based on the VMEM capacity, convolution parameters, and input size on the TPU chip, the optimal block size and segmentation strategy can be dynamically determined, which may include: Step 201: Establish a memory requirement model, and based on the memory requirement model, calculate the total memory usage of the input sub-blocks, output sub-blocks, weight data, and bias data required to process the target output block.

[0048] To accurately determine the feasibility of the block size, a memory requirement model must first be established to comprehensively cover the memory overhead of various data types required to process the target output block. Specifically, the total memory usage includes four parts: input sub-blocks, output sub-blocks, weight data, and bias data. This calculation must also consider the hardware alignment requirements of the TPU. Taking FP32 data format (each data occupies 4 bytes) as an example, the memory calculation formula is as follows: Input sub-block memory is ; in Indicates input sub-block memory. This represents the input height required to process the current output block (derived from the block size, stride, and kernel height). This indicates the input width after padding. Indicates the number of input channels Round up to a multiple of 128 to match the channel alignment requirements of TPU; Output sub-block memory is ; in, This indicates that the sub-block of memory is being output. Indicates the current attempted block size (in height or width direction). Indicates the width of the output feature map. This indicates that the number of output channels chan_out will be rounded up to a multiple of 128; Weighted data memory is The space occupied by the 1×1 sub-core weight corresponding to the current block; in, This represents the memory location of the weighted data.

[0049] Bias data memory is This corresponds to the space occupied by the bias term in the current output channel block. This indicates the biased data memory.

[0050] The above formula can accurately calculate the total memory usage of the target output block, providing a quantitative basis for subsequent block division feasibility assessment.

[0051] Step 202: Calculate the output feature map size based on the convolution parameters and input size, and determine the segmentation strategy based on the output feature map size.

[0052] In this embodiment, the convolution parameters and input size are directly input basic parameters. The convolution parameters refer to the configuration parameters that define the convolution operation itself, including: convolution height. Convolution width Vertical step size Horizontal step size Vertical filling Horizontal filling ; Input size refers to the dimensional parameters of the input data to be processed, including: the height of the input feature map. Input feature map width Number of input channels .

[0053] In addition, the basic parameters that are directly input also include the number of output channels. .

[0054] In one embodiment, the output feature map size is calculated based on the convolution parameters and the input size. The output feature map size may include the height and width of the output feature map. The processing flow may include: according to , , as well as This yields the height of the output feature map. according to , , as well as This yields the width of the output feature map.

[0055] Optional, according to Obtain the height of the output feature map; Indicates the height of the output feature map; Optional, according to Obtain the width of the output feature map; This indicates the width of the output feature map.

[0056] The choice of segmentation strategy needs to be adapted to the shape of the output feature map to balance data continuity and memory pressure. This step may include: When the height of the output feature map is greater than or equal to its width, a row segmentation strategy is adopted. This strategy divides the data into blocks along the height direction, which can maintain the continuity of the output data in memory, maximize DMA transfer efficiency, and match the TPU's optimization characteristics for continuous data access. When the height of the output feature map is less than its width, a column-segmentation strategy is adopted. This strategy divides the data into blocks along the width direction, which can reduce the width of the input data required for a single computation, effectively alleviate the capacity pressure of the TPU's limited VMEM, and avoid memory overflow caused by excessively large input data spans.

[0057] Step 203: Based on the partitioning strategy and memory requirement model, determine the maximum feasible block size by iteratively calculating within the VMEM capacity constraint using a binary search algorithm.

[0058] After defining the partitioning strategy, a binary search algorithm is used to iteratively find the maximum feasible block size within the VMEM capacity constraint, ensuring that memory usage does not exceed the limit and that the partitioning efficiency is optimal.

[0059] In one embodiment, the maximum feasible block size is determined by iterative calculation within the VMEM capacity constraint using a binary search algorithm based on the partitioning strategy and memory requirement model. This may include: The initial search interval of the binary search algorithm is determined based on either row-segmentation or column-segmentation strategies. Based on the memory requirement model, the current total memory usage corresponding to the middle block size of the current search interval is calculated iteratively. After comparing it with the VMEM capacity, the current search interval is shrunk until the maximum feasible block size is determined.

[0060] Optionally, based on the selected segmentation strategy, the initial search interval for row segmentation is... ( (to be the total height of the output feature map), the initial search interval for column segmentation is... ( (This refers to the total width of the output feature map), i.e., the left boundary. Initialized to 1, right boundary Initialize to the total size of the corresponding dimension of the output feature map; when At that time, calculate the midpoint of the current search interval. (Round up to ensure that no larger feasible dimensions are missed). This is the current intermediate block size (block height when splitting by rows, block width when splitting by columns); Will Substituting the values ​​into the memory requirement model above, calculate the corresponding current total memory usage and compare it with the TPU's VMEM capacity, VmemSize. If the total memory usage is less than or equal to VmemSize, it indicates... The corresponding block size is feasible, and the left boundary is... Updated to This allows for a larger search space; if the total memory usage exceeds VmemSize, it indicates... Exceeding memory constraints, move the right boundary. Updated to Narrowing the search scope; Repeat the above iterative process until... The search process ends at this time. (or The corresponding size is the maximum number of output rows (or columns) that can be processed under the VMEM capacity constraint, which is the maximum feasible block size.

[0061] For transpose scenarios, the maximum feasible block height needs to be limited to less than or equal to 60 to avoid discontinuous memory access due to excessively large spans, which would affect transmission efficiency. Furthermore, the binary search logic in column partitioning scenarios is consistent with row partitioning; only the input height in the memory calculation formula needs to be replaced with the input width to adapt to block calculations in the width direction.

[0062] Step 204: Hardware alignment adjustment is performed on the maximum feasible block size and number of channels to adapt to the access requirements of the TPU.

[0063] To meet the TPU's 512-byte aligned access requirement (unaligned access leads to instruction splitting, reducing transmission and computation efficiency), hardware alignment adjustments are needed to the maximum feasible block size and number of channels, specifically including: If both the number of input channels and the number of output channels are greater than or equal to 128, then the number of input channels will be... Number of output channels All data is rounded up to multiples of 128 using the ALIGN128 function to ensure that data access in the channel dimension conforms to the TPU vector instruction requirements. Furthermore, block height and width are aligned to multiples of 128 using bit operations (such as &0xFFFFFF80) to ensure that data transmission in the spatial dimension matches the processing granularity of the computing unit.

[0064] If both the number of input channels and the number of output channels are less than 128 ( and Specifically, for small channel scenarios, a preset block function is enabled, such as the dedicated block function Calc_MaxH_1x2, which is combined with 1×2 convolution kernel fusion optimization; for 1×1 convolution scenarios, a preset block function is used, such as the dedicated block function Calc_MaxH_1x1, which is adapted to its continuous memory access characteristics to further improve block efficiency.

[0065] The above alignment adjustments ensure that the segmented data can be efficiently accessed and processed by the TPU, laying the foundation for subsequent matrix multiplication operations.

[0066] See Figure 3 As shown, the calculation method of the Conv2D operator on the TPU chip also includes: Step 106: Utilize the hierarchical memory architecture and double buffering mechanism of the TPU chip to achieve parallel execution of computation and data transmission during matrix multiplication.

[0067] The core of this step is to fully utilize the hierarchical memory architecture and double buffering mechanism of the TPU chip to achieve parallel execution of computation and data transmission during matrix multiplication, thereby maximizing the masking of data transmission latency and improving hardware utilization and overall computational efficiency.

[0068] In one embodiment, the hierarchical memory architecture includes HBM, VMEM, and CMEM. Each level of memory collaborates based on its differentiated characteristics of "storage capacity, access latency, and bandwidth" to form an efficient data flow channel. HBM: As a global data storage carrier, it has the characteristics of large capacity and high bandwidth. It is used to store the original input tensors, complete weight matrices and final output results of Conv2D operators. It performs asynchronous data interaction with on-chip memory (VMEM / CMEM) through DMA. It is the "source" and "end" of data flow. VMEM: Each XYS core is equipped with a 16MB independent VMEM, which has the advantages of low latency and high bandwidth. It is mainly used to cache the current computation input sub-block, weight block and intermediate results in general convolution scenarios, reduce the frequent access to HBM, and provide high-speed data supply for vector processor and PGX matrix multiplication unit. CMEM: 32MB capacity, focusing on low-latency and highly scalable data access, with core adaptation to 1×1 convolution scenarios - the output buffer of 1×1 convolution is specifically deployed in CMEM, taking advantage of its physical characteristics of being close to the computing unit to reduce the latency of intermediate result accumulation and write-back, while also adapting to the batch DMA transfer requirements of TPU.

[0069] The three form a hierarchical data flow of "HBM→VMEM→CMEM". Data is asynchronously moved from HBM to VMEM / CMEM according to computing needs. After the computing is completed, the results are written back to HBM, ensuring that data storage and access are adapted to the needs of different computing stages.

[0070] The double buffering (Ping-pong) mechanism is the core means to achieve "parallel computation and data transmission". By dividing two buffers of equal capacity in VMEM or CMEM, the time difference between asynchronous DMA transmission and computation is used to mask the data transfer delay.

[0071] In one embodiment, the double buffering mechanism is to divide the current computation buffer and the preload buffer in VMEM or CMEM, and asynchronously preload the next block of data when computing the current block of data, and switch the buffer through pointer swapping; for 1×1 convolution scenarios, the output buffer is deployed in CMEM.

[0072] Optionally, in VMEM (general convolution scenario) or CMEM (1×1 convolution scenario), a "current computation buffer" and a "preload buffer" are divided. The two buffers are stored independently and used alternately to ensure that data loading and computation do not interfere with each other. When the TPU's computing unit (dual PGX matrix multiplication unit) performs matrix multiplication on the data in the "current computing buffer", it moves the next block of data to be computed from HBM to the "preload buffer" via DMA asynchronous transfer. After the computing unit completes the computation of the current block, it does not need to wait for the next block of data to be loaded. It directly switches the buffer through pointer swapping - the original "preload buffer" becomes the new "current computing buffer", and the original "current computing buffer" is used to receive the next block of preloaded data, realizing a seamless connection between "computation-loading". By managing the buffer's usage state through a state machine, pointer swapping operations only require modifying memory address pointers, eliminating the need for data copying and making the switching overhead negligible. In general convolution scenarios, this mechanism primarily hides the data transfer latency from HBM to VMEM; in 1×1 convolution scenarios, it can further hide the result write-back latency from CMEM to HBM, achieving 100% overlap between computation and transmission.

[0073] It should be noted that the collaboration between hierarchical memory and double buffering is not a fixed pattern, but rather dynamically adapted to different computing scenarios to ensure maximum optimization effect: In general convolution scenarios (non-1×1 convolution): the buffer is deployed in VMEM, and the double buffering mechanism focuses on the preloading of input data between HBM and VMEM. When calculating the matrix multiplication of the current input sub-block, the next input sub-block is preloaded asynchronously to avoid the computing unit being idle due to waiting for data and to give full play to the parallel computing capabilities of the dual PGX units. 1×1 convolution scenario: The output buffer is migrated to CMEM, and the double buffering mechanism achieves dual optimization. On the one hand, the input data preloading from HBM to CMEM is parallel to the matrix multiplication in CMEM; on the other hand, the accumulation of intermediate results in CMEM is parallel to the output write-back from CMEM to HBM. At the same time, combined with the continuous characteristics of 1×1 convolution data, batch DMA transfer is used to reduce startup overhead and further improve parallel efficiency. Memory access adaptation: All data storage and DMA transfers in the buffer strictly adhere to the 512-byte alignment requirement of the TPU, while also adapting to the Channels_last data format to ensure that memory access is free from instruction splitting loss, maximizing DMA transfer bandwidth and memory access efficiency.

[0074] Through the above design, the TPU's computing unit and DMA transfer unit can be in working state simultaneously, greatly reducing hardware idle time. Especially in convolution scenarios with large input size and high channel number, it can significantly reduce the proportion of data transmission, fully release the computing potential of the TPU, and achieve efficient execution of Conv2D operators.

[0075] Step 103: Based on the convolution parameter features, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios and determine the computing adaptation path of the 1×1 sub-kernel.

[0076] The core of this step is to design a hierarchical multi-branch optimization strategy based on the features of convolutional parameters. By accurately matching different computing scenarios and TPU hardware characteristics, the optimal computing adaptation path is determined for the 1×1 sub-cores decomposed in step 101, maximizing hardware utilization and computing efficiency. This multi-branch optimization strategy uses "convolutional kernel size + number of channels" as the core judgment criteria to form a three-level branch logic.

[0077] In one embodiment, based on the features of convolution parameters, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios and determine the computing adaptation path for a 1×1 sub-kernel, which may include: When the kernel size is 1×1, the 1×1 convolution branch is enabled to compute the 1×1 sub-kernel; When the kernel size is greater than 1×1 and the number of input channels and output channels are both less than or equal to 128, the small channel branch is enabled to compute the 1×1 sub-kernel, and two adjacent 1×1 sub-kernels are processed at one time. When the kernel size is greater than 1×1 and the number of input channels and / or output channels is greater than 128, the general branch is enabled to compute the 1×1 subkernel.

[0078] Optionally, when the kernel size is 1×1 (i.e. and In step 101, after decomposition, only one 1×1 sub-kernel is obtained. At this point, the 1×1 convolution branch is activated, fully adapting to the characteristics of strong data continuity and simple computational logic in this scenario, while also deeply utilizing the hierarchical memory and DMA transfer advantages of the TPU. Specific implementation points include: Output buffer deployment optimization: The output buffer of the 1×1 convolution is migrated to the TPU's CMEM (cache memory). By leveraging the low latency and high bandwidth of CMEM, the latency of intermediate result storage and access is reduced, replacing the traditional mode that relies on VMEM to store the output. Data transfer optimization: Differentiated DMA transfer strategies are adopted for different parameter combinations. When stride=1 and padding=0, the input and output data are continuously distributed in memory, and a continuous memory copy method is directly adopted without additional offset calculation, maximizing DMA transfer efficiency. In the general 1×1 convolution scenario, a batch DMA transfer mode is adopted to transfer continuous data of the Win×Cin scale in a single transaction, reducing the number of DMA startups and transfer overhead. Parallel computation and transmission: By combining the double-buffered (Ping-pong) DMA mechanism, the next block of input data is asynchronously preloaded while the matrix multiplication calculation of the current data block is being performed, so as to achieve complete overlap between computation and data transmission, mask the DMA transmission delay, and greatly improve the overall execution efficiency.

[0079] Optionally, when the convolutional kernel size is greater than 1×1 (e.g., 3×3, 5×5), and the number of input channels and output channels are both less than or equal to 128, a small-channel branch is enabled. In this scenario, the number of channels is relatively small. If the 1×1 sub-kernels are enumerated one by one according to the general logic, it will lead to idle hardware computing units. Therefore, convolutional kernel fusion optimization is used to improve computational density. Specific implementation points include: Dedicated calculation function adaptation: Design a dedicated sliding window matrix multiplication function for slide_left_aw_matmul_f32_1x2_kernel, and optimize the calculation logic for small channel scenarios; Convolutional kernel fusion processing: Load and process two adjacent 1×1 sub-kernels (i.e., a combination of sub-kernels of size 1×2) at once, reducing data transfer overhead by reducing the number of traversals of 1×1 sub-kernels and the number of weight loadings; Dual PGX Parallel Adaptation: Fully utilize the dual PGX matrix multiplication units of the TPU to distribute the computational tasks of two adjacent 1×1 sub-cores to two PGX units for parallel execution, further improving the computational efficiency in small channel scenarios and avoiding idle hardware resources.

[0080] Optionally, when the kernel size is greater than 1×1 and the number of input channels and / or output channels is greater than 128, a general branch is enabled. This branch is the complete implementation of the core logic of step 101, "Sliding Window ImplicitGEMM + Weight Blocking," adapting to complex scenarios such as large channels and large kernels. Specific implementation details include: Core algorithm execution: Strictly follows the four-layer loop structure of "output channel → input channel → convolution kernel position → spatial dimension", splits the output channel into units of 128 and the input channel into units of 256, traverses all 1×1 sub-kernels and performs sliding window matrix multiplication; Dynamic block adaptation: Integrating the dynamic block strategy determined in step 102, the data blocks are divided according to the VMEM capacity of the TPU, the input size and the convolution parameters, through row partitioning or column partitioning strategies to ensure that memory does not overflow and the computational efficiency is optimal. Full utilization of hardware resources: The parallel capabilities of the TPU's dual XYS cores and dual PGX matrix multiplication units are utilized, combined with the hierarchical memory (HBM→VMEM→CMEM) data flow design, to achieve efficient collaboration of data loading, matrix multiplication, and intermediate result accumulation. At the same time, the transmission delay is masked through a double buffering mechanism to maximize hardware utilization.

[0081] The aforementioned three-level branch optimization strategy achieves precise coverage of different computing scenarios through a hierarchical logic of "first determining the convolution kernel size, then determining the number of channels." This avoids the waste of efficiency in some scenarios caused by a "one-size-fits-all" optimization mode. It allows 1×1 convolution scenarios to enjoy the advantages of low-latency storage and continuous transmission, small-channel scenarios to improve computing density through kernel fusion, and general scenarios to ensure stability and efficiency through a complete core algorithm. Ultimately, it maximizes the release of TPU hardware performance under various combinations of convolution parameters.

[0082] Step 104: According to the optimal block size and the adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through a 1×1 sub-kernel to obtain the calculation result.

[0083] This step is the implementation stage of the core logic of the Conv2D operator Implicit GEMM. Following the optimal block size determined in step 102 and the computation adaptation path selected in step 103, it performs matrix multiplication of 1×1 sub-kernels on the original input data through a sliding window mechanism, ultimately obtaining the computation results corresponding to each sub-kernel. This process does not require explicit im2col recombination of the input data. Through the ordered flow of "offset calculation - data transmission - parallel computation - result accumulation", it achieves efficient and memory-friendly convolution computation.

[0084] In one embodiment, such as Figure 4 As shown, based on the optimal block size and adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through a 1×1 sub-kernel to obtain the calculation results, which may include: Step 401: Based on the current spatial position of the 1×1 sub-kernel, calculate the offset related to the position of the convolution kernel to determine the sliding access address of the input data.

[0085] Since the original input data has not been explicitly reconstructed, precise offset calculation is required to enable the sliding access of the 1×1 sub-kernel on the original input feature map, ensuring that each 1×1 sub-kernel can accurately match the input region corresponding to the spatial location of the original large convolutional kernel.

[0086] Offset calculation basis: combined with the spatial position of the original convolution kernel (determined by the 1×1 sub-kernel index ik after decomposition in step 101, the value of ik ranges from...). ), convolution stride Input Fill and the size after input channel alignment Construct an offset calculation model; Taking horizontal sliding as an example, core offset The calculation formula is ,in The number of input channels is aligned to a multiple of 128. Dividing this by 32 gives the addressing granularity adapted to TPU vector instructions. This offset allows for precise location of the starting access address of the current 1×1 sub-core in the width direction of the input feature map. Similarly, the vertical offset is derived to finally determine the complete sliding access address. During offset calculation, the validity of the access address is checked synchronously to avoid accessing out of bounds due to padding or step size settings, ensuring that each 1×1 sub-core only accesses the valid region in the original input data corresponding to its own spatial position.

[0087] Step 402: Perform a single DMA transfer according to the sliding access address corresponding to the offset. Large-scale continuous input data is used to move data from the TPU's HBM to the VMEM; Indicates window size. Indicates the number of input channels.

[0088] Based on the sliding access address calculated above, the input data is transferred at high speed through the DMA transfer unit of the TPU, providing data support for subsequent matrix multiplication operations.

[0089] A single DMA transfer is fixed as The scale of continuous data is determined by combining the processing granularity of the TPU dual PGX matrix multiplication unit and the VMEM cache capacity. This can avoid frequent DMA startups caused by too small a single data transfer and ensure that the transferred data can be completely stored in VMEM. Strictly adhering to the data flow rules of the TPU hierarchical memory architecture, data is precisely moved from HBM (Global Storage) to VMEM (Vector Memory) via DMA. Furthermore, the data storage format strictly adapts to the TPU's 512-byte alignment requirement and Channels_last format to avoid instruction splitting during subsequent accesses and maximize memory access efficiency. If the current adaptation path is a 1×1 convolutional branch or a general branch, the double buffering mechanism in step 106 will be activated synchronously. After the data transmission is completed, the asynchronous preloading of the next piece of data will be started immediately, laying the foundation for the parallel execution of subsequent calculations and transmissions.

[0090] Step 403: Start the dual PGX matrix multiplication unit of the TPU and perform matrix multiplication operations in parallel based on the input data in VMEM and the weight data of the 1×1 sub-core.

[0091] After the data is loaded into VMEM, the dual PGX (parallel general matrix multiplication execution unit) of the TPU is started according to the adaptation path to complete the parallel computation of the 1×1 sub-core and the input data.

[0092] When scheduling and adapting computing units, if it is a small channel adaptation path (the number of input channels and the number of output channels are both less than or equal to 128), the dedicated computing function slide_left_aw_matmul_f32_1x2_kernel is started to allocate the computing tasks of two adjacent 1×1 sub-cores to two PGX units respectively, so as to realize the parallel computing of 1×2 sub-core combination and improve the computing density of small channel scenarios. If it is a 1×1 convolution branch or a general branch, the computation task of the current block is allocated to the dual PGX units based on the channel block corresponding to the optimal block size (output channels are divided into 128 blocks and input channels are divided into 256 blocks). Each PGX unit processes the matrix multiplication of the corresponding channel block, giving full play to the parallel computing power of the dual PGX units.

[0093] by Dimensional input sub-block matrix ( , (the spatial dimensions after partitioning) and 1×1 sub-kernel weight matrix of dimension ( The output channel block size is used as input, and the two PGX units synchronously perform matrix multiplication operations. The operation process strictly follows the TPU's vector instruction set specification to ensure computational efficiency and accuracy.

[0094] Step 404: During the matrix multiplication operation, the intermediate results of each parallel computation are accumulated in real time to obtain the calculation result.

[0095] During matrix multiplication, intermediate results are accumulated in real time through the TPU's CMEM (computation memory), avoiding redundant storage of intermediate data.

[0096] By using FP32 precision to accumulate the intermediate results generated by parallel computation of dual PGX units, compared with FP16 accumulation, the precision loss caused by multiple accumulations can be effectively avoided, ensuring that the final calculation result is completely equivalent to the original convolution operation. The accumulation process is completed in CMEM, which has the characteristics of low latency and high bandwidth, and can quickly respond to the writing and reading needs of intermediate results. Especially in the 1×1 convolution branch scenario, the output buffer has been deployed in CMEM, and the accumulation result can be directly written to the output buffer, reducing data transfer overhead. The cumulative result of the current 1×1 sub-kernel will serve as the basis for the calculation results of other sub-kernels at the same spatial position. After all 1×1 sub-kernels have completed their calculations, the final cumulative sum is the convolution result of the corresponding output pixel, achieving mathematical equivalence with a single convolution operation of the original large convolution kernel.

[0097] In summary, step 104, through the orderly execution of the above four steps, not only avoids the memory redundancy problem caused by explicit im2col, but also fully adapts to the characteristics of TPU hardware through the design of dual PGX parallel computing and precise offset positioning. At the same time, it strictly follows the constraints of optimal block size and adaptation path, ensuring that the efficiency and accuracy of convolution calculation are maximized without memory overflow.

[0098] Step 105: Sum all the calculation results to obtain the final output of the Conv2D operator.

[0099] This invention provides a method for calculating Conv2D operators on a TPU chip, by processing the... Size convolution kernel decomposed into A 1×1 sub-core is used; based on the VMEM capacity, convolution parameters, and input size on the TPU chip, the optimal block size and segmentation strategy are dynamically determined to provide memory adaptation support for the computation of the 1×1 sub-core; based on the convolution parameter features, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios and determine the computation adaptation path of the 1×1 sub-core; according to the optimal block size and adaptation path, a sliding window mechanism is used to perform matrix multiplication operations on the original input data through the 1×1 sub-core to obtain the calculation results; all calculation results are accumulated to obtain the final output of the Conv2D operator.

[0100] The embodiments of the present invention will... Size convolution kernel split into Each 1×1 sub-core directly computes on the original input data via a sliding window, avoiding redundant memory expansion in the explicit im2col and HBM transfer space, thus significantly reducing memory usage. Dynamic block size calculation based on VMEM capacity, convolution parameters, and input size enables dynamic block partitioning, supporting row / column splitting and transpose scene adaptation, ensuring no memory overflow and high computational efficiency. A hierarchical multi-branch optimization strategy is employed, with dedicated optimization paths designed for 1x1 convolutions, dual PGX hardware adaptation optimization for small and medium channel scenarios, and special processing for 1x2 convolution kernels, achieving multi-dimensional deep adaptation to TPU hardware characteristics and significantly improving hardware utilization. By employing a sliding window mechanism to perform matrix multiplication operations, the parallel processing capability of the TPU's dual PGX matrix multiplication units is fully utilized, enabling concurrent processing of multiple weight blocks and matrix multiplication core calculations, efficiently converting idle hardware periods into continuous computing power.

[0101] In addition, this embodiment of the invention also sets up a hierarchical memory architecture and a double buffering mechanism to temporarily store the output of 1x1 convolution in CMEM. Combined with the double buffering mechanism, it achieves full overlap between computation and transmission, effectively reducing the proportion of DMA transmission, thereby fully releasing the low latency and high bandwidth advantages of the TPU hierarchical memory architecture.

[0102] It should be understood that the sequence number of each step in the above embodiments does not imply the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of the present invention.

[0103] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features. Such modifications or substitutions 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 invention, and should all be included within the protection scope of the present invention.

Claims

1. A method for calculating Conv2D operators on a TPU chip, characterized in that, include: To be processed Size convolution kernel decomposed into Each 1×1 sub-core; Indicates the height of the convolution kernel. Indicates the width of the convolution kernel; Based on the VMEM capacity, convolution parameters, and input size on the TPU chip, the optimal block size and segmentation strategy are dynamically determined to provide memory adaptation support for the computation of the 1×1 sub-core; Based on the features of convolution parameters, a hierarchical multi-branch optimization strategy is used to adapt to different computing scenarios and determine the computing adaptation path of the 1×1 sub-kernel. According to the optimal block size and the adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through the 1×1 sub-core to obtain the calculation result. The final output of the Conv2D operator is obtained by summing all the calculation results.

2. The method for calculating the Conv2D operator on a TPU chip according to claim 1, characterized in that, Also includes: By utilizing the hierarchical memory architecture and double buffering mechanism of the TPU chip, parallel execution of computation and data transmission is achieved during the matrix multiplication operation.

3. The method for calculating the Conv2D operator on a TPU chip according to claim 2, characterized in that, The decomposition process and matrix multiplication operation are implemented through a four-layer loop structure: output channel segmentation, input channel segmentation, convolution kernel position traversal, and spatial dimension sliding processing. The output channel is segmented based on the hardware architecture of the TPU, the base number of the output channel segments is determined, and the output channels are segmented according to the base number of the output channel segments; each time, the weight data corresponding to the output channel segment is loaded into the VMEM to adapt to the VMEM capacity constraints and the channel alignment access requirements of the TPU. The input channel segmentation is based on the hardware architecture of the TPU, determining the input channel segmentation base number, and segmenting the input channel according to the input channel segmentation base number; Each time, the original input data corresponding to the input channel block is loaded into the VMEM to match the parallel processing capability of the computing unit in the TPU; The convolution kernel position traversal is to traverse the... The spatial location of each 1×1 sub-kernel is used to transform the single convolution operation of the original large convolution kernel into batch computation of each 1×1 sub-kernel; The spatial dimension sliding process involves sliding the 1×1 sub-kernel in the height and width directions of the input feature map, and achieving continuous access to the original input data through offset calculation related to the position of the convolution kernel.

4. The method for calculating the Conv2D operator on a TPU chip according to claim 3, characterized in that, The method of dynamically determining the optimal block size and segmentation strategy based on the VMEM capacity, convolution parameters, and input size on the TPU chip includes: Establish a memory requirement model, and based on the memory requirement model, calculate the total memory usage of the input sub-blocks, output sub-blocks, weight data, and bias data required to process the target output block; Based on the convolution parameters and input size, the output feature map size is calculated, and based on the output feature map size, a segmentation strategy is determined; Based on the segmentation strategy and the memory requirement model, the maximum feasible block size is determined by iterative calculation within the VMEM capacity constraint using a binary search algorithm. The maximum feasible block size and number of channels are adjusted for hardware alignment to adapt to the access requirements of the TPU.

5. The method for calculating the Conv2D operator on a TPU chip according to claim 4, characterized in that, The step of determining a segmentation strategy based on the output feature map size includes: When the height of the output feature map is greater than or equal to its width, a row segmentation strategy is adopted; When the height of the output feature map is less than its width, a column splitting strategy is adopted.

6. The method for calculating the Conv2D operator on a TPU chip according to claim 5, characterized in that, Based on the segmentation strategy and the memory requirement model, the maximum feasible block size is determined through iterative calculation within the VMEM capacity constraint using a binary search algorithm, including: Based on the row segmentation strategy or the column segmentation strategy, determine the initial search interval for the binary search algorithm; Based on the memory requirement model, the current total memory usage corresponding to the middle block size of the current search interval is calculated iteratively. After comparing it with the VMEM capacity, the current search interval is shrunk until the maximum feasible block size is determined.

7. The method for calculating the Conv2D operator on a TPU chip according to claim 6, characterized in that, The hardware alignment adjustment includes: If both the number of input channels and the number of output channels are greater than or equal to 128, then the number of input channels and the number of output channels are rounded up to a multiple of 128, and the block size is aligned to a multiple of 128 through bit operations. If both the number of input channels and the number of output channels are less than 128, the maximum divisible size is calculated using a preset block function.

8. The method for calculating the Conv2D operator on a TPU chip according to any one of claims 1-7, characterized in that, The method of determining the computational adaptation path of the 1×1 sub-kernel based on convolution parameter features and adapting to different computational scenarios through a hierarchical multi-branch optimization strategy includes: When the kernel size is 1×1, the 1×1 convolution branch is enabled to compute the 1×1 sub-kernel; When the convolution kernel size is greater than 1×1 and the number of input channels and the number of output channels are both less than or equal to 128, the small channel branch is enabled to calculate the 1×1 sub-kernel, and two adjacent 1×1 sub-kernels are processed at one time. When the convolution kernel size is greater than 1×1 and the number of input channels and / or output channels is greater than 128, the general branch is enabled to compute the 1×1 sub-kernel.

9. The method for calculating the Conv2D operator on a TPU chip according to any one of claims 1-7, characterized in that, Based on the optimal block size and the adaptation path, a sliding window mechanism is used to perform matrix multiplication on the original input data through the 1×1 sub-kernel to obtain the calculation results, including: Based on the current spatial position of the 1×1 sub-kernel, calculate the offset related to the position of the convolution kernel to determine the sliding access address of the input data; According to the sliding access address corresponding to the offset, a single DMA transfer is performed. Large-scale continuous input data is used to move data from the TPU's HBM to the VMEM; Indicates the window size. Indicates the number of input channels; The dual PGX matrix multiplication unit of the TPU is activated, and matrix multiplication operations are performed in parallel based on the input data in the VMEM and the weight data of the 1×1 sub-core. During matrix multiplication, the intermediate results of each parallel computation are accumulated in real time to obtain the final result.

10. The method for calculating the Conv2D operator on a TPU chip according to any one of claims 2-7, characterized in that, The hierarchical memory architecture includes HBM, VMEM, and CMEM; the dual buffering mechanism is to divide the current calculation buffer and the preload buffer in VMEM or CMEM, and asynchronously preload the next block of data when calculating the current block of data, and switch the buffer through pointer swapping; for 1×1 convolution scenarios, the output buffer is deployed in CMEM.

Images