A general top-k computing device and method based on insertion-merge

Abstract

The application provides a general TopK calculation device and method based on insertion merging, relates to the technical field of integrated circuit hardware design, and comprises the following steps: receiving a TopK calculation instruction; generating a control signal sequence for controlling an operation module to perform TopK calculation based on the TopK calculation instruction; determining an operation mode to be entered based on the input control signal sequence; the operation mode comprises a pure merging mode, a pure insertion mode or an insertion merging mode; and based on the parallel calculation unit, the operation is performed on the to-be-processed data in the operation mode, and the calculation result corresponding to the TopK calculation instruction is output. The application reuses the parallel comparison unit in the existing vector calculation unit, can not only output the selected value but also output the corresponding index, the output result is stable, the number of clock cycles required by the operator is small, and therefore the support capability and the calculation efficiency for mixed expert models and various sparse strategies are effectively improved.

Description

Technical Field

[0001] This application relates to the field of integrated circuit hardware design technology, and in particular to a general TopK computing device and method based on insertion merging. Background Technology

[0002] With the rapid development of deep learning, especially large-scale language models and hybrid expert models, the scale and complexity of neural network models have increased dramatically. To maintain powerful model capabilities while controlling computational costs, hybrid expert models and various sparsity strategies (such as sparse attention and token filtering) have become important technical directions. In these architectures, TopK (selecting the K largest or smallest elements from N elements) is widely used as a core operator: for example, in the routing module of a hybrid expert model, it is necessary to select the activated Top-K experts for each input token; in sparsity strategies, it is necessary to filter out the most important Top-K tokens from a long sequence. These application scenarios typically require processing large-scale inputs (N can reach tens of thousands or even hundreds of thousands) and simultaneously obtaining the values ​​of the selected elements and their indices in the original sequence. Therefore, designing a hardware unit capable of efficiently, flexibly, and with low resource overhead processing large-scale variable-parameter TopK computations has become one of the key challenges in improving the performance of related artificial intelligence computing systems.

[0003] For the hardware implementation of TopK computing, existing technologies mainly revolve around several technical paths. For example, some schemes employ a sorting structure based on comparison-swapping networks, using multi-level fixed comparison networks to fully or partially sort the input data to achieve Top-K selection. While these methods offer high parallelism, the hardware scale increases significantly with the problem size and is difficult to flexibly adapt to changing N and K parameters. Other schemes use a structure based on streaming caching and merging, dynamically maintaining a candidate set during the data stream input process through multi-level FIFO queues or register files. This method can handle continuous data streams, but when K is large, it requires a considerable depth or number of cache units, leading to excessive on-chip storage resource consumption. In addition, there are threshold search-based schemes that filter elements by iteratively approximating the Kth largest / smallest threshold without requiring complete sorting. However, this method typically requires multiple rounds of iteration and global statistics, with order dependencies between rounds, complex control, and potentially increased hardware overhead for comparison and counting on large-scale data. Furthermore, the inability of threshold search schemes to scale to support the core sorting operator in top-P sampling also limits their application to some extent.

[0004] It is evident that existing TopK hardware implementations share common limitations: poor flexibility, high hardware cost, and large storage overhead. They typically struggle to achieve a good balance across multiple dimensions, including supporting large-scale variable inputs (N), variable number of outputs (K), simultaneous output of values ​​and indices, and low hardware overhead and low latency. Summary of the Invention

[0005] This application provides a general TopK computing device and method based on insertion merging to solve the problem that existing computing devices cannot support large-scale variable TopK tasks, resulting in high hardware costs.

[0006] In a first aspect, this application provides a general TopK computing device based on insertion merging, comprising: The instruction module is configured to receive TopK calculation instructions; the TopK calculation instructions include the total number of elements, the number of elements to be selected, and an identifier bit. The control module is configured to generate a sequence of control signals for controlling the computation module to perform TopK calculations based on the TopK calculation instructions. The computing module includes multiple parallel computing units, and the computing module is configured as follows: Based on the input control signal sequence, the operation mode to be entered is determined; the operation mode includes pure merge mode, pure insertion mode, or insertion merge mode. In the operation mode, the parallel computing unit performs calculations on the data to be processed loaded from the memory and outputs the calculation result corresponding to the TopK calculation instruction; the data to be processed includes input values ​​and indices, where the indices are either continuous indices generated based on the first input or input indices based on non-first inputs.

[0007] In some feasible embodiments, the computing unit has P paths, and the P-path computing units are connected sequentially along the (P-1)th path to the 0th path to form a data path; the register data of the computing unit in the (i+1)th path is used as the input shift data of the computing unit in the i-th path; where i∈[0, p-2]; The computing unit includes: The system comprises an output register group, a first input register group, and a second input register group; the output register group is used to store the output value and the output index; the first input register group is used to store the first input value and the first input index; and the second input register group is used to store the second input value and the second input index. The comparator is used to compare the output value, the first input value, and the second input value based on the comparison operator, and output a single-bit comparison result.

[0008] In some feasible embodiments, the computing module is further configured as follows: The single-bit comparison results output by the P-way computing unit are concatenated to generate a shift control signal; the shift control signal is used to control the output value and output index received by the output register group.

[0009] In some feasible embodiments, before the computing module outputs the calculation result corresponding to the TopK calculation instruction, it is specifically configured as follows: Parameter initialization: Initialize the total number of elements to N0 and the sorted granularity to 0. When the total number of elements is greater than a preset value, the data to be processed is divided into several data blocks; the total number of elements in each data block is less than or equal to the preset value. Calculate the single TopK calculation result for each of the data blocks; the single TopK calculation result includes the selected value and index; The results of a single TopK calculation of all the data blocks are concatenated to update the current total number of elements and the sorted granularity; When the updated total number of current elements is less than or equal to a preset value, the single TopK calculation result corresponding to the updated total number of current elements is used as the calculation result corresponding to the TopK calculation instruction.

[0010] In some feasible embodiments, the computing module determines the computing mode to be entered, specifically configured as follows: If the sorted granularity is 0, the data to be processed is preprocessed. After preprocessing, if the number of elements to be selected is less than or equal to the number of computing units, then the system is determined to enter the pure insertion mode; if the number of elements to be selected is greater than the number of computing units, then the system is determined to enter the insertion and merging mode. If the sorted granularity is not zero, and the number of elements to be selected is less than or equal to the number of computational units, then the system is determined to enter the pure insertion mode; if the number of elements to be selected is greater than the number of computational units, then the system is determined to enter the pure merge mode.

