An in-memory top-k unit based on depth-first search and an implementation method thereof

Abstract

The application discloses an in-memory Top-K unit based on depth-first search and an implementation method thereof, comprising at least two in-memory computing arrays, local sorting modules with the same number of arrays and a global decision module. The in-memory computing array is used for storing input elements to be sorted, so as to reduce data exchange overhead. The local sorting module is used for reading data from high bits to low bits in a bit serial manner and counting the number of local candidate elements. The global decision module is used for flexibly aggregating the number of local candidate elements according to grouping information to support different numbers of input elements, and generating a control signal to guide the local sorting module to exclude elements that cannot become the first K maximum values. Through the joint strategy of in-situ calculation and reconfiguration, the application realizes efficient processing of different scale Top-K computing tasks and dynamic elimination of redundant iteration processes.

Description

Technical Field

[0001] This invention relates to engineering technologies such as in-memory computing, nonlinear operations, neural networks, and high-efficiency computing circuits. In particular, it discloses an in-memory Top-K cell based on depth-first search and its implementation method, which belongs to the technical field of computing, calculation, or counting. Background Technology

[0002] In recent years, with the surge in demand for computing resources from artificial intelligence applications, the traditional von Neumann computing architecture, due to its separation of computation and storage, faces a severe "memory wall" bottleneck. Frequent data movement between processors and memory generates significant latency and power consumption. To address this challenge, the Compute-in-Memory (CIM) paradigm has emerged. It deeply integrates computation and storage functions, utilizing storage arrays to perform intensive parallel computations such as multiplication and addition in neural networks in-situ. By significantly reducing data movement, CIM effectively alleviates the memory wall problem and has become a key technology in AI deployments.

[0003] To balance model performance and computational overhead, existing neural network architectures are evolving towards sparsity. For example, to address the quadratic computational bottleneck of the Transformer architecture when processing long sequences, sparse attention mechanisms have emerged. Their core lies in dynamically selecting the most critical subset of key-value pairs from a wide range using the Top-K algorithm, leveraging the fault tolerance of neural networks to remove pruned pairs that contribute little to the result. Similarly, hybrid expert architectures, which can extend the generalization performance of models, also rely on the Top-K algorithm. This architecture does not need to activate all expert networks during inference; instead, it dynamically activates the few most relevant experts based on the input terms to complete the computational task. Top-K computation has evolved from a general auxiliary algorithm into a core computational component for building efficient neural networks.

[0004] However, existing in-memory computing architectures face fundamental challenges in handling the increasingly critical Top-K computations. While the regular physical structure and fixed data flow of in-memory computing make it well-suited for large-scale multiply-accumulate operations, it lacks the comparison and transfer capabilities necessary for Top-K computations, thus failing to natively support sorting operations. Currently, the industry-standard solution is a heterogeneous architecture, where a separate application-specific integrated circuit (ASIC) is coupled into the system to perform Top-K computations. This discrete design has two serious drawbacks: First, the independent sorting unit relies on numerous comparison and routing units, and the resulting latency, area, and power consumption become the performance bottleneck of the entire system, severely limiting the throughput of sorting tasks. Second, this approach hinders the utilization of system resources. During Top-K computations, the in-memory computing units, which occupy significant system resources, remain idle, leading to a deterioration in area efficiency. The low throughput of the independent sorting unit further amplifies this limitation.

[0005] In summary, the present invention aims to propose an in-memory Top-K unit based on depth-first search and its implementation method to overcome the above-mentioned defects. Summary of the Invention

[0006] The purpose of this invention is to address the shortcomings of the aforementioned background technology by providing an in-memory Top-K unit based on depth-first search. This unit employs a bit-serial, step-by-step counting depth-first search strategy to achieve rapid screening and early pruning of candidate elements within the local sorting module. Combined with a path-reconfigurable global decision module, the dynamic grouping and computation path reconfiguration of the local sorting module allow the computing unit to dynamically adjust its parallelism and computational granularity according to application requirements. This ensures high-efficiency in-memory computing capabilities while solving the technical problems of insufficient hardware utilization and low throughput in existing independent sorting units, thus achieving the invention's objective of combining in-memory computing with high-efficiency Top-K computing.

