An optimization method of a multi-layer long short-term memory network (LSTM) based on diagonal wave front parallelization

Abstract

The application discloses an optimization method of a multi-layer long short-term memory network (LSTM) based on diagonal wave front parallelization, and belongs to the technical field of deep learning neural network optimization. The method aims at the problem of limited parallelism caused by the inter-layer and inter-time step dependency relationship existing in the multi-layer LSTM network, and proposes a diagonal wave front scheduling strategy. The calculation task of the multi-layer LSTM is mapped to a two-dimensional space-time plane, and the cross-layer parallel execution is realized by identifying the calculation units without data dependency on the same diagonal line. The application adopts the persistent kernel function technology to reduce the kernel function startup overhead of the parallel computing device, and combines the fine-grained synchronization strategy to maximize the utilization of the computing resources under the premise of ensuring the calculation accuracy. Compared with the traditional layer-by-layer serial execution mode, the application can effectively improve the calculation parallelism and the utilization of the computing resources of the multi-layer LSTM, and is particularly suitable for the inference scene of the long sequence and multi-layer LSTM.

Description

Technical Field

[0001] This invention belongs to the field of deep learning neural network optimization technology, specifically involving an optimization method for a multilayer long short-term memory network (LSTM) based on diagonal wavefront parallelization. Background Technology

[0002] Long Short-Term Memory (LSTM) is a special type of Recurrent Neural Network (RNN). By introducing a gating mechanism, it effectively solves the gradient vanishing problem of traditional RNNs when processing long sequences, and has been widely used in natural language processing, speech recognition, time series prediction and other fields.

[0003] The computation process of LSTM inherently exhibits sequence dependencies: the computation at the current time step depends on the hidden state and cell state output of the previous time step. This dependency makes it difficult to parallelize a single-layer LSTM in the time dimension. In multi-layer LSTM networks, in addition to the dependencies between time steps, there are also inter-layer dependencies: the input of higher layers depends on the output of lower layers at the same time step.

[0004] Traditional multi-layer LSTM implementations typically employ a sequential execution strategy: the computation of all time steps in the first layer is completed before starting the computation of the second layer, and so on. While this approach is simple to implement, it fails to fully utilize the parallel computing capabilities of modern parallel computing devices (such as GPUs), resulting in low utilization of computing resources.

[0005] Some existing optimization methods attempt to improve the execution efficiency of LSTM through techniques such as batch matrix multiplication and operator fusion. For example, in the LSTM implementation of the NVIDIA cuDNN library, multiple gated matrix multiplications are merged into a single matrix operation to reduce the number of kernel function startups, but its essence is still layer-by-layer serial processing, which fails to effectively tap the cross-layer parallel potential in multi-layer LSTM.

[0006] Wavefront parallelism is a classic parallelization technique in high-performance computing, widely applied to computational problems with multidimensional data dependencies, such as template computation and LU decomposition. Its core idea is to identify independent computational units on the same antidiagonal line in a multidimensional dependency graph and execute them in parallel. However, applying wavefront parallelism to multilayer LSTM networks faces unique technical challenges: LSTM gating involves complex matrix multiplication and nonlinear activation operations, resulting in a computational granularity far greater than traditional template computation; inter-layer data transfer in multilayer LSTMs requires precise synchronization control; furthermore, implementing wavefront scheduling on parallel computing devices such as GPUs necessitates addressing engineering issues such as kernel function startup overhead and thread block synchronization.

[0007] Persistent kernel technology avoids the overhead of frequent kernel function startup and destruction by allowing thread blocks to reside on the computation unit throughout the entire kernel function's lifecycle.

[0008] Furthermore, GPU kernel function startup inherently incurs overhead. For long-sequence LSTM computations, frequent kernel function startups and synchronization operations can become a performance bottleneck. Reducing the number of kernel function startups while maintaining sufficient parallelism is one of the key challenges in improving LSTM computational efficiency. Patent application CN117272051A, published on December 22, 2023, discloses a time series forecasting method based on an LSTM optimized model. This method iteratively optimizes the hyperparameters of the LSTM model using a population algorithm, obtaining the globally optimal parameters based on fitness comparison, thus obtaining an optimized LSTM model for time series forecasting. This proposed method differs from the one described here, and it also does not implement multi-layer LSTM computation optimization based on diagonal wavefront parallelization. Patent application CN115034126A, published on September 9, 2022, discloses a method and system for optimizing an LSTM neural network model using the gray wolf algorithm. This method maps hyperparameters to the positions of individual gray wolves, identifies α wolves based on fitness objective function values, and employs a Levy flight strategy to guide the α wolves in a global search. The optimal hyperparameters are obtained through iterative updates using an improved gray wolf algorithm, and the LSTM model is then trained for air quality prediction. This proposed solution differs from the one described above, and it also does not implement multi-layer LSTM computation optimization based on diagonal wavefront parallelization.