[0011] In some feasible embodiments, the TopK calculation instruction further includes a local data identifier bit; The computing module preprocesses the data to be processed, specifically configured as follows: The input values ​​in the data to be processed are padded to an integer multiple of the number of computing units; Generate the sequential index corresponding to the input value based on the local data identifier bit.

[0012] In some feasible embodiments, the calculation module outputs the calculation result in the pure insertion mode, specifically configured as follows: Initialize the output register group; Fill the first input register group with P values ​​and their indices from the input values; The comparison result between the input value in the first input register group and the current value in the output register group is obtained based on the comparator; The output register group is updated based on the comparison results until the comparison results corresponding to all input values ​​are updated to the output register group. The calculation result is obtained by symmetrically flipping the values ​​and indices in the output register group.

[0013] In some feasible embodiments, the computation module outputs the calculation result in the pure merge mode, specifically configured as follows: Initialize the current target sorting granularity and the current sorting stage; Fill the first P values ​​and indices of the two consecutive segmented ordered values ​​in the input values ​​into the first input register group and the second input register group, respectively. The first element values ​​of the two segments are compared, the winner is moved to the end of the output register group, the corresponding remaining elements are shifted to the right, and the original elements of the output register group are shifted to the right. The calculation result is output when the target is reached in the current sorting stage and the number of output elements in the output register group reaches the number of elements to be selected.

[0014] In some feasible embodiments, the calculation module outputs the calculation result in the insertion and merge mode, specifically configured as follows: The insertion phase is performed; the insertion phase is the same as the steps performed by the operation module in the pure insertion mode, wherein, after the comparison results corresponding to all input values ​​are updated to the output register group, the data in the output register group is written to the memory; The merge phase is performed; the merge phase includes: Initialize the current target sorting granularity and the current sorting stage; Fill the first P values ​​and indices of the two consecutive segmented ordered values ​​in the input values ​​into the first input register group and the second input register group, respectively. The first element values ​​of the two segments are compared, the winner is moved to the end of the output register group, the corresponding remaining elements are shifted to the right, and the original elements of the output register group are shifted to the right. When the target is reached in the current sorting stage and the number of output elements in the output register group reaches the number of elements to be selected, the calculation result is output. Achieving the objective in the current sorting stage includes: All groups of the data to be processed have been merged, and the current target sorting stage is incremented by 1. Alternatively, the target sorting granularity is multiplied by 2, and the current target sorting stage is incremented by 1. Then, the step of filling the first P values ​​and indices of the two consecutive segments of ordered values ​​in the input values ​​into the first input register group and the second input register group is executed again.

[0015] In some feasible embodiments, the number m of elements in the last group of data to be merged in the current merging phase is determined. i The target sorting granularity g at the current stage i and the target sorting granularity g of the previous stage i-1 The relationship between them; If g i-1 <m i <2g i-1 =g i Then when the number of output values ​​and indices in the current group reaches m i The current group is considered merged when the current merge is complete. If g i =g i-1 <m i <2g i-1 Then the total number of unread elements remaining in the second input register group is m. i '-g i-1 ; If 0 <m i '≤g i-1 If the merging process is skipped, all elements in the corresponding segment sequence will be stored in memory.

[0016] Secondly, embodiments of this application provide a general TopK calculation method based on insertion merging, applied to the apparatus of the first aspect, the method comprising: Receive TopK calculation instructions; TopK calculation instructions include the total number of elements, the number of elements to be selected, and flag bits; Based on the TopK calculation instructions, a sequence of control signals is generated to control the computation module to perform TopK calculations. Based on the input control signal sequence, the operation mode to be entered is determined; the operation mode includes pure merge mode, pure insertion mode, or insertion-merge mode. In the operation mode, parallel computing units perform calculations on the data to be processed loaded from memory and output the calculation results corresponding to the TopK calculation instructions. The data to be processed includes input values ​​and indices, where the indices are either continuous indices generated based on the first input or input indices based on subsequent inputs.

[0017] The apparatus and method provided in this application have the following beneficial effects: This application can reuse the parallel comparison unit in the existing vector computing unit. The TopK-supporting computing unit also supports other types of vector computing and flexibly supports variable N, K, and largest parameters. The circuit structure provided in the embodiments of this application is not for a fixed N, K, or largest. It can output not only the selected value but also its corresponding index. The output result is stable, that is, when the values ​​are equal, the smaller index is output first. The number of clock cycles required by the operator is small, thereby effectively improving the support capability and computational efficiency for hybrid expert models and various sparse strategies. Attached Figure Description

[0018] To more clearly illustrate the technical solution of this application, the drawings used in the embodiments will be briefly introduced below. Obviously, for those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0019] Figure 1 Architecture diagram of a general TopK computing device based on insertion merging provided in this application embodiment; Figure 2 A control principle diagram of a general TopK computing device based on insertion merging provided in this application embodiment; Figure 3 A circuit diagram of the computing unit in a general TopK computing device based on insertion merging provided in the embodiments of this application; Figure 4 A flowchart illustrating the TopK operation performed by a general TopK computing device based on insertion merging, provided in an embodiment of this application; Figure 5 A flowchart illustrating the execution of a single TopK operation by a general TopK computing device based on insertion merging, as provided in an embodiment of this application; Figure 6 A flowchart of the pure insertion mode algorithm provided in the embodiments of this application; Figure 7 Example diagram of a single TopK process of a general TopK computing device provided in the embodiments of this application in pure insertion mode; Figure 8 A flowchart illustrating the algorithm for the insertion and merging pattern provided in an embodiment of this application; Figure 9 An example diagram of a single TopK process in insertion-merge mode of a general TopK computing device provided in an embodiment of this application; Figure 10 A flowchart of the pure merge mode algorithm provided in the embodiments of this application; Figure 11 An example diagram of a single TopK process in pure merge mode of a general TopK computing device provided in the embodiments of this application; Figure 12 Example diagrams illustrating the process of a general TopK computing device, provided in some embodiments of this application, calculating a TopK operation consisting of multiple single TopK operations; Figure 13 Example diagrams illustrating the process of a general TopK computing device for calculating a TopK operation consisting of multiple single TopK operations, provided in other embodiments of this application. Detailed Implementation