[0007] To achieve the above-mentioned objectives, the present invention employs the following technical solution: A depth-first search-based in-memory Top-K unit includes: an in-memory computing array for storing input elements to be sorted, wherein each row of 6T SRAM corresponds to the same bit of all input elements; The local sorting module, connected to the in-memory computing array, reads the bit sequence bit by bit, starting from the most significant bit of the input element. Each local sorting module includes: a register group for marking the state of each input element; a checker for determining whether the currently processed bit sequence is all bits 0; and a counter for counting the number of bits 1 in the current bit sequence and outputting it as the number of local candidate elements. The counter generates the number of local candidate elements based on the checker's selection signal and sends it to the global decision module. The local sorting module receives control signals from the global decision module to update the register group or stop the Top-K calculation early. The global decision module receives the output of the local sorting module and controls the local sorting module in reverse. The global decision module includes: a path-reconfigurable addition tree for accumulating the number of local candidate elements from the local sorting module; and a decision path for generating control signals. The number of decision paths is the same as the number of local sorting modules, and the decision path is determined based on the sum of the candidate elements generated by the reconfigurable addition tree.

[0008] As a further optimization scheme for in-memory Top-K units based on depth-first search, the decision path can adaptively select the computation path based on the sum of the input candidate elements. If the number of input candidate elements is 0, the adders and comparators in the decision path are bypassed, the registers in the decision path are not updated, and the output control signal remains unchanged. As a further optimization scheme for in-memory Top-K units based on depth-first search, the path-reconfigurable addition tree consists of multiple adders and has multiple operational levels, where the... Each operational level Used for The summation of inputs; the result of each operation level in the path-reconstructable addition tree is connected to a multiplexer and selected as a candidate element sum based on grouping information.

[0009] As a further optimization of the in-memory Top-K unit based on depth-first search, the checker in the local sorting module is used to determine whether the currently processed bit sequence is all bits 0. If the currently processed bit sequence is all bits 0, the checker outputs a result of 1. At this time, the local candidate element counting operation in the counter is bypassed, thereby accelerating the calculation process.

[0010] As a further optimization of the in-memory Top-K unit based on depth-first search, the register group in the local sorting module stores a 2-bit status code for each input element. The status code indicates the computational status of the element, and its least significant bit serves as an activity flag to determine whether the input element participates in the search cycle. The specific codes and corresponding states are as follows: "Activated" is encoded as "01", indicating that the element is a currently valid candidate and needs to participate in subsequent search cycles; "Sorted" is encoded as "10", indicating that the element has been determined as one of the Top-K results and will be ignored in subsequent calculations; "Excluded" is encoded as "00", indicating that the element has been pruned because it does not meet the conditions and will be ignored in subsequent calculations.

[0011] As a further optimization scheme for the in-memory Top-K unit based on depth-first search, the following operations are performed when the decision path generates the control signal: When the adder result in the decision path is greater than the sorting target K value, the corresponding local sorting module is instructed to update the status of the candidate element with the current bit 0 to "excluded"; when the adder result in the decision path is less than the sorting target K value, the corresponding local sorting module is instructed to update the status of the candidate element with the current bit 1 to "sorted", and the adder result in the decision path will be updated in the selected element count register in the decision path; when the adder result in the decision path is equal to the sorting target K value, the entire Top-K calculation process is terminated early.

[0012] As a further optimization scheme for in-memory Top-K units based on depth-first search, When the number of elements to be sorted is less than or equal to N, each local sorting module is considered an independent computation group. The number of local candidate elements output by each module is bypassed by the multiplexer in the path-reconfigurable addition tree, bypassing all computation levels, and directly fed into its corresponding decision path for independent decision-making. N is the maximum number of elements a local sorting module can sort. When the number of elements to be sorted is greater than N, multiple local sorting modules are merged into one computation group, with each group containing a specified number of local sorting modules. , To meet The smallest positive integer, Let be the number of elements to be sorted. At this point, the sum of the candidate elements output by the reconfigurable adder to the decision path is the th element. The results of each computational level.