[0009] Therefore, there is an urgent need for an efficient computational method that can effectively exploit the cross-layer parallelism of multi-layer LSTM, reduce kernel function startup overhead, and adaptively optimize according to the problem size. Summary of the Invention

[0010] To address the problems existing in the prior art, this invention provides an optimization method for a multilayer long short-term memory network (LSTM) based on diagonal wavefront parallelization. It achieves cross-layer parallelism through a diagonal wavefront scheduling strategy, reduces startup overhead by employing persistent kernel function technology, and combines an automatic tuning mechanism to adapt to different problem scales.

[0011] The technical solution of the present invention is as follows:

[0012] An optimization method for a multilayer long short-term memory (LSTM) network based on diagonal wavefront parallelization is characterized by the following steps:

[0013] Step 1, map the computational tasks of the multi-layer LSTM network onto a two-dimensional spatio-temporal plane with the time step as the horizontal axis and the network layer index as the vertical axis. Each grid node (y, t) represents the LSTM computational task at the t-th time step of the y-th layer, forming a computational unit matrix of Y×T, where Y is the number of network layers and T is the sequence length;

[0014] Step 2, according to the diagonal wavefront scheduling strategy, divide all computational tasks that satisfy the sum of the layer index y and the time step index t equal to the wavefront index w into the w-th diagonal wavefront, where w ∈ [0, Y+T−2]; for any two computational units (y1, t1) and (y2, t2) in the same wavefront w, y1 + t1 = y2 + t2 = w and (y1, t1) ≠ (y2, t2), from which it can be seen that y1 ≠ y2 and t1 ≠ t2; since the data dependence of LSTM only exists in the same-layer time dependence direction of (y, t) → (y, t+1) and the inter-layer dependence direction of (y, t) → (y+1, t), for two nodes in the same wavefront, if y1 < y2, then t1 > t2, so there is no predecessor-successor relationship between them in both the time direction and the layer direction, and there is no data dependence between any two computational units within the same wavefront and they can be executed in parallel;

[0015] Step 3, on the parallel computing device, in the way of starting a single kernel function, process each diagonal wavefront in ascending order of the wavefront index w. Multiple computational tasks within the same wavefront are executed in parallel by different parallel execution units; when the parallel computing device is a GPU, the parallel execution unit is a thread block. The single kernel function start adopts the persistent kernel function technology, and the number of thread blocks is set to the product of the hidden dimension shard number num_pid_h and the batch dimension shard number num_pid_b and does not exceed the maximum total number of thread blocks that the GPU can simultaneously reside in, to ensure that all thread blocks can be scheduled simultaneously and avoid persistent kernel deadlocks. When the product exceeds the maximum total number of thread blocks that the GPU can simultaneously reside in, the batch dimension is dynamically split into multiple batch chunks, and the kernel function completes the full-sample calculation by traversing all batch chunks in the outer loop, avoiding the overhead of frequently starting the kernel function;

[0016] Step 4, perform LSTM gating calculations within each computational unit (y, t), specifically including the following sub-steps:

[0017] Step 4.1, load the input data of the current time step. If y = 0, the input is the external sequence x t , calculate the pre-computed value input_data = ·x t ; if y > 0, the input is the hidden state of the previous layer , and the input weight matrix of the current layer The input gating value is calculated using block matrix multiplication; the pre-calculated value is input_data = · Simultaneously load the merged bias, which is the sum of the input bias and the loop bias.

[0018] Step 4.2, load the hidden state of the previous time step in the same layer. , as the loop input data.

[0019] Step 4.3, input the loop obtained in step 4.2. With the current layer's cyclic weight matrix The loop gate value is calculated by block matrix multiplication, added to the input gate value in step 4.1, and the bias is combined to obtain the original gate values ​​of input gate i, forget gate f, candidate memory g, and output gate o.

