Method for solving high-performance linear equation set running on NPU

Abstract

The invention discloses a method for solving a high-performance linear equation set running on an NPU, and belongs to the technical field of numerical calculation and high-performance calculation. Aiming at a symmetric positive definite sparse linear system generated by discretization of a two-dimensional finite element method, firstly, a global stiffness matrix and a load vector are obtained, and the global stiffness matrix and the load vector are converted into a compressed sparse row (CSR) format adaptive to NPU; before sparse matrix-vector multiplication is executed, the number of non-zero elements is analyzed according to a row offset array, a calculation task is dynamically divided into a plurality of calculation groups, short rows are merged, and long rows are segmented, so that the load of each calculation core of the NPU is balanced; meanwhile, intermediate data such as a solution vector and a residual vector are distributed and permanently stored at an NPU device end, matrix-vector multiplication, vector dot product and vector updating required by conjugate gradient iteration are executed at the device end in a closed loop mode, and a final solution is transmitted back to a host only after iteration is finished. Through resident and optimized task scheduling of the data equipment end, the data transmission and synchronization overhead between the host and the equipment is effectively reduced, and the solution performance of the sparse linear system on the NPU is remarkably improved.

Description

Technical Field

[0001] This invention relates to the fields of numerical computation and high-performance computing, specifically to a method for solving sparse linear equations for NPUs, which is particularly suitable for the rapid solution of symmetric positive definite sparse linear systems generated by Laplace equations and other systems discretized by the two-dimensional finite element method on devices such as Ascend NPUs. Background Technology

[0002] The finite element method (FEM) is a discrete numerical algorithm widely used in engineering simulation, structural analysis, heat conduction and electric field analysis. When solving partial differential equations using the FEM, a symmetric positive definite sparse linear system of equations is typically obtained. ,in The global stiffness matrix. For load vector, Let be the vector to be solved. For large-scale sparse systems generated by fine-grid discretization, direct methods often struggle to balance memory and computational costs, so iterative methods, such as the conjugate gradient method, are frequently used.

[0003] In recent years, AI-specific accelerators such as NPUs have demonstrated strong parallel capabilities in matrix computation. However, efficiently mapping the iterative solution process of sparse linear equations to NPUs still faces the following challenges: First, NPUs are acceleration chips designed for AI, and their support in traditional high-performance computing is insufficient, with inadequate default operator support, making it difficult to accelerate linear equation solutions using methods like conjugate gradients. Second, the irregular non-zero distribution of sparse matrices can easily lead to uneven computational load and low memory access efficiency. Third, the iteration process involves various operators such as sparse matrix-vector multiplication, vector dot product, and AXPY. Frequent transmission of intermediate vectors or scalars between the host and device can easily cause synchronization and data transfer overhead to become bottlenecks. Fourth, the lack of a fusion and continuous execution mechanism between dot product reduction, scalar broadcasting, and vector updates can introduce additional kernel scheduling and startup overhead.

[0004] Therefore, a high-performance linear equation solving method is needed that can keep sparse matrices and intermediate vectors resident on the NPU side and reasonably divide and schedule sparse computing tasks, so as to reduce host-device data transmission and synchronization overhead and improve NPU computing power utilization. Summary of the Invention

[0005] I. Purpose of the Invention

[0006] The purpose of this invention is to provide a high-performance method for solving linear equations running on an NPU. By constructing a sparse format for the stiffness matrix obtained from finite element discretization, dividing tasks for the NPU and storing them resident on the device, and performing the required matrix-vector multiplication, vector dot product and vector update on the NPU device, the data transmission and synchronization overhead between the host and the device is reduced, thereby improving the solution performance of sparse linear systems.

[0007] II. Technical Solution

[0008] To achieve the above objectives, the present invention adopts the following technical solution:

[0009] 1) Preprocessing of the solution matrix: Obtain the global stiffness matrix A and load vector b generated by the finite element mesh, perform block compression and rearrangement on stiffness matrix A, and map it into a compressed sparse row format adapted to the NPU matrix operation unit. This CSR format includes row offset array, column index array and numerical array.

