Method for performing number theoretic transform by tensor core of GPU and computer device

Abstract

The specification provides a method for implementing NTT by a tensor core of a GPU. First, a plurality of stored rotation factor matrices and an input data matrix in shared memory are loaded to a plurality of registers corresponding to a thread bundle. The thread bundle executes a first MMA instruction using the tensor core, takes the rotation factor matrix as a first multiplier slice corresponding to row data, and takes the input data matrix as a second multiplier slice corresponding to column data, to obtain a matrix multiplication result (a first matrix) of the rotation factor matrix and the input data matrix. A Hadamard product of the first matrix and a second rotation factor matrix is used to obtain a second matrix. Without moving data stored in a third register group, the thread bundle executes a second MMA instruction, takes data of the third register group storing the second matrix as a first multiplier slice corresponding to row data, and takes a third rotation factor matrix as a second multiplier slice, to obtain a matrix multiplication result of the second matrix and the third rotation factor matrix.

Description

Technical Field

[0001] This specification relates to the field of privacy computing technology, and more particularly to a method and computer device for performing number-theoretic transformations via the tensor core of a GPU. Background Technology

[0002] Fully homomorphic encryption (FHE) allows addition and multiplication operations to be performed directly on ciphertext to protect data privacy. Its core operation, polynomial multiplication, is usually accelerated by number-theoretic transformations (NTTs), but high-dimensional NTTs are computationally expensive and become a performance bottleneck.

[0003] Graphics Processing Units (GPUs) are commonly used to accelerate NTTs (Network Transaction Processing). Traditional methods primarily utilize Compute Unified Device Architecture (CUDA) cores for parallel computing. CUDA cores are the basic processing units in GPUs responsible for general-purpose logic operations. Traditional methods use them to execute butterfly operations in NTTs, but data needs to be frequently exchanged between registers and memory, causing computing units to be idle while waiting for data.

[0004] To improve efficiency, existing technologies introduce Tensor Cores to accelerate NTT computation. However, current solutions still rely on Shared Memory (SMEM) for data exchange and explicit matrix transposition during internal NTT computation, resulting in significant redundant memory accesses. These inefficient memory access issues prevent the full utilization of Tensor Core computing power. Therefore, a method for NTT computation that reduces intermediate data memory interactions is urgently needed. Summary of the Invention

[0005] In view of this, one or more embodiments of this specification provide a method and computer device for performing number-theoretic transformations via the tensor core of a GPU.

[0006] According to a first aspect of one or more embodiments of this specification, a method for performing number-theoretic transformations (NTTs) via the tensor core of a GPU is proposed, comprising:

[0007] Load the first rotation factor matrix, the second rotation factor matrix, the third rotation factor matrix, and the input data matrix from the shared memory into multiple registers corresponding to the thread bundle;

[0008] The thread executes a first matrix multiply-add MMA instruction, where the data in the first register group storing the first twitch factor matrix is ​​used as the first multiplier slice corresponding to the row data, and the data in the second register group storing the input data matrix is ​​used as the second multiplier slice corresponding to the column data, to obtain the first matrix; and the second matrix is ​​obtained according to the Hadamard product of the first matrix and the second twitch factor matrix; the second matrix is ​​stored in the third register group corresponding to the accumulated result slice;

[0009] The thread executes the second MMA instruction, in which the data of the third register group is used as the first multiplier slice, and the data of the fourth register group storing the third rotation factor matrix is ​​used as the second multiplier slice to obtain the third matrix;

[0010] The third matrix is ​​written into shared memory as an NTT result.

[0011] According to a second aspect of one or more embodiments of this specification, a computer-readable storage medium is provided that stores computer instructions thereon, which, when executed by a processor, implement the method of performing number-theoretic transformation NTT via the tensor core of a GPU as described in the first aspect of the embodiments of this specification.

[0012] According to a third aspect of the embodiments of this specification, a computer device is provided, the computer device comprising:

[0013] processor;

[0014] Memory used to store processor-executable instructions;

[0015] The processor executes the executable instructions to implement the method of performing number-theoretic transformation NTT via the tensor core of the GPU as described in the first aspect of the embodiments of this specification.

[0016] According to a fourth aspect of the embodiments of this specification, a computer program product is provided that, when executed by a processor, implements the method of performing number-theoretic transformation NTT via the tensor core of a GPU as described in the first aspect of the embodiments of this specification.

[0017] In one or more embodiments of this specification, multiple rotation factor matrices and an input data matrix stored in shared memory are first loaded into multiple registers corresponding to the thread bundle. The thread bundle executes a first MMA instruction using the tensor core, using the rotation factor matrices as first multiplier slices corresponding to the row data and the input data matrix as second multiplier slices corresponding to the column data, to obtain the matrix multiplication result of the rotation factor matrices and the input data matrix (the first matrix). The second matrix is ​​obtained by using the Hadamard product of the first matrix and the second rotation factor matrix. Without moving the data stored in the third register group, the thread bundle executes a second MMA instruction, using the data in the third register group storing the second matrix as the first multiplier slices corresponding to the row data and the third rotation factor matrix as the second multiplier slice, to obtain the matrix multiplication result of the second matrix and the third rotation factor matrix.

[0018] In the four-step NTT algorithm, the first step requires performing an NTT along the first dimension (e.g., column or row) of the matrix, while the fourth step requires performing an NTT along the second dimension (e.g., row or column). The aforementioned embodiments require explicit data manipulation to achieve this row-column perspective conversion. The method provided in this specification utilizes the difference in data distribution between the first and second multiplier slices in the MMA operation to achieve implicit transposition. Specifically, when the first MMA instruction is called to perform the first step of the four-step NTT, the input data matrix is ​​mapped to the second multiplier slice, and the input data participates in the calculation according to column-major logic. When the second MMA instruction is called, the second matrix stored in the third register group is directly used as the first multiplier slice storing the row data, and the second matrix participates in the calculation according to row-major logic, thus achieving the fourth step of the four-step NTT.

[0019] In this way, by changing the "role" of data in the computation process (using the data in the third register group storing the accumulated result slice as the first multiplier slice of the second MMA instruction), matrix transposition (implicit transposition) is achieved at the computational level. During implicit transposition, the data does not move within the third register group, thus eliminating the need for additional data movement or memory access throughout the process. This avoids the double shared memory traffic caused by transposition in the aforementioned methods, thereby significantly reducing memory access latency, improving the arithmetic strength of the algorithm, and achieving a highly efficient acceleration of NTT operations.

[0020] It should be understood that the above general description and the following detailed description are exemplary and explanatory only, and are not intended to limit this specification. Attached Figure Description

[0021] The accompanying drawings, which are incorporated in and form part of this specification, illustrate embodiments consistent with this specification and, together with the description, serve to explain the principles of this specification.

[0022] Figure 1 This is a flowchart illustrating a method for executing NTT via the tensor core of a GPU, as shown in this specification.

[0023] Figure 2A This is another flowchart illustrating the execution of the NTT method via the GPU's Tensor Core, as shown in this specification.

[0024] Figure 2B This is a schematic diagram illustrating the execution of the NTT method via the tensor core of a GPU, as shown in this specification.

[0025] Figure 2C This is another flowchart illustrating the execution of the NTT method via the GPU's tensor core, as shown in this specification.