[0020] Step 4.4: Perform nonlinear activation on the accumulated gating value and calculate the cell state update. and hidden state output , where σ is the sigmoid activation function and tanh is the hyperbolic tangent activation function.

[0021] Step 4.5, the block matrix multiplication in steps 4.1 and 4.3 is implemented as follows: the hidden layer dimension hidden_size is divided into k_steps = hidden_size / BLOCK_SIZE_K blocks along the summation dimension; for each block k∈[0, k_steps), the k-th block is loaded from the hidden state vector and the weight matrix respectively, the block matrix multiplication is performed, and the result is accumulated in the gate value register; after k_steps iterations of accumulation, the complete matrix multiplication operation is equivalently completed, and the final gate value is obtained.

[0022] Step 5: Establish a counter-based producer-consumer fine-grained synchronization mechanism between adjacent diagonal wavefronts. When the parallel computing device is a GPU, the specific implementation is as follows: Allocate an integer synchronization counter array global_sync of size (Y+T-1)×num_pid_b in the GPU's global memory, with the initial value of each counter being zero. After the thread block of the current wavefront w completes its computation, it first executes an internal barrier within the thread block to ensure that all write operations are completed, and then increments the synchronization counter sync[w×num_pid_b+pid_b] by 1 using an atomic increment operation with release semantics. Before the next wavefront w+1 thread block begins its computation, it polls the counter using an atomic read operation with acquisition semantics until its value reaches num_pid_h, confirming that all thread blocks of the same batch of shards in the previous wavefront have completed writing. Since the computation between different batches of shards is independent of each other, the synchronization only occurs between thread blocks with the same batch shard index.

[0023] Furthermore, this invention supports weight matrix preloading optimization: when the number of blocks k_steps after the weight matrix is ​​divided does not exceed a preset threshold, for example, when k_steps≤4, that is, when the size of the weight matrix meets the condition that it can be fully loaded into the register file, the weight matrix is ​​loaded into the register variables in blocks before the wavefront loop starts. In the calculation of each time step, the weight data in the register is directly used to perform matrix multiplication, avoiding repeated access to global memory and reducing memory access latency.

[0024] Furthermore, this invention includes an automatic configuration tuning mechanism: a predefined candidate configuration set, containing block size parameters (BLOCK_SIZE_B, BLOCK_SIZE_H, BLOCK_SIZE_K), the number of thread bundles (num_warps), and the number of pipeline stages (num_stages); firstly, configurations whose block size does not match the problem dimension are filtered based on the problem size; then, configurations with rmem > 3 are excluded using the shared memory utilization formula rmem = num_stages × n_bytes × (K × M + K × N) / shared_memory (where M = 4 × BLOCK_SIZE_H, N = BLOCK_SIZE_B, K = k_steps × BLOCK_SIZE_K, and n_bytes is the number of bytes for the data type); next, configurations with computational intensity intensity = M × N × K / (M × N + K × M + K × N) not less than 1 / 4 of the maximum value are retained; finally, when a specific (batch_size, hidden_size, ...) value is encountered for the first time... When combining dtype, the optimal configuration is selected and cached through short-sequence benchmark tests, and the cached results are directly reused for subsequent problems of the same scale.

[0025] Furthermore, this invention supports mixed-precision computation: when the input data type is FP32, the kernel function internally converts the data to FP16 or BF16 for matrix multiplication and activation function calculation, and then converts the result back to FP32 for storage. For FP16 or BF16 inputs, the original precision is used directly for calculation. In mixed-precision computation, the accumulation process of the gate value is performed at higher precision (FP32), and lower precision is used only in the intermediate calculation steps of matrix multiplication. The final result is converted back to the original precision for storage, thereby improving computational throughput while controlling precision loss.

[0026] Furthermore, the present invention handles the boundary conditions as follows: when the initial hidden state and cell state are not provided, they are initialized to all-zero tensors by default; when the sequence length T=1, it contains only y wavefronts (w from 0 to y-1), each wavefront contains 1 computation unit, degenerating into sequential computation layer by layer; when the number of network layers y=1, it contains only T wavefronts, each wavefront contains 1 computation unit, degenerating into sequential computation time step by time, consistent with the behavior of a standard single-layer LSTM.

[0027] The beneficial effects of this invention are as follows:

[0028] Compared with existing technologies, this method, based on diagonal wavefront scheduling and persistent kernel design, significantly taps into the cross-layer and fine-grained parallelism of multi-layer LSTM. By eliminating the overhead of layer-by-layer startup and the bottleneck of global synchronization, and by combining block matrix multiplication to optimize memory access efficiency, it significantly improves the utilization of computing resources and the overall throughput, and has excellent versatility and practicality. Attached Figure Description

[0029] Figure 1 This is a schematic diagram of the internal structure of an LSTM cell;

[0030] Figure 2 A schematic diagram of diagonal wavefront scheduling in a multilayer LSTM;

[0031] Figure 3 Flowchart for persistent kernel function execution;

[0032] Figure 4 This is a schematic diagram of a fine-grained synchronization mechanism. Detailed Implementation

[0033] The present invention will now be described in further detail with reference to the accompanying drawings and specific embodiments.

[0034] Example 1: Diagonal wavefront parallelization of forward propagation in a multilayer LSTM network

[0035] Reference Figure 1 After the multi-layer LSTM network is unfolded along the time axis, it forms a two-dimensional dependency structure between layers and between time steps. Figure 1The LSTM includes layers 0 and 1. Each layer contains an LSTM computation unit at each time step t. The unit uses gating mechanisms such as the forget gate (σ), input gate (σ), candidate cell state generation (tanh), and output gate (σ) to combine the hidden state h from the previous time step. t-1 and cell state C t-1 And the current input x t Update the cell state C at the current time step. t and hidden state h t In the time dimension, each computational unit in the same layer depends on the hidden state h output by that layer in the previous time step. t-1 And cell state Ct-1; in the inter-layer dimension, the input of the upper-layer computational unit comes from the hidden state h output by the lower layer at the same time step. t This bidirectional dependency constitutes the core constraint of multilayer LSTM computation and is also the starting point for the diagonal wavefront parallelization strategy adopted in this invention. This embodiment provides an optimization method for multilayer long short-term memory (LSTM) networks based on diagonal wavefront parallelization.

[0036] This embodiment uses a 4-layer LSTM network to process a sequence with a time series length of 8 as an example. The parameters are configured as batch_size=32, hidden_size=128, data type is FP16, and SM_count=84 to illustrate the implementation process of diagonal wavefront parallelized forward propagation.

[0037] Reference Figure 2 The 4-layer, 8-time-step LSTM computation task is mapped to a 4x8 matrix of computational units. The horizontal axis represents the time step t ∈ {0,1,...,7}, and the vertical axis represents the layer index y ∈ {0,1,2,3}. Each grid node (y, t) represents a computational unit. Based on the rule w = y + t, the task is divided into 11 wavefronts: wavefront 0 (w=0) contains (0,0), wavefront 1 (w=1) contains (0,1) and (1,0), wavefront 2 (w=2) contains (0,2), (1,1), and (2,0), wavefront 3 (w=3) contains (0,3), (1,2), (2,1), and (3,0) (reaching a maximum parallelism of 4), wavefronts 4 through 7 each contain 4 computational units, wavefront 8 contains 3 computational units, wavefront 9 contains 2 computational units, and wavefront 10 (w=10) contains 1 computational unit, i.e., (3,7).

[0038] Data dependency analysis is performed using wavefront 4 (w=4) as an example: This wavefront contains four computational units: (0,4), (1,3), (2,2), and (3,1). Computational unit (0,4) is dependent on... and (from wavefront 3 (0,3)); computational unit (1,3) depends , (from wavefront 3, (1,2)) and (from wavefront 3 (0,3)); computational unit (2,2) depends , (from wavefront 3, (2,1)) and (from wavefront 3 (1,2)); computational unit (3,1) depends , (from wavefront 3, (3,0)) and (from (2,1) of wavefront 3). All dependencies come from wavefront 3, and there are no dependencies within the same wavefront 4, verifying that computational units within the same wavefront can be executed in parallel.

[0039] Reference Figure 3 The execution flow of the persistent kernel function is as follows: The automatic tuning module selects the optimal configuration as BLOCK_SIZE_B=8, BLOCK_SIZE_H=32, BLOCK_SIZE_K=64, num_warps=2, num_stages=4, and calculates num_pid_b=ceil(32 / 8)=4. The preprocessing stage executes on layer 0. For matrix multiplication of x, pre-compute the gate value and store it in the ifgo tensor. The ifgo tensors of layers 1-3 are initialized to zero. The biases of each layer are merged and stacked into a tensor of [4, 512]. The stacked weight matrix is ​​a tensor of [4, 512, 128].