[0013] The present invention provides a method for implementing in-memory Top-K units based on depth-first search, comprising the following steps: Step 1: Initialize the pointer controller in the local sorting module to point to the most significant bit of the input element, initialize the value in the status register group to "01" for the "activated" state, and clear the number of selected elements to zero;

[0014] Step 2: The pointer controller drives the in-memory calculation unit to read the bit sequence composed of all elements under the current bit position, and updates the pointer controller through the subtractor; Step 3: The verifier accepts the bit sequence and the least significant bit of the status code of each element to determine whether all the "activated" bits are 0. If so, the counting operation in the local sorting module is skipped and the number of local candidate elements is set to 0. Otherwise, the counter counts the number of 1s in all the "activated" bit values, generates the number of local candidate elements and outputs it. Step 4: The path reconfigurable additive tree flexibly aggregates each local candidate element according to the grouping information. The result will be distributed to the decision path corresponding to the local sorting module. Step 5: The decision path receives the calculation result of the path reconstructable addition tree, adds it to the number of selected elements through an adder, and compares the addition result with the sorting target K value through a comparator. Step 6: Update the control register according to the comparison result. If it is less than the sorting target K value, update the "activated" element with the current bit value of 1 to "sorted", and update the number of selected elements to the sum of the results, then proceed to step 2. If it is greater than the sorting target K value, update the "activated" element with the current bit value of 0 to "excluded", then proceed to step 2. If it is equal to the sorting target K value, update the "activated" element with the current bit value of 1 to "sorted", and end the Top-K calculation.

[0015] The beneficial effects of this invention are: (1) This invention provides an in-memory Top-K unit based on depth-first search, which supports bit-serial depth-first search based on an in-memory computing array and combines a parallel global decision module for real-time pruning. Thus, the hardware can dynamically prune invalid search space layer by layer, starting from the most significant bit of the input element. This mechanism transforms a large number of comparison and swap operations into efficient counting operations, reducing computational complexity and enabling the in-memory computing unit to handle Top-K calculations, thus overcoming the limitation of low hardware utilization.

[0016] (2) The computing unit proposed in this invention has the characteristic of reconfigurable number of sorted elements. When the number of sorted elements is greater than the maximum number of sorted elements in the local sorting module, the computing grouping mechanism can flexibly adjust the parallelism and granularity of the computing according to the application requirements of the upstream, which greatly improves the utilization of hardware resources and adaptability to different computing tasks. This solves the problem of low efficiency caused by the fixed function of dedicated sorting units in traditional heterogeneous computing, and significantly improves the overall throughput. Attached Figure Description

[0017] Figure 1 This is a schematic diagram of the overall architecture of an in-memory Top-K unit based on depth-first search proposed in this invention.

[0018] Figure 2 This is a schematic diagram of the structure of the local sorting module involved in the present invention.

[0019] Figure 3 This is a schematic diagram of the global decision-making module involved in the present invention.

[0020] Figure 4 This is a schematic diagram of the path-reconfigurable additive tree involved in the present invention.

[0021] Figure 5 This is a flowchart of the in-memory Top-K unit based on depth-first search proposed in this invention. Detailed Implementation

[0022] The technical solution of the invention will now be described in detail with reference to the accompanying drawings.

[0023] like Figure 1 As shown, the present invention proposes an overall architecture for an in-memory Top-K unit based on depth-first search, used to compute the operation of finding K maximum values ​​in a neural network. The in-memory Top-K unit includes: at least two in-memory computing arrays, a local sorting module with the same number of arrays, and a global decision module.

[0024] The in-memory computing array can handle multiply-accumulate operations, and the results are dequantized and directly stored in 6T SRAM for use as operands in Top-K calculations. This mechanism reduces data movement between internal and external memory, and in conjunction with the local sorting module and the global decision module, this computing unit can achieve in-situ Top-K calculations.