[0026] Figure 2D This is a schematic diagram illustrating another method of performing the NTT method via the GPU's Tensor Core, as shown in this specification.

[0027] Figure 3 This is a hardware structure diagram of a computer device shown in this specification. Detailed Implementation

[0028] Exemplary embodiments will now be described in detail, examples of which are illustrated in the accompanying drawings. When the following description relates to the drawings, unless otherwise indicated, the same numerals in different drawings denote the same or similar elements. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with one or more embodiments of this specification. Rather, they are merely examples of apparatuses and methods consistent with some aspects of one or more embodiments of this specification as detailed in the appended claims.

[0029] It should be noted that the steps of the corresponding methods are not necessarily performed in the order shown and described in this specification in other embodiments. In some other embodiments, the methods may include more or fewer steps than described in this specification. Furthermore, a single step described in this specification may be broken down into multiple steps in other embodiments; and multiple steps described in this specification may be combined into a single step in other embodiments.

[0030] Fully homomorphic encryption (FHE) is a cryptographic technique that allows algebraic operations such as addition and multiplication to be performed directly on the ciphertext. The decrypted result is identical to the result of performing the same operations on the plaintext. This technique effectively protects data privacy and security during processing, and therefore has broad application prospects in sensitive fields such as cloud computing, medical data analysis, and financial risk control.

[0031] In fully homomorphic encryption schemes, the core operations involve the multiplication of high-dimensional polynomials. A polynomial is an algebraic expression composed of variables and coefficients through addition, subtraction, and multiplication operations, typically represented as: The form is , where x is a variable. The coefficients are denoted as C(x) and B(x). The multiplication of two polynomials A(x) and B(x) is represented as C(x) = A(x) ⋅ B(x). The coefficients of the resulting polynomial are calculated by discrete convolution of the input polynomial coefficients. Directly calculating polynomial multiplication has high time complexity. Moreover, as the security parameter increases, the polynomial order N is usually high, resulting in a huge performance overhead for direct calculation, becoming a major bottleneck restricting system efficiency.

[0032] To accelerate polynomial multiplication, existing techniques typically employ the Number Theoretic Transform (NTT). NTT is a special type of Discrete Fourier Transform (DFT) defined on the ring of integers. It utilizes number theory properties to transform a polynomial from a coefficient representation to a point-valued representation. Point-valued representation involves selecting a specific set of numerical values, substituting them into the polynomial, calculating and recording the function values ​​at these points, thus uniquely identifying the polynomial using these point-value pairs.

[0033] According to the convolution theorem, polynomial multiplication corresponds to point-value multiplication after NTT (Network Transformation). Specifically, in coefficient representation, multiplying two polynomials requires complex convolution operations; however, in point-value representation, polynomial multiplication simplifies to direct multiplication of corresponding values. By first performing NTT to convert the data to point-value form, completing simple element-wise multiplication, and then performing inverse number theory transformation (Inverse NTT, INTT) to restore the result to coefficient form, the time complexity of polynomial multiplication can be reduced from... The computation time is reduced to O(NlogN). Although NTT significantly reduces the number of operations, NTT itself still involves a large amount of computation, making its time consumption dominate the entire homomorphic encryption process and constituting a performance bottleneck that restricts the practical deployment of homomorphic encryption technology.

[0034] In its hardware implementation, NTT employs a divide-and-conquer strategy, breaking down the long NTT transformation process into multiple levels of small-scale computations. The core operation at each level is called a "butterfly operation," which selects two elements from the data sequence and performs modular multiplication and addition operations with fixed coefficients to generate a new intermediate result. Because the data pairs processed by all butterfly operations within the same level do not overlap, these operations are logically independent and can be executed in parallel without conflict. This characteristic makes NTT highly suitable for accelerating processing through parallel hardware.

[0035] GPUs, with their massively parallel architecture, are often used to accelerate NTT (Network Tolerance) computations. Traditional methods primarily utilize Compute Unified Device Architecture (CUDA) cores for parallel processing. CUDA cores are the fundamental processing units within a GPU responsible for general-purpose logic operations. In traditional implementations, each CUDA thread is responsible for executing a portion of the butterfly operations in the NTT algorithm.

[0036] However, CUDA Core is a general-purpose computing unit designed based on scalar or vector architecture. This means it can only process a single data element (scalar) or a small one-dimensional array (vector) at a time, unlike dedicated hardware which can perform batch parallel processing on large blocks of two-dimensional matrix data. Under this execution model, each thread must strictly follow a "load-compute-store" data flow pattern when performing butterfly operations: First, the thread needs to load the two data elements involved in the operation and their corresponding twiddle factors from global or shared memory into a register; then, it performs modular multiplication and modular addition in the register; finally, it stores the intermediate results back into shared or global memory for the next level of operation. Since the latency of memory access in GPUs is much higher than that of arithmetic operations, this frequent data movement in each iteration causes the computing unit to often be in a state of waiting for data to be ready, limiting the overall performance improvement of NTT.

[0037] To improve efficiency, existing technologies introduce Tensor Cores to accelerate NTT (Network Transformation) operations. Tensor Cores are dedicated hardware processing units integrated into GPUs, whose core function is to accelerate large-scale matrix multiply-accumulate (MMA) operations. While Tensor Cores possess extremely high theoretical computing power, their efficient operation relies on continuous and well-organized large-scale matrix data input. However, the butterfly operation involves array operations, which are not well-suited to Tensor Cores. To adapt to the matrix operation characteristics of Tensor Cores, a four-step NTT algorithm is typically used to transform long sequences of one-dimensional polynomial transformations into two-dimensional matrix operations. For ease of understanding, the four-step NTT algorithm is described in detail here.

[0038] It should be noted that the four-step decomposition method described below originates from the Discrete Fourier Transform (DFT) over the complex domain, where the twitch factor is the complex root of unity. In contrast, the Number Theory Transform (NTT) used in scenarios such as fully homomorphic encryption is defined over an integer ring modulo q, with the twitch factor being the primitive root modulo q. All operations are performed within a finite integer domain and do not involve complex or floating-point numbers. Despite the different mathematical foundations, the structured approach of the four-step decomposition can be transferred to NTT for adaptation to hardware acceleration units such as Tensor Cores.

[0039] Suppose a tensor x(j) with n elements is transformed into another tensor y(s) with n elements via an NTT transformation, where the rotation factor is the transformation matrix. The transformation can then be written as: The transformation matrix is ​​related to the tensor size n, such as... , 0 ≤ s < n. When n can be decomposed into the product of two factors n1 and n2, n = n1 × n2, then j and s can be expressed as j = j1 + j2 * n1, s = s2 + s1 * n2. Furthermore, x and y are expressed in two-dimensional form: x(j) = x(j1, j2), y(s) = y(s2, s1), 0 ≤ j1 < n1-1, 0 ≤ j2 < n2-1, 0 ≤ s1 < n1-1, 0 ≤ s2 < n2-1. Typically, j1 and j2 are as close as possible, s1 and s2 are as close as possible, and n1 and n2 are as close as possible. Therefore, the above tensor transformation process can be represented as:

[0040] y(s2, s1) = .

[0041] Under the four-step decomposition method, we have:

[0042] The first step is to perform tensor transformations on n1 groups of n2 points, such as x1(j1, s2) = ;