[0010] 2) Matrix transformation and mapping: Before the NPU performs matrix-vector multiplication, the number of non-zero elements in each row is analyzed based on the row offset array. The computation task is dynamically divided into multiple computation groups so that the total number of non-zero elements in each computation group is approximately adapted to the load capacity of each NPU computation core.

[0011] 3) Allocate and store intermediate vectors: Allocate and persistently store solution vectors on the NPU device. residual vector Search direction vector With intermediate vector ;

[0012] 4) Solving and optimizing linear equations: Perform conjugate gradient calculation on the NPU device, using the CANN operator for computation. , And calculate the step size parameter. This allows for vector updates to be performed on the NPU device; the updated residual norm is calculated on the device and compared with a preset convergence threshold to determine whether to terminate the iteration.

[0013] 5) Computation End and Output: x is only transmitted from the NPU device back to the host for output after the iteration ends; high performance is achieved by keeping CSR format data and vectors resident on the device during the iteration process and executing sparse matrix and vector operators on the device. This reduces the data transmission and synchronization overhead between the host and the device.

[0014] Based on the above technical solution, the present invention can further adopt the following optimized features:

[0015] (1) NPU includes, but is not limited to, Ascend NPU; the device-side calculation of the dot product can be achieved by calling the element-wise multiplication operator and the reduction summation operator through the NPU runtime interface;

[0016] (2) The construction of CSR format may include: estimating the number of non-zero elements on the NPU and pre-allocating the row offset array and column index array of A to reduce dynamic expansion and memory copying;

[0017] (3) The process of dynamically dividing multiple calculation groups may include: for short rows with fewer than a preset threshold, merging multiple short rows into the same calculation group; for long rows with significantly more than the threshold, dividing the long row into multiple sub-segments, and assigning each sub-segment as an independent calculation task to different calculation cores for processing.

[0018] (4) In execution Previously, multiple threads of the NPU could collaboratively work together to adjust the search direction vector. Frequently accessed elements are preloaded and cached in on-chip high-speed storage; during computation... During the process, each thread accesses the on-chip cache based on the column index array. Data to reduce random access latency; in calculating vector inner products and At that time, the merging of the block parts can be completed through the parallel reduction instructions inside the device;

[0019] (5) In parallel reduction operators, FP32 numerical precision can be used first to ensure numerical stability, and then reduced to FP16 at appropriate positions to improve computation speed.

[0020] (6) Replace some scalars that need to reside in memory with vectors of one-dimensional length to reduce data exchange between scalars stored on the NPU chip and CPU memory;

[0021] (7) Combine sparse matrix-vector multiplication, vector inner product, and vector update into a continuous device-side execution operator to complete the operation. After calculation, the intermediate results stored in registers or on-chip caches are directly reused. Computation and vector updates; multiple iterations can be executed continuously via device-side graph execution mode or loop kernel mode, reducing host-side kernel scheduling and startup overhead; during iteration... None of them send back to the host; the final solution is only output after the convergence condition is met.

[0022] III. Beneficial Effects

[0023] Compared with the prior art, the present invention has at least the following beneficial effects:

[0024] (1) For the first time, an iterative solver for sparse linear equations in traditional high-performance computing was implemented on an NPU, verifying the practicality and feasibility of applying the NPU to high-performance computing;

[0025] (2) By using the sparse format of CSR and the intermediate vector to reside on the device side, the data transmission and synchronization waiting between the host and the device are significantly reduced;

[0026] (3) By partitioning the computing groups based on non-zero distribution, the load balance of each computing core of the NPU is improved, and the parallel efficiency caused by long tail lines is reduced.

[0027] (4) By using strategies such as FP32 / FP16 mixed precision and scalar vectorization, the execution efficiency of the device can be further improved while ensuring convergence stability. Detailed Implementation

[0028] The present invention will be further described below in conjunction with the implementation of this project. It should be understood that the following embodiments are only used to explain the present invention and are not intended to limit the scope of protection of the present invention; various modifications or substitutions can be made to the embodiments by those skilled in the art without departing from the technical solution of the present invention, and all such modifications or substitutions should fall within the scope of protection of the present invention.