[0020] To enable those skilled in the art to better understand the technical solutions in this application, the technical solutions in the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments in this application, all other embodiments obtained by those of ordinary skill in the art without creative effort should fall within the scope of protection of this application.

[0021] With the rapid development of large-scale models, some models adopt a hybrid expert model architecture to pursue stronger model capabilities while avoiding excessively high computational costs. Expert routing is a key component, typically used to select a certain number of experts for each token; TopK is the core operator in this architecture. On the other hand, many sparse strategies involve the TopK operator, which selects a certain number of tokens from a sequence of length N tokens. Since the sequence length can reach the maximum sequence length of the model, the parameter N in TopK (selecting the largest or smallest K elements from N elements) can reach the order of 100,000, placing high demands on the hardware implementation of the TopK operator. Furthermore, in these application scenarios, the most needed resource is often the index returned by TopK; therefore, how to return the indices of the largest or smallest K elements while performing TopK is also a challenge. In addition, avoiding massive hardware resource overhead and excessive clock cycles is another significant challenge for the TopK computation unit.

[0022] To address the aforementioned issues, this application provides a general TopK computing device based on insertion-merge. By efficiently reusing the parallel comparison unit in the existing vector computing unit, a TopK computing unit based on "insertion-merge" is proposed. It can flexibly support variable parameters such as N, K, and largest, requires fewer clock cycles for the operator, and can output the selected value and index, thus effectively improving the support capability and computational efficiency for hybrid expert models and various sparse strategies.

[0023] See Figure 1 This is an architecture diagram of a general TopK computing device based on insertion merging provided in the embodiments of this application.

[0024] like Figure 1 As shown, the general TopK computing device based on insertion merging provided in this application includes: The instruction module is configured to receive TopK calculation instructions; the TopK calculation instructions include the total number of elements, the number of elements to be selected, and a flag bit.

[0025] The instruction module can be configured with instruction decoding logic, which is responsible for parsing the instruction word into the control signals and operation parameters required by TopK.

[0026] The control module is configured to generate a sequence of control signals for controlling the computation module to perform TopK calculations based on TopK calculation instructions.

[0027] The control module can be configured with TopK control logic, which mainly includes the TopK state machine and the signals driven by this state machine.

[0028] The computation module, including multi-path parallel computing units, is configured as follows: Based on the input control signal sequence, the operation mode to be entered is determined; the operation mode includes pure merge mode, pure insertion mode, or insertion-merge mode; the operation module can perform a variety of different operations based on different configurations.

[0029] In any operation mode, the operation module can perform calculations on the data to be processed loaded from memory based on parallel computing units and output the calculation results corresponding to the TopK calculation instructions; wherein, the data to be processed includes input values ​​and indices, and the indices are either continuous indices generated based on the first input or input indices based on non-first inputs.

[0030] See Figure 2 The diagram shows the control principle of a general TopK computing device based on insertion merging provided in this application embodiment; see also Figure 3 This is a circuit diagram of the computing unit in the general TopK computing device based on insertion merging provided in the embodiments of this application. It should be noted that the accompanying drawings of this application only show the circuit structure related to the TopK function, and other unrelated circuit structures are not shown.

[0031] Depend on Figure 2 and Figure 3 As shown, in some embodiments, the computing module includes P parallel computing units, wherein the P computing units can be configured to... Figure 2 As shown, the data path is formed by sequentially connecting the P-1th path to the 0th path; the register data of the calculation unit of the i+1th path is used as the input shift data of the calculation unit of the ith path; where i∈[0, p-2].

[0032] TopK primarily utilizes the fp32 comparator and six 32-bit registers (the remaining circuitry is not shown). Two 32-bit registers form the output register group, used to store the output value and output index; two 32-bit registers form the first input register group (A), and two 32-bit registers form the second input register group (B), both used to store the input value and input index.

[0033] For example, in the TopK operation, the six 32-bit registers in the calculation unit are used to store the output value (fp32), the output index (uint32), the input value A (fp32), the input index A (uint32), the input value B (fp32), and the input index B (uint32), respectively.

[0034] The input signals received by the computing unit mainly include: 1. left_0~left_5: This is the register data of the leftmost calculation unit of the current calculation unit, used for data shifting.

[0035] 2. bc_0, bc_1: These are the input values ​​from calculation unit 0, and the data from index register A or B.

[0036] 3.new: This is new data from SRAM.

[0037] 4. Shift_ctrl_0~Shift_ctrl_2: Control the MUX at the input terminals of: output value and index register; input value and index register A; and input value and index register B, respectively, controlling whether it receives left_x, bc_x, or new (x refers to any number). Specifically: (1) When shift_ctrl_0=2'b00, the output value and the index register receive the original value; when shift_ctrl_0=2'b01, the output value and the index register receive bc_0 and bc_1; when shift_ctrl_0=2'b11, the output value and the index register receive left_0 and left_1.

[0038] (2) When shift_ctrl_1=1'b0, the input value is received and the index register A receives the original value; when shift_ctrl_1=1'b1, the input value is received and the index register A receives left_2 and left_3.

[0039] When shift_ctrl_2=1'b0, the input value is received and the index register B receives the original value; when shift_ctrl_2=1'b1, the input value is received and the index register B receives left_4 and left_5.

[0040] In the embodiments of this application, see Figure 3 The FP32 comparator can interpret and compare inputs op1[31:0] and op2[31:0] according to the IEEE 754 single-precision floating-point method. Based on the comparison operators (greater than / greater than or equal to / less than / less than or equal to, given by the control logic), it determines the winner. If op1 wins, then op3 is the winner; if op2 wins, then op4 is the winner. Figure 3 In the diagram, comp_result is actually the least significant bit (single-bit signal) of winner when op3 and op4 are 32'd1 and 32'd0 respectively. It indicates whether op1 is greater than / greater than or equal to / less than / less than or equal to op2 under the current comparison operator (greater than / greater than or equal to / less than / less than or equal to). For ease of understanding, it is distinguished from winner[31:0].

[0041] After the comparator performs the comparison operation, the output signal of the computing unit includes: The winner and comp_result provided by the fp32 comparator, as well as reg_0~reg_5: are the data in the six 32-bit registers in this calculation unit.