[0043] The second step is to perform matrix transpose, such as x2(s2,j1) = x1(j1,s2);

[0044] The third step is to calculate the Hadamard product, such as x3(s2,j1) = x2(s2,j1). ;

[0045] The fourth step is to perform tensor transformations on n2 groups of n1 points, such as y(s2, s1) = .

[0046] Furthermore, in some instructions, the positions of steps two and three can be interchanged. Based on the above principles, it can be seen that when n, n1, and n2 are determined, the transformation matrices that serve as rotation factors can also be determined, such as... , , , All of these can be determined. Typically, the rotation factor can be related to at least one of the following: rotation type (such as NTT type, FFT type, etc.), transformation parameters (such as e, i, etc.), and number of elements (such as n).

[0047] Based on the above four-step decomposition principle, the transformation process of a vector with length N (i.e., the number of elements contained in the vector) can be viewed as N1×N2 dimensional, and decomposed into N1 tensor transformations each applied to N2 points, and N2 tensor transformations each applied to N1 points, thus performing a two-dimensional decomposition of the constructed computational unit. Here, N1×N2=N, and N, N1, and N2 are usually powers of 2. In optional embodiments, N1 and N2 can be made as close as possible to improve processing efficiency. Specifically, when N is an even power of 2, N1 can be equal to N2; for example, when N=4096, N1=N2=64. When N is an odd power of 2, one of N1 and N2 is the square root of 2N, and the other is the square root of N / 2; for example, when N=8192, N1=128, N2=64, and so on. It is understandable that when N is not a power of 2, methods such as padding with zeros can be used to find the smallest power of 2 greater than N.

[0048] The four-step NTT transforms the NTT into an operation between the rotation factor matrix and the input data matrix, thus facilitating acceleration using the tensor core used for matrix multiplication and accumulation.

[0049] Furthermore, the input data matrix of the tensor core is limited in size, while the order N of the polynomial is often large (from 8192 to 65536) when applied to homomorphic encryption, resulting in a large input data matrix, which is not conducive to directly using the tensor core for acceleration. Moreover, if only a single thread block is used to complete the NTT, the parallelism is low, and the on-chip resources (such as shared memory) of the GPU corresponding to a single thread block cannot support the completion of large-scale NTT.

[0050] Therefore, this specification employs a two-stage divide-and-conquer strategy for homomorphic encryption scenarios: a large NTT computation is split into two main stages, and the computational tasks of these two stages are distributed to multiple different GPU thread blocks for parallel processing. To ensure that the computational load of these two stages is as balanced as possible, we design the decomposition parameters N1 and N2 to be as close in size as possible. It should be noted that this divide-and-conquer strategy is not limited to homomorphic encryption scenarios, but can also be used in other scenarios requiring large-scale NTT computation.

[0051] To refine the computational granularity and better adapt to the hardware, the NTT decomposition scheme can be fixed as N1 = B×n, based on the size of the matrix processed by the tensor core. 13 And N2 = B×n 23Where B is a value determined by the size of the matrix processed by the tensor core. For example, for a tensor core that processes double-precision floating-point numbers, it processes the multiplication of an 8×4 matrix and a 4×8 matrix, so B can be set to 8*8=64. In this specification, the NTT calculation part with base B is defined as the internal NTT, and the internal NTT part is specifically handled by the tensor core.

[0052] The radix-B NTT is also implemented based on the aforementioned four-step NTT. In an optional implementation, during the radix-B NTT processing using the Tensor Core, the Tensor Core is first used to perform the NTT transformation in the first step of the four-step NTT, i.e., calculating the product of the rotation factor matrix and the input matrix using the Tensor Core, and storing the calculated data in the register corresponding to the thread bundle. Then, the second step of the four-step NTT is performed, which involves matrix transposition. This process generally involves loading the data from the register corresponding to the thread bundle into shared memory, and then reloading it from shared memory back into the register according to the transposed layout, thus completing the logical row and column swap. Next, the data stored in the register is processed to complete the third step of the four-step NTT, the Hadamard product, and the fourth step, the NTT transformation.

[0053] During transposition, data needs to be loaded from registers into shared memory, and after processing, the transposed data in shared memory is loaded back into registers. This requires frequent data transfer between registers and the relatively slow shared memory. Tensor cores will frequently be limited in processing speed due to waiting for data, and memory access latency will become the dominant performance bottleneck, preventing the powerful computing capabilities of tensor cores from being effectively utilized.

[0054] Based on this, this specification describes the NTT processing procedure described above. First, the multiple rotation factor matrices and the input data matrix stored in shared memory are loaded into multiple registers corresponding to the thread bundle. Then, the thread bundle uses the tensor core to execute the first MMA instruction, using the rotation factor matrix as the first multiplier slice corresponding to the row data and the input data matrix as the second multiplier slice corresponding to the column data, to obtain the matrix multiplication result of the rotation factor matrix and the input data matrix (the first matrix). Then, the second matrix is ​​obtained by using the Hadamard product of the first matrix and the second rotation factor matrix. The second matrix is ​​stored in the third register group as a slice of accumulated results. Then, without moving the data stored in the third register group, the thread bundle executes the second MMA instruction, using the data in the third register group as the first multiplier slice corresponding to the row data and the third rotation factor matrix as the second multiplier slice, to obtain the matrix multiplication result of the second matrix and the third rotation factor matrix (the third matrix). Finally, the third matrix is ​​written from the registers to shared memory.

[0055] In the four-step NTT algorithm, the first step requires performing an NTT along the first dimension (e.g., column or row) of the matrix, while the fourth step requires performing an NTT along the second dimension (e.g., row or column). The aforementioned embodiments require explicit data manipulation to achieve this row-column perspective conversion. The method provided in this specification utilizes the difference in data distribution between the first and second multiplier slices in the MMA operation to achieve implicit transposition. Specifically, when the first MMA instruction is called to perform the first step of the four-step NTT, the input data matrix is ​​mapped to the second multiplier slice, and the input data participates in the calculation according to column-major logic. When the second MMA instruction is called, the second matrix stored in the third register group is directly used as the first multiplier slice storing the row data, and the second matrix participates in the calculation according to row-major logic, thus achieving the fourth step of the four-step NTT.

[0056] In this way, by changing the "role" of data in the computation process (using the data in the third register group storing the accumulated result slice as the first multiplier slice of the second MMA instruction), matrix transposition (implicit transposition) is achieved at the computational level. During implicit transposition, the data does not move within the third register group, thus eliminating the need for additional data movement or memory access throughout the process. This avoids the double shared memory traffic caused by transposition in the aforementioned methods, thereby significantly reducing memory access latency, improving the arithmetic strength of the algorithm, and achieving a highly efficient acceleration of NTT operations.

[0057] Next, we will combine Figure 1 This specification describes a method for executing NTT using the tensor core of a GPU. Figure 1 As shown, it includes the following steps:

[0058] Step 101: Load the first rotation factor matrix, the second rotation factor matrix, the third rotation factor matrix, and the input data matrix from the shared memory into multiple registers corresponding to the thread bundle.

[0059] First, let's explain the various terms involved in step 101.