[0025] The local sorting module is tightly integrated with the in-memory computing array, and can perform parallel reading of bit-serial data. The block diagram of the local sorting module is shown below. Figure 2As shown, the array consists of a pointer controller, a register set, a checker, and a counter. The pointer controller reads the bit sequence bit by bit, starting from the most significant bit of the input element, by driving the sub-lines of the in-memory computation array. For example, for a 32-bit input element, the initial value of the pointer controller is 31, and its count is decremented by 1 after each search cycle. The register set stores a 2-bit status code for each input element, where "01" indicates "activated," "00" indicates "excluded," and "10" indicates "sorted." The least significant bit of the status code is considered the activity flag; if it is 0, the corresponding element is not included in the calculation. The checker receives the read bit sequence and the corresponding status code to quickly determine if all "activated" bits are all 0. The counter records the number of 1s among all "activated" bits; this result is called the number of local candidate elements.

[0026] The structural block diagram of the global decision-making module is as follows: Figure 3 As shown, it consists of a path-reconfigurable adder tree and decision paths, used to accept the number of candidate elements from each local sorting module and perform a decision. The path-reconfigurable adder is used to flexibly accumulate the number of local candidate elements sent by each local sorting module according to the grouping information. Its specific structure is as follows. Figure 4 As shown. This addition tree has multiple operation levels, the first being... The result of each operation level is The sum of the input elements. Based on the number M of local sorting modules in the computation group, the reconstructable addition tree selects the nth... The results of each computational level are sent to the decision path, where To meet The smallest positive integer. For example, if a computation group has 3 local sorting modules, then the result of the 2nd operation level is selected and sent to the corresponding decision path. Decision paths are used to execute decisions, and their number is equal to the total number of local sorting modules. Each decision path includes a register for recording the number of selected elements, an adder for adding the number of selected elements and the sum of the candidate elements, and a comparator for comparing the adder result with the sorting target K value. The comparison result guides the local sorting module state update.

[0027] The overall workflow diagram of the in-memory Top-K unit of this invention is as follows: Figure 5As shown, firstly, the pointer controllers and register groups in each local sorting module are reset to their initial values. The initial value of the pointer controller is equal to the bit width of the element to be sorted minus 1, and the initial value of the register group is 01, indicating "activated". Next, the local sorting module obtains the bit sequence by driving the bit lines of the in-memory computing array through the pointer controller. After reading, the value in the pointer controller is decremented by 1. Then, the checker in the local sorting module uses the AND logic of the tree structure to determine whether all "activated" bits in the bit sequence are all 0. If all are 0, the counter in the local sorting module is bypassed, and the number of local candidate elements is output as 0; otherwise, the counter is not counted. The counter counts and outputs the number of local candidate elements. Then, the global sorting module receives candidate elements from all local sorting modules and selects different levels of calculation results based on grouping information. The sum of candidate elements is distributed to decision paths corresponding to each local sorting module. Finally, in each decision path, the sum of candidate elements is added to the number of selected elements and compared with the sorting target value K. If it equals K, the corresponding local sorting module stops early. If it does not equal K, a control signal is generated based on the comparison result, and the corresponding local sorting module updates its register group accordingly. If the sum is less than K, the number of selected elements in the decision path is also updated. In the next cycle, each local sorting module drives the bit line to read the next bit sequence and begins a new calculation iteration.