[0042] See also Figure 2 In a P-way parallel computing unit, between two adjacent computing units (the i-th and i+1-th paths), the outputs reg_0~reg_5 of the i+1-th path can be used as inputs left_0~left_5 to the i-th computing unit, forming a data right-shift path. The left_0~left_5 of the P-th computing unit are driven by either 32'hFFFF_FFFF or 32'h7FFF_FFFF, depending on the largest value. The winner between the outputs reg_2 and reg_4 of the 0-th computing unit is used as bc_0, which can optionally be input to the output value registers of all P computing units; the winner between reg_3 and reg_5 is used as bc_1, which can optionally be input to the output index registers of all P computing units.

[0043] After each computational unit completes its comparison operation, its output single-bit comparison result, `comp_result`, can be concatenated, with 1'b0 appended to the most significant bit as a shift control signal (P+1 bits in total). This shift control signal is used to control the signals (output values ​​and indices) received by the output register groups in the P computational units. Specifically, this P+1 bit signal can be split into P groups of adjacent 2-bit signals, each overlapping by 1 bit, and used as the control signal `shift_ctrl_0` for the P computational units.

[0044] In some embodiments, during the TopK operation, the output of the computation module comes from the output (value / index) register group of the P-way computation unit. Because only a unidirectional right shift path is introduced between computation units, the output element arrangement direction is reversed between the pure insertion mode and other modes (in the pure insertion mode without symmetrical flipping, the K largest elements when largest=1 are arranged in ascending order, and the K smallest elements when largest=0 are arranged in descending order; in other modes, the K largest elements when largest=1 are arranged in descending order, and the K smallest elements when largest=0 are arranged in ascending order). To unify the output order (unified as the K largest elements when largest=1 are arranged in descending order; and the K smallest elements when largest=0 are arranged in ascending order), a symmetrical flipping operation needs to be performed in the pure insertion mode.

[0045] See Figure 4 This is a flowchart illustrating the TopK operation performed by a general TopK computing device based on insertion merging, as provided in an embodiment of this application.

[0046] like Figure 4 As shown, the TopK operation consists of one or more "single TopK" operations and their accompanying load / store operations. Here, N is the total number of TopK elements, K is the number of elements to be selected, largest is the size flag, local is the local data flag for a single TopK operation, and G is the sorted granularity of a single TopK operation.

[0047] Accordingly, before the computation module outputs the calculation result corresponding to the TopK calculation instruction, it is specifically configured as follows: S11: Parameter initialization, initialize the total number of current elements to N0, and initialize the sorted granularity to 0 (indicating unordered).

[0048] S12: When the total number of elements is greater than a preset value, the data to be processed is divided into several data blocks; the total number of elements in each data block is less than or equal to the preset value.

[0049] If N0 ≤ N_limit, only one TopK operation is needed. However, if N0 is greater than N_limit, the input value is divided into blocks, resulting in blocks with a total number of elements N0' = N_limit and several remaining blocks with less than N_limit (N0' < N_limit). Each of these remaining blocks undergoes a single TopK operation. Note that the parameters N = N0' (N' is the total number of elements in the block) and K = K' (K' is the smallest power of 2 not less than K) for this single TopK operation. If the remaining block's N' ≤ G, its TopK operation is skipped, and the corresponding data is directly moved to the target address.

[0050] In this embodiment of the application, the total number of elements that can be processed in a single TopK iteration is limited by the SRAM size (e.g., 1024). The upper limit of N in a 1024 SRAM word TopK is 8192 (called the upper limit of N in a single TopK). Therefore, when the total number of TopK elements N is greater than the upper limit of N in a single TopK (denoted as N_limit), it is necessary to perform block-based TopK.

[0051] It should be noted that, for the mathematical correctness of the TopK operation, the TopK operation N in this embodiment can support a maximum of 2. 32 (Due to register bit width limitations), when N≤N_limit, it can support any 1≤K≤N, while when N>N_limit, it can support any 1≤K≤(N_limit / 2). (In MoE, the K required for routing is generally below 16, while the K required for token sparsity in sparse strategies is generally 2048 or below. For example, N_limit=8192 of 1024×1024 SRAM can cover almost all scenarios).

[0052] S13: Calculate the single TopK calculation result for each data block; the single TopK calculation result includes the selected value and index.

[0053] S14: Concatenate the results of a single TopK calculation for all data blocks to update the current total number of elements and the sorted granularity. Specifically, concatenate the elements (value pairs and index pairs) selected from all blocks, update parameter N0 to the current total number of elements, and update parameter G to K'.

[0054] S15: When the updated total number of current elements is less than or equal to the preset value, the single TopK calculation result corresponding to the updated total number of current elements will be used as the calculation result corresponding to the TopK calculation instruction.

[0055] After repeating steps S12-S14 until N0 < N_limit, the final calculation result can be obtained. It should be noted that in the first round of calculation, the input of a single TopK iteration only contains the value and no index, while in all subsequent rounds, the input of a single TopK iteration includes both the value and the index.

[0056] See Figure 5 This is a flowchart illustrating a single TopK operation performed by a general TopK computing device based on insertion merging, as provided in an embodiment of this application.

[0057] In this embodiment, a single TopK operation can support any N and K where 1 ≤ K ≤ N ≤ N_limit, where N_limit is the upper limit of N for a single TopK operation. A single TopK operation can be controlled by a single instruction, which includes: the total number of elements N, the number of elements to be selected K, the size flag (largest), the sorted granularity G, the local data flag (local), and the SRAM starting address (sram_addr). This embodiment primarily targets IEEE 754 single-precision floating-point numbers.

[0058] like Figure 5 As shown, in a single TopK operation, when the computation module determines the computation mode to be entered, it is specifically configured as follows: S21: Preprocess the data to be processed when the sorted granularity is 0; S22: After preprocessing, if the number of elements to be selected is less than or equal to the number of computation units, then it is determined to enter the pure insertion mode; if the number of elements to be selected is greater than the number of computation units, then it is determined to enter the insertion merge mode. S23: If the sorted granularity is not 0, and the number of elements to be selected is less than or equal to the number of computation units, then it is determined to enter the pure insertion mode; if the number of elements to be selected is greater than the number of computation units, then it is determined to enter the pure merge mode.