[0060] Shared memory refers to the on-chip static random access memory (SRAM) in a GPU that can be shared by all threads within a thread block. A thread block can include one or more warps. A warp is the smallest set of threads that the GPU schedules and executes, typically consisting of a fixed number of threads (e.g., 32 threads) that execute locksteps under the same instructions.

[0061] Thread registers are high-speed storage spaces allocated to each thread in the GPU, used to store computational variables and intermediate results, and cannot be directly accessed by other threads. Multiple registers corresponding to a thread bundle refer to the registers belonging to the multiple threads within that thread bundle.

[0062] The first, second, and third twisting factor matrices are multiple twisting factor matrices involved in the four-step NTT method. The first twisting factor matrix is ​​the first step in the four-step NTT, specifically the matrix used in column-oriented NTT. As shown in the formula above, x1(j1, s2) = The column-oriented NTT can be written as the product of the first rotation factor matrix and the input data matrix.

[0063] The second twitch factor matrix is ​​the twitch factor matrix used when performing the Hadamard product in the four-step NTT method. The third twitch factor matrix is ​​the twitch factor matrix used in the fourth step of the four-step NTT method, i.e., the matrix used in the row-oriented NTT.

[0064] The input data matrix is ​​obtained by decomposing the vector to be subjected to NTT into a two-dimensional matrix. Specifically, the input data matrix x(j) can be represented in matrix form as x(j) = x(j1, j2). Where n = n1 × n2, j = j1 + j2 * n1, 0 ≤ j1 < n1-1, 0 ≤ j2 < n2-1.

[0065] In step 101, the data required for the thread to execute the NTT method is loaded from shared memory into a register. This method requires the use of a tensor core to encrypt the matrix multiplication process. However, the GPU's tensor core and general-purpose computing units can only directly operate on data in registers and cannot directly process data in shared memory. Therefore, by loading the data from shared memory into the thread's private register, MMA instructions can be executed to perform matrix multiplication.

[0066] Furthermore, when loading data into registers, the loaded data needs to conform to the characteristics of each fragment. Therefore, a pre-analysis of the data layout can be performed before the above steps are executed. Specifically, this involves pre-analyzing how the data fragments of the GPU's Tensor Core matrix multiplication-accumulation (MMA) operation are distributed among the threads within the thread bundle. Taking an 8x4x8 dimension MMA operation supported at FP64 precision as an example, it is necessary to clarify how the data in input fragments A and B and accumulation fragment C are mapped to the private registers of the 32 threads, and how the result data of output fragment D is distributed among the threads. This analysis is the foundation for subsequent implementation of register-level data flow.

[0067] Accordingly, during data loading, based on the pre-analysis results, the thread bundle reads the input data required for NTT calculation with a base of B (e.g., B=64) from shared memory and loads this data into the private registers of each thread within the thread bundle. Details on how this loading is performed are provided below and will not be elaborated upon here.

[0068] Step 103: The thread beam executes the first MMA instruction to obtain the first matrix.

[0069] The first MMA instruction uses the data in the first register group storing the first rotation factor matrix as the first multiplier slice corresponding to the row data, and uses the data in the second register group storing the input data matrix as the second multiplier slice corresponding to the column data.

[0070] In addition, the first matrix can be stored in the third register group corresponding to the slice of the accumulated result.

[0071] This step can be performed by a tensor core, a GPU hardware unit specifically designed for MMA calculations. The MMA instruction is used to calculate D = A × B + C, where A, B, C, and D are all matrices. Specifically, the first multiplier slice corresponds to slice A, the second multiplier slice corresponds to slice B, and the accumulated result slice corresponds to slice D.

[0072] In an optional implementation, the aforementioned slice C can be a zero matrix, i.e., a matrix in which all elements are 0. Correspondingly, the zero matrix can be loaded into multiple registers corresponding to the thread bundle. In the first MMA instruction, the data in the fifth register group storing the zero matrix is ​​used as the addition slice. This addition slice is also the slice C mentioned above.

[0073] Both the first and second register groups include several registers, and the registers included in the first and second register groups do not overlap. According to the multiplier-sliced ​​storage method, the first register group includes registers corresponding to each of the several threads within a thread bundle. For example, if a thread bundle includes 32 threads, the first register group could include one register corresponding to each of the 32 threads, for a total of 32 registers. The second register group is similar, also including registers corresponding to each of the several threads within a thread bundle.

[0074] The first multiplier slice is the slice corresponding to the row data, that is, the first multiplier slice is stored first according to the row. For example, as shown in Figure 2, in the first multiplier slice shown in Figure 2, the registers used to store the first row of matrix data are the registers of threads 0, 1, 2 and 3, and the second row of data corresponds to the registers of threads 4, 5, 6 and 7 respectively.

[0075] The first multiplier slice is the slice corresponding to the column data, meaning the second multiplier slice is stored primarily by column. For example, as shown in Figure 2, in the second multiplier slice shown in Figure 2, the registers used to store the first column data are the registers of threads 0, 1, 2, and 3, and the registers used to store the second column data are the registers of threads 4, 5, 6, and 7.

[0076] The first rotation factor matrix corresponds to the rotation factor matrix used in the first step of the four-step NTT method. Step 103 completes the first step of the four-step NTT, namely, the NTT transformation in the first direction for the input data matrix.

[0077] In an optional implementation, after step 103 is completed, modulo reduction can also be performed, that is, the result is reduced by modulo q.

[0078] Step 105: Obtain the second matrix based on the Hadamard product of the first matrix and the second rotation factor matrix.

[0079] The second matrix is ​​stored in the third register group corresponding to the slice of the accumulated result.

[0080] The second twitch factor matrix corresponds to the twitch factor matrix used in the third step of the four-step NTT method mentioned above to calculate the Hadamard product. The Hadamard product is the element-wise multiplication of two matrices. The second matrix is ​​the result of the Hadamard product of the first matrix and the second twitch factor matrix. Through the operation in step 105, the process of calculating the Hadamard product in the third step of the four-step NTT method can be completed.

[0081] For step 105, step 105 may not be executed by the tensor core, but rather the thread bundle may use the GPU's general computing units to perform the above-mentioned Hadamard product operation.

[0082] Furthermore, both the first and second matrices can be stored in the third register group of the accumulated result slice, meaning that the above operations do not change the data storage location, and the data is not transferred between shared memory and registers.

[0083] Step 107: The thread beam executes the second MMA instruction, in which the data of the third register group is used as the first multiplier slice, and the data of the fourth register group storing the third rotation factor matrix is ​​used as the second multiplier slice, to obtain the third matrix.

[0084] The third rotation factor matrix corresponds to the rotation factor matrix used in the fourth step of the four-step NTT. Step 107 completes the fourth step of the four-step NTT, that is, the NTT transformation in the second direction is performed on the input data matrix.

[0085] Furthermore, similar to step 103, step 107 can be executed by the Tensor Core, whose MMA instructions are used to calculate D = A × B + C. Therefore, during data loading in step 101, the zero matrix can also be loaded into multiple registers corresponding to the thread bundle. In the second MMA instruction, the data in the fifth register group storing the zero matrix is ​​used as the addition slice, i.e., slice C.