[0040] In the main loop of the kernel function, each thread block traverses 11 wavefronts: when w > 0, it polls the synchronization counter of the previous wavefront through atomic read operations, waiting for its value to reach 4 (i.e., the value of num_pid_h); it identifies the computational units (layer, ss) to be processed in the current wavefront (satisfying layer + ss = w); it loads the gated pre-computation values, and when layer > 0, it loads the hidden state of the previous layer and performs block matrix multiplication (k_steps = ceil(128 / 64) = 2), loads the hidden state of the previous time step in the same layer and performs block matrix multiplication, accumulates it, and then executes the activation function to update the cell state and hidden state; after executing the barrier within the thread block, it releases the semantically incrementing synchronization counter of the current wavefront. After the kernel function finishes execution, it extracts the output sequence and final state from the last layer.

[0041] Example 2: Implementation of a fine-grained synchronization mechanism

[0042] Reference Figure 4 The synchronization counter array `global_sync` has a dimension of 11 × 4 = 44 (wavefront number 11 × num_pid_b = 4), and all initial values ​​are 0. For example... Figure 4As shown in the upper part, all hidden dimension thread blocks (pid_h=0, 1, ..., N) responsible for the same batch of slices pid_b in the same wavefront w are computed in parallel. After each thread block is completed, the hidden state and cell state are written to the hidden state data area of ​​global memory, and then the corresponding counter Sync[w×num_pid_b+pid_b] is incremented by atomic addition operation with release semantics. Figure 4 As shown in the lower half, before the next wave, the thread block w+1 obtains the semantic reading counter before computation. If its value does not reach num_pid_h, it polls and checks the counter status until the counter value reaches num_pid_h (indicating that all hidden dimension fragments of the same batch of slices in the previous wave have been written), before it can enter the computation of the next wave. Taking the thread block (pid_b=1, pid_h=2) as an example: after wave front 0 finishes the computation of (0,0), it increments global_sync[0×4+1]=global_sync[1]. Other hidden dimension thread blocks of the same batch of slices also increment the counter. Before wave front 1 starts, it polls and checks the status of global_sync[1] until its value reaches 4 (equal to num_pid_h), confirming that all hidden dimension fragments of pid_b=1 in wave front 0 have been completed. Then it executes the corresponding computation unit and increments global_sync[1×4+1]=global_sync[5]. Subsequent waves follow the same pattern until all waves are completed. Thread blocks with different pid_b use independent counters, do not need to be synchronized with each other, and can be executed in complete parallel.

[0043] Example 3: Selection process for automatic tuning configuration

[0044] For the issues with batch_size=32, hidden_size=128, and dtype=fp16, the automatic tuning module filters configurations using the following steps: It generates 2304 candidate configurations (6×6×2×4×8); initially filters out configurations with BLOCK_SIZE_H≥256 and BLOCK_SIZE_B≥64, retaining approximately 800; calculates shared memory usage using the formula, excluding configurations with rmem>3, retaining approximately 400; calculates the computational intensity (intensity) value for all remaining configurations, identifies the maximum value (intensity_max), and excludes inefficient configurations with intensity < intensity_max / 4. Computational intensity reflects the computational load per unit of data access and is a key indicator of configuration efficiency. Insufficient intensity means that memory overhead is too high relative to computational load, failing to fully utilize GPU computing power. Using 1 / 4 of the maximum intensity as a threshold, while ensuring the number of candidate configurations, obviously inefficient configurations are eliminated, resulting in approximately 200 configurations. Based on tile_size=4×32×8=1024, configurations with num_warps=2 and num_stages∈{4,5} are matched, and approximately 30 are retained. Short sequences with seq_len=6 are used for benchmark testing, and the configuration with the shortest execution time (BLOCK_SIZE_B=8, BLOCK_SIZE_H=32, BLOCK_SIZE_K=64, num_warps=2, num_stages=4) is selected and cached for reuse in subsequent problems of the same size.

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

Claims