[0059] Specifically, when the sorted granularity (0 when the input values ​​are unordered; m when the input values ​​are ordered in segments of 2m elements) is 0, the input only contains the original values. In this case, depending on whether K is greater than P, it enters either "pure insertion" mode or "insert-merge" mode.

[0060] When the sorted granularity is not 0, the input includes the sorted values ​​and their indices in the original sequence. In this case, depending on whether K is greater than P, it enters either "pure insertion" mode or "pure merge" mode.

[0061] In some embodiments, preprocessing can be performed on the data before executing a single TopK operation on the data to be processed. Specifically, when the TopK calculation instruction includes a local data flag, the computation module preprocesses the data to be processed, specifically configured as follows: S31: Fill the input values ​​in the data to be processed to an integer multiple of the number of computing units; S32: When the local data identifier indicates that the current input is raw local data, a continuous index corresponding to the original input sequence is generated based on the starting position of the data block and the position of the element within the data block; If the data to be processed already contains an index, use that index as the input index.

[0062] Specifically, let N be the length of the TopK input sequence in a single iteration, K be the length of the target sequence, N' be the length of the input sequence after preprocessing and padding to a multiple of P, and K' be the length of the target sequence after padding to a multiple of P. When largest=1, the final output is the K' largest values ​​among the N input values ​​and their indices in the original sequence; when largest=0, the final output is the K' smallest values ​​among the N input values ​​and their indices in the original sequence.

[0063] The following is a detailed description of the three different operation modes proposed in the embodiments of this application.

[0064] When entering pure insertion mode, the operation of the device in this embodiment is as follows: When K is less than or equal to P, the system enters pure insertion mode. In this mode, all input values ​​and index registers used are input values ​​and index register A, while register B is not used. The algorithm flow in this mode is as follows: Figure 6 As shown. The details of each step are as follows: 1. Input: The input consists of preprocessed numerical values ​​and indices.

[0065] 2. Initialize the values ​​and indices in the output registers. Specifically, when larger=1, fill P 32-bit output value registers with 32'hFF80_0001; when larger=0, fill all output value registers with 32'h7F80_0001. Fill P 32-bit output index registers with 32'h0000_0000.

[0066] 3. Fill a new batch of P values ​​and indices into the input registers. Specifically, P input values ​​are filled into P 32-bit input value registers, and P input indices are filled into P 32-bit input index registers.

[0067] 4. The current value is compared in parallel with the P values ​​in the output register. Specifically, the data in the 0th input value register is compared in parallel with the P data in the output value register to obtain a comparison result of P bits. When larger=1, the comparison operator is greater than; when larger=0, the comparison operator is less than.

[0068] 5. Insert the current value and index into the correct positions of the P values ​​and indices in the output register, and shift the remaining elements to the right. That is, based on the comparison results above, insert the 0th input value and the data in the index register into the correct positions in the output value and index register, respectively.

[0069] For example, when P=4, if the comparison result is 4'b0011, it will be extended from 1'b0 to 5'b00011 as the output shift control signal: For the 0th output value and index register, bits [1:0] are taken, i.e., 2'b11, thus receiving data from the 1st output value and index register; for the 1st output value and index register, bits [2:1] are taken, i.e., 2'b01, thus receiving data from the 0th input value and index register; for the 2nd output value and index register, bits [3:2] are taken, i.e., 2'b00, thus maintaining the original value; for the 3rd output value and index register, bits [4:3] are taken, i.e., 2'b00, thus maintaining the original value. This achieves the goal of inserting the current value and index into the correct position in the output value and index register. Except for the last one, all input values ​​and index registers receive data from the previous input value and index register. The last input value and index register receives 32'hFFFF_FFFF or 32'h7FFF_FFFF depending on the largest parameter.

[0070] 6. As long as there are still valid elements in the input value and index register, repeat the above operation until the current batch of P elements has been processed.

[0071] 7. As long as there are still batches that have not been processed, the above operation must be repeated for the next batch until all batches have been processed.

[0072] 8. Post-processing: Let's assume `largest` = 1. After the above process is completed, the data in the output value registers 0~P-1 consists of the K largest elements out of N elements arranged in ascending order. To maintain descending order, a symmetrical flip is required. That is, the i-th output value and the index register are output as the P-1-i-th data. The fundamental reason for this operation is that only a unidirectional data transfer channel is introduced between adjacent registers. Therefore, either the pure insertion mode or one of the other two modes must introduce this post-processing operation to ensure the consistency of the ascending and descending order of the output data.

[0073] like Figure 7 The image shows an example of a single TopK process in pure insertion mode, where the comparator parallelism P=32, the total number of elements N=476, the number of elements to be selected K=29, and the size flag largest=1. Note that each line contains P=32 elements (value-index pairs), and the output element count is a multiple of P, resulting in a target sequence length of K'=32.

[0074] When entering the insertion and merging mode, the operation process of the device in this embodiment of the application is as follows: When the sorted granularity G is 0 and K is greater than P, the algorithm enters the insertion-merge mode. In this mode, the algorithm flow is as follows: Figure 8 As shown, the overall process can be divided into an insertion stage and a merging stage. Let the length of the input sequence in this TopK test be N, the length of the target sequence be K, the length of the input sequence after preprocessing and padding to a multiple of P be N', and the length of the target sequence after padding to a multiple of P be K'. It should be noted that the parallelism P of the comparator is generally a power of 2.

[0075] Insertion stage. The function of the insertion stage is to sort the input sequence of length N' into a segmented ordered sequence with each segment consisting of P elements.

[0076] See Figure 8 The insertion process is similar to the aforementioned insertion mode process, except that: 1. In the insertion phase, after sorting each batch of P elements, the current output value and the data in the index register are written to SRAM, and the output value and index register are reinitialized. This operation is not involved in the insertion mode.

[0077] 2. In the insertion process, when larger=1, the comparison operator is less than or equal to; when larger=0, the comparison operator is greater than or equal to, which is the opposite of the aforementioned insertion mode.