[0086] In step 103 above, the first rotation factor matrix is ​​used as the first multiplier slice, and the input data matrix is ​​used as the second multiplier slice, performing a column-oriented NTT. In step 107, the second matrix is ​​used as the first multiplier slice, and the third rotation factor matrix is ​​used as the second multiplier slice. Thus, implicit transposition is achieved by utilizing the data distribution difference between input slices A and B in the MMA operation. Specifically, the data processed in the previous steps and still stored in the thread registers (located in slice D) is directly used as the input to the second MMA instruction. By treating this data as slice A when calling the MMA instruction (while in the first MMA instruction, the original input data matrix is ​​treated as slice B), matrix transposition is achieved at the computational level, without requiring additional data movement or memory access throughout the process.

[0087] Furthermore, the accumulated result slices are stored in row-major order, as shown in Figure 2. During the second MMA instruction, the second matrix is ​​used as the first multiplier slice corresponding to the row data. This ensures that the intermediate result (second matrix) is always stored in the register, eliminating the need for data movement between registers or moving data to shared memory, thus accelerating the NTT calculation process.

[0088] The third matrix can also be stored in a register as a slice of the accumulated result. In an optional implementation, after executing step 107, the third matrix can be modulo-reduced to obtain the final calculation result.

[0089] Step 109: Write the third matrix as the NTT result into shared memory.

[0090] Specifically, the third matrix stored in the register can be written back to the shared memory, completing a thread-level internal NTT accelerated calculation.

[0091] Through the above steps, the entire internal NTT computation with base B, except for initial data loading and final result writing back, involves all intermediate calculations, data processing, and cross-thread data exchange within the thread bundle's registers. In particular, the implicit transpose method in step 107 cleverly utilizes the data flow characteristics of the tensor core, avoiding the double shared memory traffic caused by transposition in traditional methods. This significantly reduces memory access latency, improves the algorithm's arithmetic strength, and achieves highly efficient acceleration of NTT operations.

[0092] In one alternative implementation, the matrix multiplication dimension that the tensor core can perform may not match the size of the radix B. For example, the tensor core can compute 8×4×8 matrix multiplications, that is, multiplying an 8×4 matrix by a 4×8 matrix, with a radix size B of 64. This results in the original input data matrix, the rotation factor matrix, and the final output result (the third matrix) all being 8×8 in size.

[0093] To enable the tensor kernel to perform the above calculations, the input data matrix and rotation factor matrix can be split to complete the calculations.

[0094] Specifically, for the first MMA instruction in step 103, the input data matrix is ​​split into a first sub-matrix and a second sub-matrix by row, and stored in the first sub-register group and the second sub-register group in the second register group, respectively. Correspondingly, the first rotation factor is split into a third sub-matrix and a fourth sub-matrix by column, and stored in the third sub-register group and the fourth sub-register group in the first register group, respectively.

[0095] The first MMA instruction consists of the third MMA instruction and the fourth MMA instruction.

[0096] First, the product of the third submatrix and the first submatrix (i.e., the first intermediate matrix) is calculated using the third MMA instruction, where the zero matrix can be used as an addition slice. Specifically, the thread executes the third MMA instruction, using the data from the third sub-register group as the first multiplier slice, the data from the first sub-register group as the second multiplier slice, and the data from the fifth register group storing the zero matrix as the addition slice, to obtain the first intermediate matrix; the first intermediate matrix is ​​stored in the sixth register group corresponding to the accumulated result slice.

[0097] Then, the product of the fourth submatrix and the first submatrix is ​​calculated using the fourth MMA instruction. Furthermore, the first intermediate matrix is ​​used as an addition slice, allowing the matrix multiplication and merging of the two matrix multiplication results to be completed via the fourth MMA instruction. Specifically, the thread executes the fourth MMA instruction, where the data from the fourth sub-register group is used as the first multiplier slice, the data from the second sub-register group is used as the second multiplier slice, and the data from the sixth register group is used as the addition slice to obtain the first matrix.

[0098] In one optional implementation, the input data matrix comprises m rows, with the first m / 2 rows corresponding to the first submatrix and the last m / 2 rows corresponding to the second submatrix; the first rotation factor matrix comprises m columns, with the first m / 2 columns corresponding to the third submatrix and the last m / 2 columns corresponding to the fourth submatrix. This facilitates reading and loading into the register. In the case of B=64 mentioned earlier, m can be 8.

[0099] The second MMA instruction in step 107 is similar to the first MMA instruction. The input second matrix to be calculated for NTT can be split into a fifth sub-matrix and a sixth sub-matrix according to columns and stored in the fifth sub-register group and the sixth sub-register group in the third register group, respectively. Correspondingly, the third rotation factor is split into a seventh sub-matrix and an eighth sub-matrix according to columns and stored in the seventh sub-register group and the eighth sub-register group in the fourth register group, respectively.

[0100] The second MMA instruction consists of the fifth MMA instruction and the sixth MMA instruction.

[0101] First, the product of the fifth and seventh sub-matrices (i.e., the second intermediate matrix) is calculated using the fifth MMA instruction, where the zero matrix is ​​used as the addition slice. Specifically, the thread executes the fifth MMA instruction, using the data from the fifth sub-register group as the first multiplier slice, the data from the seventh sub-register as the second multiplier slice, and the data from the fifth register group storing the zero matrix as the addition slice, to obtain the second intermediate matrix; the second intermediate matrix is ​​stored in the seventh register group corresponding to the slice of the accumulated result.

[0102] Then, the product of the sixth and eighth submatrices is calculated using the sixth MMA instruction, and this product is combined with the second intermediate matrix. Specifically, the thread executes the sixth MMA instruction, where the data from the sixth sub-register group is used as the first multiplier slice, the data from the eighth sub-register is used as the second multiplier slice, and the data from the seventh register group is used as the addition slice to obtain the third matrix. The seventh sub-register stores the second intermediate matrix.

[0103] In one alternative implementation, the second matrix comprises m columns, with odd-numbered columns corresponding to the fifth submatrix and even-numbered columns corresponding to the sixth submatrix; the third rotation factor matrix comprises m rows, with odd-numbered rows corresponding to the seventh submatrix and even-numbered rows corresponding to the eighth submatrix. In the case of B=64 mentioned earlier, m can be 8.

[0104] The reason for dividing the results into odd and even numbers is due to the storage characteristics of the accumulated result slices. See Figure 2 for details. In the accumulated result slices, the data stored in the first row corresponds to the two registers of threads 0, 1, 2, and 3, respectively. The data distribution of the first multiplier slice means that it can only include one register corresponding to each thread. Therefore, for the second matrix, the fifth and sixth sub-matrices are obtained by intermittent reading. Correspondingly, the third rotation factor matrix is ​​also split in a similar way to ensure the accuracy of matrix multiplication.

[0105] In another alternative implementation, there may be a mismatch between the number of bits in the input data and the number of bits that the tensor core can process. For example, the input data may be 32 bits, while the tensor core supports double-precision floating-point arithmetic. It should be noted that although double-precision floating-point numbers include 64 bits, several bits are used to represent the sign and exponent, leaving only 53 bits actually usable to represent the data. Therefore, when executing MMA instructions, if the elements in the matrix (e.g., the input data matrix, the various twitch factor matrices, and the second matrix) are 32 bits, the product of two 32-bit data points might result in a 64-bit product, exceeding the range of values ​​that double-precision floating-point numbers can represent.