1. An optimization method for a multilayer long short-term memory (LSTM) network based on diagonal wavefront parallelization, characterized in that, It includes the following steps: Step 1: Map the computational tasks of the multi-layer LSTM network onto a two-dimensional spatio-temporal plane with the time step as the horizontal axis and the network layer index as the vertical axis. Each grid node (y, t) represents the LSTM computational task at the t-th time step of the y-th layer, forming a computational unit matrix of Y×T, where Y is the number of network layers and T is the sequence length; Step 2: According to the diagonal wavefront scheduling strategy, divide all computational tasks that satisfy the sum of the layer index y and the time step index t equal to the wavefront index w into the w-th diagonal wavefront, where w∈[0, Y+T−2]; For any two computational units (y1, t1) and (y2, t2) in the same wavefront w, it satisfies y1+t1=y2+t2=w and (y1, t1)≠(y2, t2), from which it can be seen that y1≠y2 and t1≠t2; Since the data dependence of LSTM only exists in the same-layer time dependence direction of (y, t)→(y, t+1) and the inter-layer dependence direction of (y, t)→(y+1, t), for two nodes in the same wavefront, if y1<y2, then t1>t2. Therefore, there is no predecessor-successor relationship between them in both the time direction and the layer direction, and there is no data dependence between any two computational units within the same wavefront and they can be executed in parallel; Step 3: On the parallel computing device, in the way of starting with a single kernel function, process each diagonal wavefront in ascending order of the wavefront index w. Multiple computational tasks within the same wavefront are executed in parallel by different parallel execution units; When the parallel computing device is a GPU, the parallel execution unit is a thread block. The single kernel function startup adopts the persistent kernel function technology, and the number of thread blocks is set to the product of the hidden dimension shard number num_pid_h and the batch dimension shard number num_pid_b and does not exceed the maximum total number of thread blocks that the GPU can resident simultaneously, to ensure that all thread blocks can be scheduled simultaneously and avoid persistent kernel deadlocks. When the product exceeds the maximum total number of thread blocks that the GPU can resident simultaneously, the batch dimension is dynamically split into multiple batch blocks batch_chunks, and the kernel function completes the full-sample calculation by traversing all batch blocks in the outer loop, avoiding the overhead of frequently starting the kernel function; Step 4: Perform LSTM gating calculations within each computational unit (y, t), specifically including the following sub-steps: Step 4.1: Load the input data for the current time step. If y=0, the input is the external sequence x. t Calculate the pre-computed value input_data = ·x t If y > 0, the input is the hidden state of the previous layer. , and the current layer input weight matrix The input gating value is calculated using block matrix multiplication; the pre-calculated value is input_data = · Simultaneously load the merged bias, which is the sum of the input bias and the loop bias. Step 4.2, load the hidden state of the previous time step in the same layer. , as the loop input data. Step 4.3, input the loop obtained in step 4.

2. With the current layer's cyclic weight matrix The loop gate value is calculated by block matrix multiplication, added to the input gate value in step 4.1, and the bias is combined to obtain the original gate values ​​of input gate i, forget gate f, candidate memory g, and output gate o. Step 4.4: Perform nonlinear activation on the accumulated gating value and calculate the cell state update. and hidden state output , where σ is the sigmoid activation function and tanh is the hyperbolic tangent activation function. Step 4.5: The block matrix multiplication in Step 4.1 and Step 4.3 is implemented as follows: Divide the hidden layer dimension hidden_size into k_steps = hidden_size / BLOCK_SIZE_K blocks along the summation dimension according to the block size BLOCK_SIZE_K; For each block k∈[0, k_steps), load the k-th block from the hidden state vector and the weight matrix respectively, perform block matrix multiplication and accumulate the results into the gating value register; After k_steps iterations of accumulation, equivalently complete the complete matrix multiplication operation to obtain the final gating value. Step 5: Establish a counter-based producer-consumer fine-grained synchronization mechanism between adjacent diagonal wavefronts. When the parallel computing device is a GPU, the specific implementation is as follows: Allocate an integer synchronization counter array global_sync of size (Y+T-1)×num_pid_b in the GPU's global memory, with the initial value of each counter being zero. After the thread block of the current wavefront w completes its computation, it first executes an internal barrier within the thread block to ensure that all write operations are completed, and then increments the synchronization counter sync[w×num_pid_b+pid_b] by 1 using an atomic increment operation with release semantics. Before the next wavefront w+1 thread block begins its computation, it polls the counter using an atomic read operation with acquisition semantics until its value reaches num_pid_h, confirming that all thread blocks of the same batch of shards in the previous wavefront have completed writing. Since the computation between different batches of shards is independent of each other, the synchronization only occurs between thread blocks with the same batch shard index.

Images