[0078] The merge stage. The merge stage merges the ordered segments of P elements each pairwise, ultimately obtaining the K largest / smallest values / indices. The merge stage can be divided into multiple merge phases, each with a target sorting granularity g. i In the i-th merging phase, every two pairs of lengths g i-1 The ordered subsequences are merged into a single sequence of length g. i The subsequences are sorted starting from i=1, with initial g0=P and g1=2P. i-1 When K is less than K, g i =2g i-1 When g i-1 When K is greater than or equal to, g i =g i-1 , while i increases by 1 with each stage. The ( ) The ) phase ends, the merging ends, and the final output is obtained. Let's consider using h instead. i Let the coverage granularity of the i-th merge stage be described (i.e., the coverage range of the output of the i-th merge stage in the initial input of this TopK, ignoring the part at the end that cannot reach g or h, which will be considered in special case handling). Then, initially h0=P, h1=2P, h i =2h i-1 Therefore, when g i-1 When K is less than K, g i =h i When g i-1When K is greater than or equal to, g i <h i That is, g i with h i All multiply by 2 as the stage number i increases, but when the current target sorting granularity g i When g is greater than or equal to K, i No longer multiplied by 2 (frozen), but h i Continue multiplying by 2.

[0079] See Figure 8 The sub-steps in the merging process are described below (for the sake of illustration, we assume that N' is an ideal power of 2; the general case where N' is not a power of 2 will be explained in the special case handling section): 1. Preprocessing: Before entering the merging stage, initialize the current target sorting granularity to 2P, and the current sorting stage is stage 1.

[0080] 2. Read in two new consecutive segments (A and B, where A and B are identifiers for the two consecutive segmented ordered sequences for distinction) of ordered numerical values, the first P values ​​of each index, and the index, respectively, into input registers A and B. Regarding the specific meanings of A and B, for example, when g... i =h i When the result is 4P, all the results from the previous round are grouped into groups of 4P. Among the 4P elements, the first 2P and the last 2P elements are ordered. A refers to the first 2P elements and B refers to the last 2P elements. Input register A is used to store the data and index corresponding to the first 2P elements in the merge stage, and input register B is used to store the data and index corresponding to the last 2P elements in the merge stage.

[0081] 3. The header elements of the two segments are compared. The winning element is moved to the end of the output register, and the corresponding remaining elements are shifted to the right. The original elements in the output register are also shifted to the right. "Comparing the header elements of the two segments" refers to comparing the 0th input value register A and the 0th input value register B. When larger=1, the comparison operator is greater than or equal to; when larger=0, the comparison operator is less than or equal to. "Moving the winning element to the end of the output register and shifting the corresponding remaining elements to the right" means that, based on the comparison result, the data in the 0th input value and index register of the winning element (A or B) is moved to the (P-1)th output value and index register, while the data in the 1st to (P-1)th input value and index registers are moved to the 0th to (P-2)th input value and index registers. The (P-1)th input value and index register is filled with 32'hFFFF_FFFF or 32'h7FFF_FFFF, depending on larger. "Shifting the original elements in the output register to the right" means that the data in the 1st to (P-1)th output value and index registers are moved to the 0th to (P-2)th output value and index registers.