[0106] Therefore, in the above case, one of the matrices in the input matrix can be split into the high 16 bits and the low 16 bits, and then the multiplication of the 16-bit data and the 32-bit data can be calculated. The result is 48 bits, which does not exceed the range of values ​​that double-precision floating-point numbers can represent.

[0107] Here, you can choose to split the rotation factor matrix, since the rotation factor matrix may be reused when calculating the internal NTT multiple times. Splitting the reused rotation factor matrix will improve processing speed.

[0108] Here, we will take the example of a tensor core that supports double-precision floating-point operations, where the length of the elements in the tensor core input data matrix, the first rotator matrix, and the third rotator matrix is ​​32 bits, to explain the first MMA instruction and the second MMA instruction mentioned above.

[0109] For the first MMA instruction in step 103, the aforementioned first register used to store the first rotation factor matrix can store the high 16 bits of data corresponding to the first rotation factor. The first MMA instruction is used to calculate the product of the high 16 bits of the first rotation factor and the input data matrix with 32-bit elements.

[0110] Correspondingly, the thread bundle can also execute the seventh MMA instruction to use the data in the eighth register group corresponding to the first twitch factor (lower 16 bits) as the first multiplier slice and the data in the second register group as the second multiplier slice to obtain the fourth matrix. The seventh MMA instruction can be used to calculate the product of the lower 16 bits of the first twitch factor and the 32-bit input data matrix (the fourth matrix).

[0111] Then, the two product results can be bit-merged, that is, the first matrix and the fourth matrix can be bit-merged. The bit-merged result is then updated to the first matrix, so that the first matrix can be used to continue to complete the subsequent steps.

[0112] The second MMA instruction in step 107 is similar to the first MMA instruction. The fourth register group, which was previously used to store the third twisting factor matrix, stores the high 16 bits of data corresponding to the third twisting factor. The second MMA instruction is used to calculate the product of the second matrix and the high 16 bits of the third twisting factor.

[0113] The second matrix can be 32-bit. Although the NTT method mentioned earlier may yield 64-bit data, since the NTT is a linear transformation performed over a finite field, the modulo reduction of the second matrix described above can map the second matrix to a 32-bit width.

[0114] Correspondingly, the thread bundle can also execute the eighth MMA instruction, using the data from the third register group as the first multiplier slice and the data from the ninth register group storing the lower 16 bits corresponding to the third twitch factor as the second multiplier slice, to obtain the fifth matrix. The product of the second matrix and the lower 16 bits of the third twitch factor (the fifth matrix) can then be calculated using the eighth MMA instruction.

[0115] Then, the product of the second MMA instruction and the eighth MMA instruction can be bit-merged, that is, the third matrix and the fifth matrix can be bit-merged, and the result of the bit-merging can be updated to the third matrix.

[0116] The method described in this specification will now be explained with reference to a specific embodiment. This embodiment uses the application of this method to the calculation of polynomial multiplication of order N in a homomorphic encryption scheme, where N is a power of 2, as an example. As mentioned earlier, the NTT process can be divided into an inner NTT and an outer NTT, where each inner NTT has a base of 64 and is processed using a tensor core. Each NTT is processed by a single thread bundle, so the polynomial multiplication is executed by N / p thread bundles, which can be processed in parallel. Any thread bundle performs a base-p NTT calculation. In this embodiment, p is 64. Here, p has the same meaning as the base B mentioned earlier.

[0117] In this method, a tensor core that supports double-precision floating-point operations is used as an example. The input data matrix, the first rotation factor matrix, and the third rotation factor matrix are all 32-bit data.

[0118] For each internal NTT, the execution process is as follows: Figure 2A As shown in Figure 2C. Before performing the following steps, data loading can be performed first, loading the data needed for the computation from shared memory into the private registers of each thread in the thread bundle.

[0119] First, such as Figure 2AAs shown, the internal NTT input is an 8×8 matrix. The first twiddle factor matrix is ​​split into high 16 bits and low 16 bits, referred to as TFM (Twiddle Factor Matrix) in the diagram. Then, four MMA instructions, MMA1a, MMA1b, MMA2a, and MMA2b, are executed. MMA1a and MMA1b are used to calculate the matrix multiplication result of the high 16 bits of the first twiddle factor matrix and the input data matrix. Correspondingly, in MMA1a and MMA1b, the registers corresponding to each slice store the high 16 bits of the TFM data and the complete 32-bit data of the input data matrix. MMA2a and MMA2b are used to calculate the matrix multiplication result of the low 16 bits of the first twiddle factor matrix and the input data matrix. Correspondingly, in MMA2a and MMA2b, the registers store the low 16 bits of the TFM data and the complete 32-bit data of the input data matrix.

[0120] The execution methods of the four MMA instructions, MMA1a, MMA1b, MMA2a, and MMA2b, are as follows: Figure 2B As shown. Figure 2B In the matrix, the numbers represent the thread IDs within a warp that performed the access operation; in other words, the thread number within the warp corresponding to the register storing the data at that location. The dashed boxes correspond to unaccessed elements. So, from... Figure 2B As can be seen, the MMA1a and 2a instructions are used to calculate the product of the first four columns of the first twist factor matrix and the first four rows of the input data matrix; this product is called the intermediate result. The MMA1b and 2b instructions are used to calculate the product of the last four columns of the first twist factor matrix and the last four rows of the input data matrix, and they sum this product with the intermediate result to obtain the first step of the four-step NTT calculation, i.e., the NTT-8 result. According to the matrix multiplication rules, through the above splitting, it can still be guaranteed that the NTT-8 result = the unsplit TFM × the unsplit input data matrix.

[0121] like Figure 2B As shown, among these four MMA instructions, TFM serves as slice A (i.e., the first multiplier slice), and the storage rule corresponding to slice A is row-major. For example... Figure 2B In the TFM, the data in the first row is stored in the registers of threads 0, 1, 2, and 3, respectively. The input data matrix is ​​used as slice B (i.e., the second multiplier slice), and slice B is stored in column-major order, for example... Figure 2B In the input data matrix, the data in the first column is stored in the registers of threads 0, 1, 2 and 3 respectively.

[0122] During the calculation, in MMA1a and MMA1b, the zero matrix is ​​used as slice C. In MMA1b and 2b, the intermediate results obtained from MMA1a and MMA1b are used as slice C. This allows matrix multiplication and summation to be performed using MMA instructions, ensuring the accuracy of the output results. Furthermore, slice D and slice C have the same storage rules, meaning the register set storing the intermediate results can be directly used as slice C.

[0123] like Figure 2A As shown, after executing the four MMA instructions MMA1a, MMA1b, MMA2a, and MMA2b, the corresponding NTT-8 results for the high-order bits and the low-order bits are obtained. Figure 2A Step 3 in the process can perform a bit-merge on the two NTT-8 results, that is, merge the NTT-8 results corresponding to the high-order bits and the NTT-8 results corresponding to the low-order bits to restore the complete integer. For example... Figure 2A Section 4.1 allows for modulo reduction (ModRed) on the merged data, which involves performing a reduction operation on the merged result by modulo q to map the integer operation result back to the valid representation range of a finite integer ring, thereby ensuring the closure of the algebraic structure and preventing numerical overflow.