[0028] Assume that each local sorting module in the Top-K unit of memory can process four 8-bit unsigned numbers at a time. When processing the Top-2 calculation of four elements, each calculation group contains only one local sorting module. Assume the four elements to be sorted are A = 01110011 (binary) = 115 (decimal), B = 00101010 (binary) = 42 (decimal), C = 00111000 (binary) = 56 (decimal), and D = 00010110 (binary) = 22 (decimal). At the start of the sorting calculation, the pointer controller points to the most significant bit, i.e., the 7th bit, and the register for the number of selected elements in the decision path is cleared. At this time, the state of all four elements (A, B, C, and D) is 01, i.e., "activated". In the first computation cycle, the local sorting module reads the bit sequence {0 0 0 0} from the bit lines of the in-memory computation array via the pointer controller. This bit sequence, after passing through the AND logic tree in the checker, results in a value of 1 (all bits are 0). This signal drives the selector in the counter to output a candidate element count of 0. In this grouping case, the reconfigurable adder tree in the global decision module sends the input of the local sorting module as the sum of candidate elements to the decision path. Because the sum of candidate elements is 0, subsequent computation operations in this cycle are bypassed to optimize energy consumption. In the second computation cycle, the local sorting module reads the bit sequence {1 0 0 0}. The checker determines that this bit sequence is not entirely composed of bits 0, and the counter counts the number of 1s among the active elements, which is 1 (i.e., the number of local candidate elements is 1). Since the local sorting module within the group is 1, the sum of candidate elements equals the number of local candidate elements, which is 1. Adding the sum of candidate elements to the selected elements yields a result of 1, which is compared to the sorting target value K, and is found to be less than K. In this case, the number of selected elements in the decision path is updated to 1, and the local sorting module updates the state of element A corresponding to bit 1 to 10, i.e., "sorted". In the third calculation cycle, the local sorting module reads the bit {1 1 1 0}. The checker determines that the bit sequence is not all bits 0, counts the number of 1s in the active elements to be 2, the reconfigurable adder is bypassed, and the sum of the output candidate elements is 2. Adding the sum of the candidate elements to the selected elements yields a result of 3, which is compared with the sorting target value K, and the result is greater than K. In this case, the local sorting module updates the state of element D corresponding to bit 0 to 00, i.e., "excluded". In the fourth calculation cycle, the local sorting module reads the bit {1 0 1 1}. The checker determines that the bit sequence is not all bits 0, counts the number of 1s in the active elements to be 1, the reconfigurable adder is bypassed, and the sum of the output candidate elements is 1. The sum of the candidate elements is added to the selected elements, resulting in 2. This result is compared with the sorting target value K, and the result is equal to K. In this case, the sorting calculation stops early, and the local sorting module updates the state of element C corresponding to bit 1 to 10, i.e., "sorted".

[0029] The number of unsorted elements in the Top-K units of this memory can be reconstructed. Assume the number of local sorting modules in the Top-K units of this memory is M, where M is a power of 2, and the maximum number of sorted elements in each local sorting unit is N. Then the range of the number of unsorted elements supported by the Top-K units in this memory is: When the number of elements to be sorted is less than or equal to N, the computation group contains only one local sorting module, and the reconfigurable addition tree in the global decision module is bypassed. When the number of elements to be sorted, P, is greater than N, the number of local sorting modules in the computation group is... , To meet The smallest positive integer is , at which point the sum of the candidate elements output by the reconstructable adder to the decision path is the . The results of each computational level.

[0030] This invention is not limited to the preferred embodiments described above. Anyone inspired by this invention can derive various other forms of depth-first in-memory Top-K cells for Transformers. All equivalent variations and modifications made within the scope of the claims of this invention should be considered within the scope of this invention.

Claims

1. A memory-based Top-K system based on depth-first search, characterized in that, include: An in-memory computing array is used to store the input elements to be sorted, with each row of 6T SRAM corresponding to the same bit of all input elements; A local sorting module, connected to the in-memory computing array, is used to read the bit sequence bit by bit, starting from the most significant bit of the input element. Each local sorting module includes: a register group for marking the state of each input element; a checker for determining whether the currently processed bit sequence is all bits 0; and a counter for counting the number of bits 1 in the current bit sequence and outputting it as the number of local candidate elements. The counter generates the number of local candidate elements based on the selection signal of the checker and sends it to the global decision module. The local sorting module receives control signals from the global decision module to update the register group or stop the Top-K calculation early. A global decision module receives the output of a local sorting module and controls the local sorting module in reverse. The global decision module includes: a path reconfigurable addition tree for accumulating the number of local candidate elements from the local sorting module; and a decision path for generating control signals. The number of decision paths is the same as the number of local sorting modules, and the decision path makes a decision based on the sum of candidate elements generated by the reconfigurable addition tree. The path-reconfigurable addition tree consists of multiple adders and has multiple operation levels, where the i-th operation level is used for... Summing the inputs, The result of each operation level in the path-reconstructable addition tree is connected to a multiplexer and selected as the sum of the candidate elements according to the grouping information. When the decision path generates control signals, it performs the following operations: When the adder result in the decision path is greater than the sorting target K value, it instructs the corresponding local sorting module to update the status of the candidate element with the current bit 0 to "excluded"; when the adder result in the decision path is less than the sorting target K value, it instructs the corresponding local sorting module to update the status of the candidate element with the current bit 1 to "sorted", and at the same time, the adder result in the decision path will be updated in the selected element count register in the decision path; when the adder result in the decision path is equal to the sorting target K value, the entire Top-K calculation process is terminated in advance.