[0029] I. Application Scenarios and Discrete Modeling

[0030] Can be combined Figure 1 and Figure 2 Understanding the implementation process of this invention, wherein Figure 1 The process of solving a system of linear equations running on an NPU is shown. Figure 2 The macroscopic operational architecture of the linear equation system solution method of this invention is illustrated. This project employs the two-dimensional finite element method for the region... Triangular or quadrilateral meshes are generated, and Gaussian integrals are performed on each element to assemble the sparse, symmetric, positive definite global stiffness matrix. With load vector Taking the conjugate gradient solution method as an example, we solve the linear equation system.

[0031] II. Sparse Format Construction and Preprocessing

[0032] To adapt to sparse computation in NPUs, this invention converts the stiffness matrix A into CSR format. The CSR format includes:

[0033] (1) Numeric array A_val: stores non-zero element values ​​row by row;

[0034] (2) Column index array A_col: The column number corresponding to A_val;

[0035] (3) Row offset array A_row: records the starting position of each row in A_val

[0036] In one implementation, the number of non-zero elements in A can be estimated, and A_val and A_col can be pre-allocated to reduce dynamic expansion and memory copying. After the conversion is complete, the CSR data can reside on the NPU device after a one-time transmission, providing input for SpMV in subsequent iterations.

[0037] III. Dynamic Task Partitioning and Mapping for NPU

[0038] To improve the parallel efficiency of sparse computation, this invention counts the number of non-zero elements in each row based on the adjacent differences of A_row before performing sparse matrix-vector computation, and dynamically divides the sparse matrix-vector computation into multiple computation groups according to the following rules:

[0039] (1) Short row merging: rows with fewer than the threshold number of non-zero elements are merged into the same calculation group to reduce thread idle time;

[0040] (2) Long row segmentation: For rows with a number of non-zero elements significantly higher than the threshold, the row is divided into multiple sub-segments, and each sub-segment is assigned to different computing cores for parallel processing;

[0041] (3) Intra-group load balancing: control the total number of non-zero elements in each computing group to be within the target range so that it approximately matches the throughput of the NPU computing core and the on-chip storage capacity.

[0042] The above dynamic partitioning can reduce the load imbalance caused by uneven non-zero distribution and improve the overall throughput of sparse matrix-vector computation.

[0043] IV. Device-side resident storage and iteration process

[0044] This invention allocates and keeps the storage resident on the NPU device side. Intermediate vectors; during iteration, calculations should be completed on the device side as much as possible, and only after the iteration is complete should the vectors be shared. Return the host output. A typical conjugate gradient iteration process can be as follows: Figure 1 As shown:

[0045] (1) Initialization: Settings It is a zero vector. , ;

[0046] (2) Iterative calculation:

[0047] a) Calculation ;

[0048] b) Calculation and , and by Obtain the step size parameter;

[0049] c) Update ,renew ;

[0050] d) Calculate the updated residual norm And compare it with the convergence threshold;

[0051] e) Calculation ,renew ;

[0052] (3) End output: After the convergence condition is met or the maximum number of iterations is reached, the device output will be terminated. Copy back to the host and output. Attached Figure Description

[0053] Figure 1 This is a schematic diagram of the linear equations solution process running on the NPU according to the present invention;

[0054] Figure 2 This is a flowchart illustrating the macroscopic operational architecture of the linear equation system solution method of this invention.

Claims