[0082] 4. When the output register is full (i.e., as the merging process proceeds, all P output values ​​and the index register are filled with valid data), the data in the output values ​​and index register are written to SRAM. At this point, if the current sorting stage has reached the final goal ( If the number of output elements reaches K (greater than or equal to K), then the output ends, and this TopK iteration is complete.

[0083] 5. If this is not the last output of the current TopK sequence, then it is necessary to consider whether the current group has been merged. The current group's output value and index number reach g. i When this happens, the merging of the current group is considered complete. For example, when g... i =2P, h i When the count is 4P, if the current group has already output 2P data points, the merging is considered complete; otherwise, it is not complete.

[0084] 6. If the current group has been merged, then it is necessary to consider whether all groups have been merged. Here, for example, when the input sequence length N' = 16 × 4P, g i =2P, h i When N' = 4P (assuming N' is an ideal power of 2 for the purpose of introduction; the general case where N' is not a power of 2 will be introduced later), once all 16 groups of 2P data have been output, it is considered that the merging of all groups is complete.

[0085] 7. If all groups have been merged, it is necessary to consider whether the current target sorting granularity has reached the final goal. If the current target sorting granularity has not yet reached the final goal, it is multiplied by 2; if the final goal has been reached, the current target sorting granularity remains unchanged. Regardless of whether this condition is true or not, the current sorting stage is incremented by 1. This is the update of metadata between merging stages, that is, the update logic of the target sorting granularity g and the coverage granularity h.

[0086] 8. If the current merging is not yet complete, or the output register is not full, it is necessary to continue considering whether there is a situation where input register A (or B) is empty and there are still remaining unread elements (i.e., as the merging proceeds, if the input value and index register A (or B) are both filled with 32'hFFFF_FFFF or 32'h7FFF_FFFF (no valid data), and the segmented ordered sequence A (or B) has not been completely read). If this condition is met, the next batch of P unread values ​​from the segmented ordered sequence A (or B) is read from SRAM, indexed to the input value, and then to index register A (or B). Note that since only the head element of either A or B is moved to the (P-1)th output register each time, at most one of A and B will trigger this condition.

[0087] 9. Regardless of whether input register A (or B) is empty and there are still unread elements, the comparison of the first element values ​​of the two segments must continue. The loop can only be exited when the output register is full, the current group merge is complete (entering the next group or the next merge stage), or the current sorting stage reaches its final goal. And the number of output elements reaches K (greater than or equal to K) (that is, this is the last time to write the output value and index register to SRAM in this TopK).

[0088] In the above insertion and merge pattern, some special cases may arise. For example, in general, N' is not a power of 2, so during the merge process, the last group may have fewer than 2g elements. i-1 This situation is referred to here as a "special case".

[0089] Let m be the number of input elements (value pairs, index pairs) in the j-th group of the i-th stage. ij ,0≤j≤N' / h i The number of input elements in the i-th stage is n. i It should be noted that only the j-th = N' / h i The last group may have a "special case". Let m be the number of input elements in the last group of the i-th stage. i '.

[0090] This can be further subdivided into three cases: (1)g i-1 <m i <2g i-1 =g i At this point, in the last group, A corresponds to a complete sequence of length g. i-1 Given an ordered sequence, B corresponds to a fragment of length m. i '-g i-1 An ordered sequence. And m i 'Small than the target sorting granularity g of the current stage i The difference between its merging process and the aforementioned normal process is as follows: ① Due to m i <2g i-1 When determining whether the condition "input register B is empty and there are still unread elements" is true, the total number of elements in B is m. i '-g i-1 (instead of the normal g) i-1 ).

[0091] ②Because m i ' <g i When determining whether the current group has been merged, if the output value and index of the current group reach m... i'(instead of the normal g)' i When the condition is met, the current group is considered to have been merged.

[0092] (2)g i =g i-1 <m i <2g i-1 At this point, in the last group, A corresponds to a complete sequence of length g. i-1 Given an ordered sequence, B corresponds to a fragment of length m. i '-g i-1 An ordered sequence. And m i 'Greater than the target sorting granularity g of the current stage i The difference between its merging process and the aforementioned normal process is as follows: ① Due to m i <2g i-1 When determining whether the condition "input register B is empty and there are still unread elements" is true, the total number of elements in B is m. i '-g i-1 (instead of the normal g) i-1 ).

[0093] 0 <m i '≤g i-1 At this point, in the last group, A corresponds to the missing element of length m. i The ordered sequence is ', while B is missing. Figure 8 The condition "whether the number of elements in the current group is greater than the target sorting granularity of the previous stage" is designed for this purpose (where, if the current stage is stage 1, the target sorting granularity of the previous stage refers to the target sorting granularity of the insertion stage, i.e., P). The merging process differs significantly from the aforementioned normal process. Since B is missing, A does not need to be merged with B. Therefore, the values ​​and indices corresponding to the ordered sequence A in the current group are directly moved to the SRAM target address, skipping the intermediate merging process.

[0094] like Figure 9 The image shows an example of a single TopK process in the insertion merge pattern, where the comparator parallelism P=32, the total number of elements N=476, the number of elements to be selected K=62, and the size flag largest=1. Note that each row contains P=32 elements (value-index pairs).

[0095] In a single TopK sorted sequence using the insertion-merge pattern, the original data first undergoes an insertion stage, transforming it into a segmented, ordered sequence with a segment size of P=32. Then, it goes through a merge stage comprising four sorting phases to obtain the final result. It's important to note that in sorting phase 2, since the segmented, ordered segment size already exceeds K'=64, the target sorting granularity is fixed to the smallest power of 2 of K, which is 64. Finally, the output sequence, after padding to a multiple of P, has a length of K'=64.

[0096] When entering pure merge mode, the operation of the device in this embodiment of the application is as follows: When the sorted granularity G is not 0 and K is greater than P, the algorithm enters pure merge mode. This mode is primarily designed for block-based TopK partitioning. The algorithm flow in this mode is as follows: Figure 10 As shown, the workflow is similar to the merge step in the insert-merge pattern, with the only difference being: 1. When initializing the target sorting granularity, directly initialize it to the input sorted granularity G.

[0097] 2. When initializing the current merge phase, it is directly initialized to (log2G-log2P+1) (note that G and P are both powers of 2).

[0098] 3. Between merge phases, the target sorting granularity is frozen and no longer multiplied by 2.

[0099] The remaining procedures are the same and will not be repeated here. (Step (...)) In the first phase, when the number of output elements reaches K, the merging ends.

[0100] like Figure 11 The image shows an example of a single TopK process in pure merge sort mode, where the comparator parallelism P=32, the total number of elements N=476, the number of elements to be selected K=62, the size flag largest=1, and the sorted granularity G=64. Note that each row contains P=32 elements (value-index pairs).

[0101] In a single TopK run in pure merge mode, the sorting phase is initialized to (log2G-log2P+1)=2. After three phases of merging, the final output is obtained, and the number of output elements is a multiple of P. The length of the target sequence is K'=64.

[0102] Figure 12 The example illustrates a TopK process with comparator parallelism P=32, total number of elements N=32680, number of elements to be selected K=22, and size flag largest=1. Each arrow represents a single TopK. Since K=22≤32=P, all single TopKs in this example are pure insertion mode single TopKs.

[0103] Figure 13 The diagram illustrates the TopK process with comparator parallelism P=32, total number of elements N=32680, number of elements to be selected K=3803, and size flag largest=1. Each solid arrow represents a single TopK, while dashed arrows indicate direct movement to the target address because the total number of elements in the block N0' < the sorted granularity G. Since K=3803>32=P, a single TopK in this example consists of insertion-merge mode and pure merge mode TopK. The first round of a single TopK is generally insertion-merge mode (not all are insertion-merge mode; for example, in some configurations, if the total number of elements in the block of the last remaining group N0'≤P=32, the single TopK is a pure insertion mode single TopK with parameters N=N0' and K=N0'), while subsequent rounds of single TopK are pure merge mode.

[0104] As can be seen from the above technical solutions, the general TopK computing device based on insertion merging provided in this application embodiment, by setting hierarchical and segmented parallel merging logic, freezes and stops expanding after the target sorting granularity reaches K, only retaining the normal expansion of the covering granularity. This not only makes full use of parallel computing resources to improve computing efficiency, but also avoids sorting redundant elements that exceed the TopK requirement range, reducing unnecessary computing and storage overhead. It can adapt to TopK computing requirements with different input sizes, different K values, and different initial ordering levels, and has stronger versatility. The hardware implementation has better resource utilization and computing efficiency.

[0105] Meanwhile, this application can reuse the parallel comparison unit in the existing vector computing unit. The TopK-supporting computing unit also supports other types of vector computing and flexibly supports variable N, K, and largest parameters. The circuit structure provided in the embodiments of this application is not for a fixed N, K, or largest. It not only outputs the selected value but also its corresponding index. The output result is stable, that is, when the values ​​are equal, the smaller index is output first. The number of clock cycles required by the operator is small, thereby effectively improving the support capability and computational efficiency for hybrid expert models and various sparse strategies.

[0106] The above specific embodiments further illustrate the purpose, technical solution and beneficial effects of this application. It should be understood that the above are only specific embodiments of this application and are not intended to limit the scope of protection of this application. Any modifications, equivalent substitutions, improvements, etc., made on the basis of the technical solution of this application should be included within the scope of protection of this application.

Claims

1. A general TopK computing device based on insertion merge, characterized in that, include: The instruction module is configured to receive TopK calculation instructions; the TopK calculation instructions include the total number of elements, the number of elements to be selected, and an identifier bit. The control module is configured to generate a sequence of control signals for controlling the computation module to perform TopK calculations based on the TopK calculation instructions. The computing module includes multiple parallel computing units, and the computing module is configured as follows: Based on the input control signal sequence, the operation mode to be entered is determined; the operation mode includes pure merge mode, pure insertion mode, or insertion merge mode. In the operation mode, the parallel computing unit performs calculations on the data to be processed loaded from the memory and outputs the calculation result corresponding to the TopK calculation instruction; the data to be processed includes input values ​​and indices, wherein the indices are either continuous indices generated based on the first input or input indices based on non-first inputs.