[0124] Then, as Figure 2A As shown in section 4.2, element-wise multiplication (EWMult) can be performed on the modulo reduction result. That is, as mentioned in step 105, the first matrix is ​​multiplied digit-wise with the second rotation factor (calculating the Hadamard product) to obtain the second matrix. Figure 2A The text also mentions 4.3 implicit transpose, which utilizes the data distribution differences between input fragments A and B in MMA operations.

[0125] After execution Figure 2A After the first step, you can proceed with the next step. Figure 2C The steps are shown. (With) Figure 2A similar, Figure 2A This corresponds to the four-step NTT process: the first step calculates the NTT in the first direction, the second step is the transpose, and the third step calculates the Hadamard product. Figure 2B This corresponds to the calculation of the NTT in the second direction in the fourth step of the four-step NTT process.

[0126] like Figure 2C As shown, to complete the fourth step of the four-step NTT, four MMA instructions, MMA3a, MMA3b, MMA4a, and MMA4b, need to be executed. Among them, MMA3a and MMA3b are used to calculate the matrix multiplication result of the high 16 bits of the second matrix and the third twisting factor matrix, and MMA4a and MMA4b are used to calculate the matrix multiplication result of the low 16 bits of the second matrix and the third twisting factor matrix.

[0127] like Figure 2D As shown, with Figure 2D The method shown differs in that, in order to achieve implicit transpose without moving data, it is necessary to obtain partition A and partition B based on the storage method of the second matrix. Specifically, as follows... Figure 2B As shown, in the second matrix, the data in the first row is stored in the registers of thread 001122233 respectively. Figure 2B The diagram shows the storage method of the first matrix. Modulo reduction and other operations do not change the data storage method, therefore the storage method of the second matrix is ​​the same as that of the first matrix. In slice A, the four data items in the first row need to be stored in the registers corresponding to the threads numbered 0, 1, 2, and 3, respectively. Therefore, for the second matrix, the following method is used... Figure 2D The interval method is used to obtain fragment A, and correspondingly, according to Figure 2D The interval method shown is used to obtain slice B. This allows for implicit transposition without moving the data in the registers, and the correct calculation results can be obtained based on this.

[0128] like Figure 2C As shown, after the four MMA instructions MMA3a, MMA3b, MMA4a and MMA4b are calculated, bit combining and modulo reduction can be performed to obtain the output matrix of the internal NTT, which is an 8×8 matrix.

[0129] After obtaining the internal NTT output matrix, the data can be written back to shared memory from the register.

[0130] The method described above employs a computational pipeline where all intermediate steps, from data loading to result writing back, are completed within the thread's private registers. This includes MMA calculations, bit merging, modulo reduction, multiplication with a twiddle factor, and implicit transpose. This maximizes the temporal and spatial locality of data, minimizes data access latency, and makes the computation process virtually unrestricted by memory bandwidth.

[0131] Furthermore, instead of achieving matrix transposition through conventional data movement (such as using shfl instructions or shared memory), it leverages the hardware characteristic of different data layouts for input slices A and B in thread registers during Tensor Core MMA operations. By changing the "role" of the data in the computation flow (i.e., directly using the output slice D of the previous stage as the input slice A of the next stage), the transposition effect is achieved at both the logical and computational levels, while the physical data never leaves the registers. This completely eliminates the shared memory (SMEM) access bottleneck caused by matrix transposition in traditional NTT divide-and-conquer algorithms. This not only saves a significant number of memory access cycles but also eliminates double the SMEM bandwidth usage.

[0132] This approach leverages the powerful computational throughput of the tensor core while maintaining the overall time complexity of the algorithm at an efficient O(N log N).

[0133] Corresponding to the embodiments of the methods described above, this specification also provides embodiments of the apparatus and the computer equipment on which it is applied.

[0134] This specification provides a device for executing NTT via the tensor core of a GPU, comprising:

[0135] The loading module is used to load the first rotation factor matrix, the second rotation factor matrix, the third rotation factor matrix, and the input data matrix from shared memory into multiple registers corresponding to the thread bundle;

[0136] The first MMA instruction execution module is used to execute the first matrix multiplication-addition MMA instruction in a thread bundle. The data of the first register group storing the first rotation factor matrix is ​​used as the first multiplier slice corresponding to the row data, and the data of the second register group storing the input data matrix is ​​used as the second multiplier slice corresponding to the column data to obtain the first matrix. The second matrix is ​​obtained by the Hadamard product of the first matrix and the second rotation factor matrix. The second matrix is ​​stored in the third register group corresponding to the accumulation result slice.

[0137] The second MMA instruction execution module is used to execute the second MMA instruction in the thread bundle. The data of the third register group is used as the first multiplier slice, and the data of the fourth register group storing the third rotation factor matrix is ​​used as the second multiplier slice to obtain the third matrix.

[0138] A storage module is used to write the third matrix as an NTT result into shared memory.

[0139] The specific implementation process of the functions and roles of each module in the above device can be found in the implementation process of the corresponding steps in the above method, and will not be repeated here.

[0140] For the device embodiments, since they basically correspond to the method embodiments, the relevant parts can be referred to in the description of the method embodiments. The device embodiments described above are merely illustrative. The modules described as separate components may or may not be physically separate, and the components shown as modules may or may not be physical modules, that is, they may be located in one place or distributed across multiple network modules. Some or all of the modules can be selected to achieve the purpose of the solution in this specification according to actual needs. Those skilled in the art can understand and implement this without creative effort.

[0141] like Figure 3 As shown, Figure 3A hardware structure diagram of a computer device is shown. The device may include: a processor 1010, a memory 1020, an input / output interface 1030, a communication interface 1040, and a bus 1050. The processor 1010, memory 1020, input / output interface 1030, and communication interface 1040 are internally connected to each other via the bus 1050.

[0142] The processor 1010 can be implemented using a general-purpose GPU, CPU (Central Processing Unit), microprocessor, application-specific integrated circuit (ASIC), or one or more integrated circuits, and is used to execute relevant programs to implement the technical solutions provided in the embodiments of this specification. The processor implements the above-described methods by running executable instructions.

[0143] The memory 1020 for storing processor-executable instructions can be implemented in the form of ROM (Read Only Memory), RAM (Random Access Memory), static storage device, dynamic storage device, etc. The memory 1020 can store the operating system and other applications. When the technical solutions provided in the embodiments of this specification are implemented through software or firmware, the relevant program code is stored in the memory 1020.

[0144] The input / output interface 1030 is used to connect input / output modules to realize information input and output. Input / output modules can be configured as components within the device (not shown in the figure) or externally connected to the device to provide corresponding functions. Input devices may include keyboards, mice, touchscreens, microphones, various sensors, etc., while output devices may include displays, speakers, vibrators, indicator lights, etc.

[0145] The communication interface 1040 is used to connect a communication module (not shown in the figure) to enable communication between this device and other devices. The communication module can communicate via wired means (such as USB, Ethernet cable, etc.) or wireless means (such as mobile network, WIFI, Bluetooth, etc.).

[0146] Bus 1050 includes a pathway for transmitting information between various components of the device, such as processor 1010, memory 1020, input / output interface 1030, and communication interface 1040.

[0147] It should be noted that although the above-described device only shows the processor 1010, memory 1020, input / output interface 1030, communication interface 1040, and bus 1050, in specific implementations, the device may also include other components necessary for normal operation. Furthermore, those skilled in the art will understand that the above-described device may only include the components necessary for implementing the embodiments of this specification, and not necessarily all the components shown in the figures.