1. A high-performance method for solving linear equations running on an NPU, characterized in that, The method includes the following steps: Preprocessing steps for solving the matrix: Obtain the global stiffness matrix generated from the finite element mesh. and load vector For the stiffness matrix Block compression and rearrangement are performed, mapping the data to a compressed sparse row format adapted to the NPU matrix operation unit. This format includes row offset arrays, column index arrays, and numerical arrays. Furthermore, if the form of the global stiffness matrix is ​​similar to a portion of the solver's built-in form, the global stiffness matrix is ​​not passed. Instead, a matrix-free method is used to solve it; Matrix transformation and mapping steps: Before the NPU performs matrix-vector multiplication, the number of non-zero elements in each row is analyzed based on the row offset array. The computation task is dynamically divided into multiple computation groups so that the total number of non-zero elements in each computation group is approximately adapted to the load capacity of each NPU computation core. The steps for allocating and storing intermediate vectors are as follows: Allocate and persistently store the solution vectors on the NPU device. Intermediate vector , and ; The steps for solving and optimizing the linear equation system are as follows: The linear equation system is solved on the NPU device, utilizing the CANN operator for computation. , And calculate the step size parameter. Therefore, vector updates are performed on the NPU device side, and the updated residual norm is calculated on the device side and compared with a preset convergence threshold to determine whether to terminate the iteration. End of calculation and output steps: Only after the iteration is completed, output the... The method transmits the output from the device back to the host. The method achieves high performance by keeping the CSR format data and the vector resident on the device during the iteration process and executing the sparse matrix and vector operators on the device. This reduces the data transmission and synchronization overhead between the host and the device.

2. The method according to claim 1, characterized in that, The NPU includes, but is not limited to, the Ascend NPU, and the device-side calculation of the dot product is implemented by calling the element-wise multiplication operator and the reduction summation operator through the runtime interface of the NPU.

3. The method according to claim 1, characterized in that, The construction of the CSR format includes: estimating the number of non-zero elements on the NPU and pre-allocating the row offset array and column index array of A to reduce dynamic expansion and memory copying.

4. The method according to claim 1, characterized in that, Matrix-free methods implement the product of linear operators and vectors using operators during iteration, without explicitly storing discrete matrices.

5. The method according to claim 1, characterized in that, In the matrix transformation and mapping steps, the process of dynamically dividing multiple computation groups includes: after analyzing the number of non-zero elements in each row, for short rows with fewer than a preset threshold, merging multiple short rows into the same computation group; for long rows with significantly more than the threshold, dividing the long row into multiple sub-segments, with each sub-segment being assigned as an independent computation task to different computation cores for processing.

6. The method according to claim 1, characterized in that, The steps for solving and optimizing the linear equation system also include: performing... Previously, multiple threads of the NPU collaboratively preloaded frequently accessed elements in the search direction vector p and cached them in the NPU's on-chip high-speed storage area; during computation... During the process, each thread accesses the vector in the on-chip cache based on the column index array. Data, thereby reducing the latency caused by random memory access; in calculating the vector inner product and At this time, the merging of the block partial sums is completed through the parallel reduction instructions or register mixed instructions inside the NPU device, thereby efficiently completing the calculation of scalar values ​​at the device end.

7. The method according to claim 1, characterized in that, The solution and optimization steps for the linear equation system also include: when using the parallel reduction operator inside the NPU device, using the numerical precision of FP32 format as much as possible to ensure that the numerical precision does not cause divergence during the calculation process, and reducing the precision to FP16 in appropriate places to improve the calculation speed.

8. The method according to claim 1, characterized in that, The solution and optimization steps for the linear equations also include replacing some scalars that need to reside in memory with vectors of one-dimensional length, minimizing data exchange between scalars stored on the NPU chip and CPU memory, and reducing performance loss caused by data transfer.

9. The method according to claim 1, characterized in that, In the solution and optimization steps of the linear equations, the SpMV, vector inner product, and vector update operations are combined into a set of continuous device-side execution operators. This ensures that the solution process is completed within the specified steps. After calculation, the intermediate results stored in registers or on-chip caches are directly reused. The computation and vector update process avoids writing intermediate results back to CPU memory; multiple iterative steps are executed continuously on the NPU through device-side graph execution mode or loop kernel mode, reducing kernel scheduling and startup overhead on the host side; vector solutions are performed during the iteration process. Intermediate vector and The solution is not sent back to the host; it is only output after the convergence condition is met, thereby maximizing the performance improvement brought about by data residency on the device.

Images