2. The general TopK computing device based on insertion merging according to claim 1, characterized in that, The computing unit has P paths, and the P-path computing units are connected sequentially along the P-1th path to the 0th path to form a data path. The register data of the computing unit of the (i+1)th path is used as the input shift data of the computing unit of the i-th path; Where i∈[0, p-2]; The computing unit includes: The system comprises an output register group, a first input register group, and a second input register group; the output register group is used to store the output value and the output index; the first input register group is used to store the first input value and the first input index; and the second input register group is used to store the second input value and the second input index. The comparator is used to compare the output value, the first input value, and the second input value based on the comparison operator, and output a single-bit comparison result.

3. A general TopK computing device based on insertion merging according to claim 2, characterized in that, The computing module is also configured to: The single-bit comparison results output by the P-way computing unit are concatenated to generate a shift control signal; the shift control signal is used to control the output value and output index received by the output register group.

4. A general TopK computing device based on insertion merging according to claim 3, characterized in that, Before the computation module outputs the calculation result corresponding to the TopK calculation instruction, it is specifically configured as follows: Parameter initialization: Initialize the total number of elements to N0 and the sorted granularity to 0. When the total number of elements is greater than a preset value, the data to be processed is divided into several data blocks; the total number of elements in each data block is less than or equal to the preset value. Calculate the single TopK calculation result for each of the data blocks; the single TopK calculation result includes the selected value and index; The results of a single TopK calculation of all the data blocks are concatenated to update the current total number of elements and the sorted granularity; When the updated total number of current elements is less than or equal to a preset value, the single TopK calculation result corresponding to the updated total number of current elements is used as the calculation result corresponding to the TopK calculation instruction.

5. A general TopK computing device based on insertion merging according to claim 4, characterized in that, The computation module determines the computation mode to be entered, and is specifically configured as follows: If the sorted granularity is 0, the data to be processed is preprocessed. After preprocessing, if the number of elements to be selected is less than or equal to the number of computational units, then it is determined to enter the pure insertion mode; If the number of elements to be selected is greater than the number of computation units, then the insertion and merging mode is entered. If the sorted granularity is not 0, and the number of elements to be selected is less than or equal to the number of computational units, then it is determined to enter the pure insertion mode. If the number of elements to be selected is greater than the number of computation units, then the system will proceed to the pure merge mode.

6. A general TopK computing device based on insertion merging according to claim 5, characterized in that, The TopK calculation instruction also includes local data identifier bits; The computing module preprocesses the data to be processed, specifically configured as follows: The input values ​​in the data to be processed are padded to an integer multiple of the number of computing units; When the local data identifier indicates that the current input is original local data, a continuous index corresponding to the original input sequence is generated based on the starting position of the data block and the position of the element within the data block; When the data to be processed already contains an index, the index is used as the input index.

7. A general TopK computing device based on insertion merging according to claim 5, characterized in that, The calculation module, in the pure insertion mode, outputs the calculation result, specifically configured as follows: Initialize the output register group; Fill the first input register group with P values ​​and their indices from the input values; The comparison result between the input value in the first input register group and the current value in the output register group is obtained based on the comparator; The output register group is updated based on the comparison results until the comparison results corresponding to all input values ​​are updated to the output register group. The calculation result is obtained by symmetrically flipping the values ​​and indices in the output register group.

8. A general TopK computing device based on insertion merging according to claim 7, characterized in that, The computation module, in the pure merge mode, outputs the calculation result, specifically configured as follows: Initialize the current target sorting granularity and the current sorting stage; Fill the first P values ​​and indices of the two consecutive segmented ordered values ​​in the input values ​​into the first input register group and the second input register group, respectively. The first element values ​​of the two segments are compared, the winner is moved to the end of the output register group, the corresponding remaining elements are shifted to the right, and the original elements of the output register group are shifted to the right. The calculation result is output when the target is reached in the current sorting stage and the number of output elements in the output register group reaches the number of elements to be selected.

9. A general TopK computing device based on insertion merging according to claim 8, characterized in that, The computation module, in the insertion and merge mode, outputs the calculation result, specifically configured as follows: The insertion phase is performed; the insertion phase is the same as the steps performed by the operation module in the pure insertion mode, wherein, after the comparison results corresponding to all input values ​​are updated to the output register group, the data in the output register group is written to the memory; Perform the merge phase; the merge phase includes: Initialize the current target sorting granularity and the current sorting stage; Fill the first P values ​​and indices of the two consecutive segmented ordered values ​​in the input values ​​into the first input register group and the second input register group, respectively. The first element values ​​of the two segments are compared, the winner is moved to the end of the output register group, the corresponding remaining elements are shifted to the right, and the original elements of the output register group are shifted to the right. When the target is reached in the current sorting stage and the number of output elements in the output register group reaches the number of elements to be selected, the calculation result is output. Achieving the objective in the current sorting stage includes: All groups of the data to be processed have been merged, and the current target sorting stage is incremented by 1. Alternatively, the target sorting granularity is multiplied by 2, and the current target sorting stage is incremented by 1. Then, the step of filling the first P values ​​and indices of the two consecutive segments of ordered values ​​in the input values ​​into the first input register group and the second input register group is executed again.

10. A general TopK calculation method based on insertion merging, applied to the apparatus of any one of claims 1 to 9, characterized in that, The method includes: Receive TopK calculation instructions; the TopK calculation instructions include the total number of elements, the number of elements to be selected, and an identifier bit; Based on the TopK calculation instructions, a control signal sequence is generated to control the operation module to perform TopK calculations; Based on the input control signal sequence, the operation mode to be entered is determined; the operation mode includes pure merge mode, pure insertion mode, or insertion merge mode. In the aforementioned operation mode, parallel computing units perform calculations on the data to be processed loaded from memory and output the calculation results corresponding to the TopK calculation instructions. The data to be processed includes input values ​​and indices, where the indices are either continuous indices generated based on the first input or input indices based on subsequent inputs.

Images