[0148] This specification also provides a computer program product that, when executed by a processor, implements the aforementioned method of executing NTT via the tensor core of a GPU.

[0149] This specification also provides a computer-readable storage medium storing a computer program that, when executed by a processor, implements the aforementioned method of executing NTT via the tensor core of a GPU.

[0150] Computer-readable media includes both permanent and non-permanent, removable and non-removable media that can store information using any method or technology. Information can be computer-readable instructions, data structures, modules of programs, or other data. Examples of computer storage media include, but are not limited to, phase-change memory (PRAM), static random access memory (SRAM), dynamic random access memory (DRAM), other types of random access memory (RAM), read-only memory (ROM), electrically erasable programmable read-only memory (EEPROM), flash memory or other memory technologies, CD-ROM, digital versatile optical disc (DVD) or other optical storage, magnetic tape, magnetic magnetic disk storage or other magnetic storage devices, or any other non-transferable medium that can be used to store information accessible by a computing device. As defined herein, computer-readable media does not include transient computer-readable media, such as modulated data signals and carrier waves.

[0151] It should also be noted that the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Without further limitation, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes said element.

[0152] The foregoing has described specific embodiments of this specification. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps recited in the claims may be performed in a different order than that shown in the embodiments and may still achieve the desired result. Furthermore, the processes depicted in the drawings do not necessarily require the specific or sequential order shown to achieve the desired result. In some embodiments, multitasking and parallel processing are possible or may be advantageous.

[0153] The user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, data stored, data displayed, etc.) involved in this manual are all information and data authorized by the user or fully authorized by all parties. The collection, use and processing of related data shall comply with the relevant laws, regulations and standards of the relevant countries and regions, and corresponding operation portals shall be provided for users to choose to authorize or refuse.

Claims

1. A method for performing number-theoretic transformations (NTTs) via the tensor core of a GPU, comprising: Load the first rotation factor matrix, the second rotation factor matrix, the third rotation factor matrix, and the input data matrix from the shared memory into multiple registers corresponding to the thread bundle; The thread executes the first matrix multiply-add MMA instruction, in which the data of the first register group storing the first rotation factor matrix is ​​used as the first multiplier slice corresponding to the row data, and the data of the second register group storing the input data matrix is ​​used as the second multiplier slice corresponding to the column data, to obtain the first matrix; The second matrix is ​​obtained by taking the Hadamard product of the first matrix and the second rotation factor matrix; the second matrix is ​​stored in the third register group corresponding to the slice of the accumulated result. The thread executes the second MMA instruction, in which the data of the third register group is used as the first multiplier slice, and the data of the fourth register group storing the third rotation factor matrix is ​​used as the second multiplier slice to obtain the third matrix; The third matrix is ​​written into shared memory as an NTT result.

2. The method according to claim 1, wherein, Also includes: Load the zero matrix into multiple registers corresponding to the thread bundle; In the first MMA instruction and the second MMA instruction, the data in the fifth register group storing the zero matrix is ​​used as the addition slice.

3. The method according to claim 1, wherein, The input data matrix is ​​split into a first sub-matrix and a second sub-matrix according to the rows, and stored in the first sub-register group and the second sub-register group in the second register group, respectively; the first rotation factor is split into a third sub-matrix and a fourth sub-matrix according to the columns, and stored in the third sub-register group and the fourth sub-register group in the first register group, respectively. The thread bundle executes the first matrix multiply-add MMA instruction, including: The thread bundle executes the third MMA instruction, wherein the data of the third sub-register group is used as the first multiplier slice, the data of the first sub-register group is used as the second multiplier slice, and the data of the fifth register group storing the zero matrix is ​​used as the addition slice to obtain the first intermediate matrix; the first intermediate matrix is ​​stored in the sixth register group corresponding to the accumulation result slice; The thread bundle executes the fourth MMA instruction, in which the data of the fourth sub-register group is used as the first multiplier slice, the data of the second sub-register group is used as the second multiplier slice, and the data of the sixth register group is used as the addition slice to obtain the first matrix.

4. The method according to claim 3, wherein, The input data matrix consists of m rows, with the first m / 2 rows corresponding to the first submatrix and the last m / 2 rows corresponding to the second submatrix; the first rotation factor matrix consists of m columns, with the first m / 2 columns corresponding to the third submatrix and the last m / 2 columns corresponding to the fourth submatrix.

5. The method according to claim 1, wherein, The second matrix is ​​split into a fifth sub-matrix and a sixth sub-matrix according to columns and stored in the fifth sub-register group and the sixth sub-register group in the third register group, respectively; the third rotation factor is split into a seventh sub-matrix and an eighth sub-matrix according to rows and stored in the seventh sub-register group and the eighth sub-register group in the fourth register group, respectively. The thread bundle executes the second MMA instruction, including: The thread bundle executes the fifth MMA instruction, wherein the data of the fifth sub-register group is used as the first multiplier slice, the data of the seventh sub-register is used as the second multiplier slice, and the data of the fifth register group storing the zero matrix is ​​used as the addition slice to obtain the second intermediate matrix; the second intermediate matrix is ​​stored in the seventh register group corresponding to the accumulation result slice; The thread bundle executes the sixth MMA instruction, in which the data of the sixth sub-register group is used as the first multiplier slice, the data of the eighth sub-register is used as the second multiplier slice, and the data of the seventh register group is used as the addition slice to obtain the third matrix.

6. The method according to claim 5, wherein, The second matrix consists of m columns, with odd-numbered columns corresponding to the fifth submatrix and even-numbered columns corresponding to the sixth submatrix; the third rotation factor matrix consists of m rows, with odd-numbered rows corresponding to the seventh submatrix and even-numbered rows corresponding to the eighth submatrix.

7. The method according to claim 1, wherein, The tensor core is a tensor core that supports double-precision floating-point operations; the length of the elements in the input data matrix and the first rotation factor matrix is ​​32 bits; The first register group stores the high 16 bits of data corresponding to the first rotation factor; The method further includes: The thread executes the seventh MMA instruction, using the data of the eighth register group corresponding to the first rotation factor (lower 16 bits) as the first multiplier slice, and the data of the second register group as the second multiplier slice, to obtain the fourth matrix; Perform bit merging on the first and fourth matrices, and update the first matrix with the result of the bit merging.

8. The method according to claim 7, wherein, The length of each element in the third rotation factor matrix is ​​32 bits; the fourth register group stores the high 16 bits of data corresponding to the third rotation factor. The method further includes: The thread executes the eighth MMA instruction, using the data in the third register group as the first multiplier slice and the data in the ninth register group that stores the lower 16 bits corresponding to the third twitch factor as the second multiplier slice, to obtain the fifth matrix; Perform bit merging on the third and fifth matrices, and update the third matrix with the result of the bit merging.

9. The method according to claim 1, wherein, The method is used to calculate polynomial multiplication of order N in homomorphic encryption schemes, where N is a power of 2; the polynomial multiplication is executed by N / p thread bundles, and any thread bundle performs NTT calculation of radix-p.

10. A computer device, comprising: processor; Memory used to store processor-executable instructions; The processor implements the method as described in any one of claims 1-9 by executing the executable instructions.

Images