2. The in-memory Top-K system based on depth-first search according to claim 1, characterized in that, The decision path can adaptively select the calculation path based on the sum of the input candidate elements. If the number of input candidate elements is 0, the adder and comparator in the decision path are bypassed, the registers in the decision path are not updated, and the output control signal remains unchanged.

3. The in-memory Top-K system based on depth-first search according to claim 2, characterized in that, The verifier in the local sorting module is used to determine whether the currently processed bit sequence is all bits 0. If the currently processed bit sequence is all bits 0, the verifier outputs a result of 1. At this time, the local candidate element counting operation in the counter is bypassed, thereby accelerating the calculation process.

4. The in-memory Top-K system based on depth-first search according to claim 3, characterized in that, The register group in the local sorting module stores a 2-bit status code for each input element. The status code indicates the calculation status of the element, and its least significant bit serves as an activity flag to determine whether the input element participates in the search cycle. The specific codes and corresponding states are as follows: "Activated" is encoded as "01", indicating that the element is a currently valid candidate and needs to participate in subsequent search cycles; "Sorted" is encoded as "10", indicating that the element has been determined as one of the Top-K results and will be ignored in subsequent calculations; "Excluded" is encoded as "00", indicating that the element has been pruned because it does not meet the conditions and will be ignored in subsequent calculations.

5. The in-memory Top-K system based on depth-first search according to claim 4, characterized in that, When the number of elements to be sorted is less than or equal to N, each local sorting module is regarded as an independent computing group. The number of local candidate elements output by the module bypasses all operation levels through the multiplexer in the path reconstructable addition tree and is directly sent to the corresponding decision path for independent decision-making. N is the maximum number of sorted elements of the local sorting module. When the number of elements to be sorted is greater than N, multiple local sorting modules are merged into a single computation group, with each computation group containing the following number of local sorting modules: j is satisfied The smallest positive integer, P, is the number of elements to be sorted. At this time, the sum of the candidate elements output by the reconfigurable adder to the decision path is the result of the j-th computation level.

6. The implementation method of an in-memory Top-K system based on depth-first search according to any one of claims 1 to 5, characterized in that, Includes the following steps: Step 1: Initialize the pointer controller in the local sorting module so that it points to the most significant bit of the input element, initialize the value in the status register group to "01" for the "activated" state, and clear the number of selected elements to zero; Step 2: The pointer controller drives the in-memory computing array to read the bit sequence composed of all elements under the current bit position, and updates the pointer controller through the subtractor; Step 3: The verifier accepts the bit sequence and the least significant bit of the status code of each element to determine whether all the "activated" bits are 0. If so, the counting operation in the local sorting module is skipped and the number of local candidate elements is set to 0. Otherwise, the counter counts the number of 1s in all the "activated" bit values, generates the number of local candidate elements and outputs it. Step 4: The path reconfigurable additive tree flexibly aggregates each local candidate element according to the grouping information, and the aggregated element results are distributed to the decision path corresponding to the local sorting module. Step 5: The decision path receives the calculation result of the path reconstructable addition tree, adds it to the number of selected elements through an adder, and compares the addition result with the sorting target K value through a comparator. Step 6: Update the control register according to the comparison result. If it is less than the sorting target K value, update the "activated" element with the current bit value of 1 to "sorted", and update the number of selected elements to the sum of the results, then proceed to step 2. If it is greater than the sorting target K value, update the "activated" element with the current bit value of 0 to "excluded", then proceed to step 2. If it is equal to the sorting target K value, update the "activated" element with the current bit value of 1 to "sorted", and end the Top-K calculation.

Images