Method for operating DeePMD-kit model in Shenwei super computer

Abstract

The invention relates to the technical field of computers, in particular to an operation method of a DeePMD-kit model in a Shenwei super computer, six core groups CG of the Shenwei super computer are respectively responsible for reasoning of six sub-regions of an MPI process, each core group CG comprises an MPE and 64 CPEs, and the MPE contains a 256-bit SIMD unit; each CPE contains a 512-bit SIMD (Single Instruction Multiple Data) unit; each CPE comprises a local data memory, and each local data memory LDM realizes continuous data transmission through a DMA strategy and a main memory; the CPEs in the same core group CG realize data transmission through remote memory access (RMA); the reasoning method of each core group CG comprises the following steps: accelerating a time-consuming operator through an SACA, and distributing a calculation-intensive part of the time-consuming operator to 64 CPEs by MPE; after calculation is completed, the MPE synchronizes calculation results of the 64 CPEs and stores the results in a main memory; when the SACA runs, the SIMD is parallel, and a 512-bit vectorization instruction is sent out through the CPE. By means of the technical scheme, the method has the advantage that the DeePMD-kit model can be applied to the Shenwei computer.

Description

Technical Field

[0001] The present invention relates to a DeePMD-kit model, and in particular to an operation method of the DeePMD-kit model on a Shenwei supercomputer. Background Art

[0002] The DeePMD-kit model is an open source NNMD package based on the deep potential (DP) method. It supports training DP models and performing DP model inference to derive potential energy surfaces (PES). Currently, the DeePMD-kit model is implemented as a force field plugin for LAMMPS for large-scale molecular dynamics simulations. The calculation process of the DeePMD-kit model is as follows: Figure 1 As shown, LAMMPS completes one step of MD simulation based on these statistics and updates the atomic positions.

[0003] Specifically, in a one-step simulation, the DP model takes L(i) as input and then predicts the potential energy surface (PES), atomic forces, and stress tensors. (a) The ProdEnvMat operator maps the neighbor list L(i) to the environment matrix (b) The descriptor network transforms s(r ij )and Map to It preserves physical symmetry during model training; (c) The Tabulate operator compresses the descriptor network using Weierstrass approximation; (d) The fitted network predicts the atomic energy Ei; (e) The TabulateGrad operator calculates the gradient of the descriptor network; (f) The ProdForce operator calculates the gradient of the descriptor network according to r j The gradient of ) Calculate the atomic force F i ; (g) The ProdVirial operator calculates Virial Ξ; (h) Performs ab initio molecular dynamics and analyzes the physical or chemical properties of the simulation area.

[0004] The DeePMD-kit model contains custom TensorFlow operations, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial. The calculation process of DeePMD-kit is as follows Figure 1 As shown, it mainly consists of two parts: descriptor D and fitting network N. D is responsible for extracting the input environment matrix (denoted as ) to obtain the symmetry-preserving feature. It can be expressed as:

[0005]

[0006] in, Expressed as an embedding matrix, Nm Indicates the maximum number of neighbors in the neighbor table. G i The sub-matrix represents the M < A vector of rows of the environment matrix It is expressed as:

[0007]

[0008] Among them, r ij =r j -r i , is the relative distance between atom i and its neighbor atom j, s(r ij )=w(r ij ) / |r ij |, where w(r ij ) is represented as a gate function when |r ij | <r c When , it decays smoothly from 1 to 0. The matrix G i Fitting s(r ij ) is generated, expressed as:

[0009]

[0010] Each layer is a fully connected network. Each layer increases the dimension of the input data until the final network dimension M is obtained.

[0011] The fitting network N learns the local environment features and atomic energy contributions E obtained from D i Therefore, the PES of the entire system is expressed as the sum of the atomic energy contributions:

[0012] E=∑ i E i ;

[0013] The entire network is also composed of a fully connected network and uses Tanh as the network activation function. By combining D and N, the DeePMD-kit model can achieve first-principles accuracy by fitting first-principles data generated from scientific computing software such as VASP, QuantumEspresso and PWmat. The force operation of atom i is obtained during the back propagation process (such as Figure 1 (f) is expressed as:

[0014]

[0015] The Virial operation of the system is as follows Figure 1 As shown in (g), the computational virial operator in the model is expressed as:

[0016] Ξ=∑ i Ri F i ;

[0017] The floating-point operations of the embedded network account for more than 95% of the total FLOPs. A model compression algorithm is designed to reduce the overall amount of calculation. The modified module is represented as follows:

[0018] G i =a0+a1x+a2x 2 +a3x 3 +a4x 4 +a5x 5 ;

[0019] Where x = s(r {ij} ). The DP model is trained and compressed, which reduces the computational cost by 82% in MD simulation (inference). The compressed model is constructed through Tabulate and TabulateGrad operations, such as Figure 1 (c)(e) It is worth noting that the compressed model maintains a high accuracy. For example, the average root mean square error (RMSE) of the energy of a water system with 100 water molecules is 2.0*10 -3}eV / atom, the RMSE of atomic force is

[0020] However, the processor used in the computer to which the existing DeePMD-kit model is applied is usually a CPU processor or a GPU processor.

[0021] The new generation of Sunway supercomputers is equipped with a new generation of heterogeneous multi-core processors (i.e., SW26010-pro). Each computing node is interconnected through a fat tree topology, and 256 nodes in each rack form a supernode. Each node is equipped with a SW26010-pro processor, which contains 6 core groups (CGs) and has a theoretical peak performance of approximately 14TFLOPS under double-precision computing. Each core group contains a management processing unit (MPE) and a grid of 8$\times$8 computing processing units (CPEs), connected to a 16GB DDR4 main memory with a bandwidth of 51.2GB / s. Each CPE has a 256KB manually controlled local data memory (LDM). MPE and CPE have 256-bit and 512-bit SIMD units, respectively.

[0022] The SW26010-pro processor supports direct memory access (DMA), which can transfer continuous data between LDM and main memory, with a theoretical bandwidth of 307GB / s. In addition, data is transferred between CPEs in the same core group through remote memory access (RMA), with a bandwidth of about 400GB / s. The programming interface of the new generation of Shenwei is SACA, which provides a basic compilation environment, basic libraries and runtime support.

[0023] Although previous studies have successfully optimized the performance of the DeePMD-kit model on ARM and NVIDIA GPU platforms, there are still many challenges in migrating the DeePMD-kit model to the new Sunway supercomputer.

[0024] The architecture of the SW26010-pro processor is significantly different from mainstream CPUs and GPUs. Unlike typical multi-core CPUs that execute instructions simultaneously on homogeneous cores, or GPUs that use thread blocks to organize parallel threads, the SW26010-pro uses an MPE-CPE coordination mechanism, where one MPE controls 64 CPEs. Therefore, programmers must carefully determine the division and allocation of tasks.

[0025] Therefore, it is urgent to develop a DeePMD-kit model for application in the Sunway supercomputer.

[0026] SW26010-pro does not optimize main memory access through multi-level cache like CPU, but equips each CPE with programmable fast LDM, similar to manually controlled L1 cache, which makes cache strategy more flexible and memory access more configurable. In addition, unlike GPU using shared memory for thread data sharing, or CPU relying on cache for inter-thread communication, SW26010-pro does not provide physical shared memory for CPE. Instead, data exchange between different CPEs is managed by RMA mechanism, and this uniqueness creates a unique programming framework. We should make full use of these characteristics of SW26010-pro architecture to obtain excellent performance.

[0027] Currently, the resource utilization of DeePMD-kit operators is quite low. Unlike implementations using OpenMP (GPU) and CUDA (GPU), the design of MPE and CPE code requires careful consideration. MPE calls 64 CPEs for calculations via the \textit{athread_spawn()} instruction, and each CPE performs calculations and memory accesses simultaneously. However, previous work did not optimize the DeePMD-kit operators for the SW26010-pro architecture, resulting in all operators being executed only on MPE. In addition, some new features of CPE have not been effectively utilized, including memory access through DMA, on-chip data sharing through RMA, and 512-bit vectorization. For example Figure 3 The performance analysis of the DeePMD-kit time-consuming operators shown in shows that the achieved peak performance and memory bandwidth of each operator are significantly lower than the theoretical peak performance (14 TFLOPS) and peak memory bandwidth (300 GB / s).

[0028] The DeePMD-kit model relies on multiple compute-intensive TensorFlow custom operators and official operators; custom operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial; official operators include MatMul, Tanh, TanhGrad, Slice, and Pad.

[0029] Therefore, it is crucial to optimize each DeePMD-kit operator for the architecture of the Sunway supercomputer to ensure high resource utilization and computational efficiency. Summary of the invention

[0030] In order to solve the problems in the related technology, the present application provides a DeePMD-kit model applied to the SW26010-pro processor of the Sunway supercomputer, which solves the problem that the existing DeePMD-kit model cannot be applied to the SW26010-pro processor in the new generation of Sunway supercomputers.

[0031] The technical solution is as follows:

[0032] The method for running the DeePMD-kit model on the Shenwei supercomputer is characterized in that the six core groups CG of the Shenwei supercomputer are respectively responsible for the reasoning of the six sub-areas of the MPI process; the programming interface of the Shenwei supercomputer is SACA;

[0033] Each of the core groups CG includes a management processing unit MPE connected to the main memory DDR and 64 computing processing units CPE, the management processing unit MPE contains a 256-bit SIMD unit; each computing processing unit CPE contains a 512-bit SIMD unit; each computing processing unit CPE contains a local data memory LDM, each local data memory LDM cancels the cache configuration, and each local data memory LDM realizes continuous data transmission through the DMA strategy and the main memory DDR; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA;

[0034] Among them, the reasoning method of each core group CG is:

[0035] SACA is used to accelerate time-consuming operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive parts of the time-consuming operators to 64 computation processing units CPE according to the atomic ID. Each computation processing unit CPE calculates the corresponding divided blocks respectively.

[0036] After the calculation processing unit CPE completes the calculation, the management processing unit MPE executes the athread join() instruction to synchronize the calculation results of the 64 calculation processing units CPE, and stores the results in the main memory DDR;

[0037] SIMD parallelism during SACA runtime: The computational processing unit (CPE) issues 512 vectorized instructions to improve the parallel efficiency of custom operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial.

[0038] Through the above technical solution, the DeePMD-kit model can be applied in the Sunway supercomputer through MPI, SACA and SIMD parallelism; and by using 512-bit vectorized instructions, the parallel efficiency of the DeePMD-kit operator can be improved.

[0039] Preferably, the molecular system is evenly divided into multiple sub-regions according to active MPI processes; the Sunway supercomputer starts 6 MPI processes on the SW26010-pro node, each core group CG holds an MPI process and is responsible for calculating one sub-region of all atoms.

[0040] Preferably, the computational efficiency and bandwidth utilization of the Tabulate operator and the TabulateGrad operator are optimized by overlapping computation with memory access and kernel fusion. The optimization process is as follows:

[0041] Step 1, each computing processing unit CPE initiates an asynchronous memory access request by issuing a CRTS_dma_iget() or CRTS_dma_iput() instruction; after the DMA request is initiated, the computing processing unit CPE synchronously waits for memory access and checks the DMA status by executing a dma_wait() instruction. Before executing the dma_wait() instruction, the calculations of the Tabulate operator and the TabulateGrad operator will not be blocked, and calculations that have no data dependency on the current calculation pipeline are allowed to be executed before executing the dma_wait() instruction;

[0042] The asynchronous DMA transfer between the local data memory LDM and the main memory DDR is synchronized by executing the dma_wait() instruction during the manual control of the overlapping efficiency;

[0043] Each computing processing unit CPE is responsible for executing multiple descriptor vectors D(R i ), input s(r ij ) and its corresponding table coefficients are divided into two buffers, each of which independently processes a portion of the descriptor vector D(R i ), the processing of the buffer includes calculation and asynchronous DMA access between the local data memory LDM and the main memory DDR, and the calculation data of the two buffers are synchronized through the dma_wait() instruction;

[0044] When the first buffer completes the DMA transfer from the main memory DDR to the local memory and starts computing, the second buffer initiates a DMA request to load s(r ij ) and table coefficients; when the second buffer completes the DMA request and starts calculation, the first buffer request will be D(R i ) written back to main memory;

[0045] Step 2: The TabulateGrad operator reuses the G matrix previously calculated by the Tabulate operator.

[0046] Through the above technical solution, by overlapping the calculation and memory access, because the Sunway supercomputer supports high-bandwidth asynchronous DMA, it can efficiently access continuous data blocks in the main memory, so that the operator can maximize the 300GB / s DMA bandwidth and improve the throughput;

[0047] By making the TabulateGrad operator reuse the G matrix previously calculated by the Tabulate operator, frequent memory accesses when calculating the G matrix can be omitted, thereby eliminating redundant calculations of the Tabulate operator, thereby saving 54.6% of the execution time.

[0048] Preferably, the computational efficiency of the ProdEnvMat operator is optimized by data layout transformation and branch elimination, and the optimization process is as follows:

[0049] Step 1: Convert the AoS variable into the SoA format, perform numerical shuffling through the VSHFW of the Sunway supercomputer, and store it back into the AoS format after vectorization.

[0050] In VSHFW, the source vector is divided into 16 32-bit words, the shuffle mask includes 16 5-bit components, and each component is responsible for specifying the word position in the source vector. By executing VSHFW, the 16 words of the target vector are allocated according to the mask;

[0051] Step 2: Calculate a conditional mask for the SIMD vector, multiply the SIMD vector by the conditional mask, and derive the branch output.

[0052] Through the above technical solution, through the setting of data layout conversion, the parallel loading, storage and computing capabilities of the Sunway supercomputer can be fully utilized, and then the numerical shuffle can be performed efficiently, so that the throughput can be increased by 4 times through 256-bit VSHFW and by 8 times through 512-bit VSHFW;

[0053] By setting the score elimination, the divergence can be reduced, thereby eliminating redundant calculations, thereby optimizing the computational efficiency of the ProdEnvMat operator in the DeePMD-kit model.

[0054] Preferably, in the CPE array architecture, each cluster includes four computing processing units CPE and 16 slave cluster management SCMs. The 16 slave cluster management SCMs in each CPE array are interconnected through an on-chip network, and each slave cluster management SCM includes an RMA engine and a DMA engine.

[0055] Preferably, the write conflict of the ProdVirial operator is optimized through the replication and reduction strategy. The optimization process is as follows:

[0056] Step 1: Create a virial replica named virial_rep for each computing processing unit CPE and initialize its elements to 0;

[0057] Calculate the increment tmp_v of each computing processing unit CPE on virial and add it to virial_rep;

[0058] After each computing processing unit CPE calculates virial_rep, it synchronizes the calculation results of 64 computing processing units CPE by executing the CRTS_ssync_array() instruction;

[0059] Step 2: The first computing processing unit CPE in each cluster aggregates the partial values ​​through two rounds of operations that are quickly executed by the RMA engine in the cluster;

[0060] The first computational processing unit CPE in the first column of the CPE array derives the intra-row partial reduction by performing the RMA engine reduction on the clusters sharing the same row number;

[0061] The first computing processing unit CPE in the first column performs partial reduction through the RMA engine, and the first computing processing unit CPE of the first cluster obtains the final reduction result and writes it back to the main memory DDR through the DMA strategy.

[0062] Through the above technical solution, by setting the copy and reduce strategy, the integrity of the code before the CRTS_ssync_array() instruction can be guaranteed. Through the three-step RMA-based reduction method, redundant global synchronization can be avoided, thereby improving parallel efficiency.

[0063] Preferably, the TensorFlow operator is optimized by fitting the network optimization and parallelizing the tensor operator. The optimization process of the TensorFlow operator is as follows:

[0064] Step 1, optimizing the fitting network: the fitting network includes a plurality of fully connected layers and a Tanh activation function, which are used to obtain corresponding gradients for force calculation during the back propagation process; the fully connected layer of the Shenwei supercomputer includes a MatMul operator and an Add operator;

[0065] Fuse the MatMul operator and the Add operator into the GEMM in Eigen, and use swBLAS in the Shenwei supercomputer to synchronously accelerate the Tanh and TanhGrad operators through SACA;

[0066] Step 2: Parallelize tensor operators: Parallelize Pad and Slice operators through SACA, and use DMA strategy to optimize the bandwidth of main memory DDR.

[0067] Through the above technical solution, by fusing the MatMul operator and the Add operator into the GEMM in Eigen, using the swBLAS of the Sunway supercomputer, and accelerating the Tanh and TanhGrad operators through SACA, the computing power of the CPE in the SW26010-pro processor in the Sunway computer can be fully utilized, thereby optimizing the computing efficiency of the TensorFlow operator in DeePMD-kit;

[0068] By parallelizing Pad and Slice operators through SACA, the main memory DDR bandwidth utilization can be maximized, thereby optimizing TensorFlow operators.

[0069] Preferably, the bandwidth of the Tabulate operator is optimized by a mixed-precision tabulatetion method, which is as follows:

[0070] Step 1: Create low-order weights a0, a1, a2, and a3 using bfloat16, and high-order weights a4 and a5 using float32.

[0071] Step 2: Design a mixed-precision Tabulate kernel using SACA and vectorized intrinsics.

[0072] Through the above technical solution, the DeePMD-kit descriptor is directly fitted from scratch by using low-precision table weights in mixed precision, which can minimize the compression error. By using bfloat16 to create low-order weights, overflow problems in the training process can be effectively prevented. By utilizing the efficient and low-precision computing power of the CPE equipped with the Sunway supercomputer, this can save about 3 times the memory usage during the Tabulate operator calculation process, thereby increasing the throughput by about 23 times.

[0073] Preferably, the calculation workflow of the optimized Tabulate operator is as follows:

[0074] Step 1. Calculation and with Perform mixed precision vector-vector multiplication to obtain the activation matrix a ij ;

[0075] Step 2: ij with a ij Perform matrix-vector multiplication to obtain the descriptor vector

[0076] Through the above technical solution, by calculating The time-intensive loop before calculating a ij, and converts the Tabulate operation into two matrix-vector multiplications, which reduces M×38 redundant calculations, thereby optimizing the computational efficiency of the Tabulate operator.

[0077] In summary, the beneficial effects of the present invention are:

[0078] 1. Through MPI, SACA and SIMD parallelism, the DeePMD-kit model can be applied in the Sunway supercomputer; and by using 512-bit vectorized instructions, the parallel efficiency of the DeePMD-kit operator can be improved;

[0079] 2. By optimizing the performance of Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, MatMul, Tanh, TanhGrad, Slice, and Pad operators, the utilization of the new generation of Shenwei resources can be maximized.

[0080] It is to be understood that the foregoing general description and the following detailed description are exemplary only and are not restrictive of the invention. BRIEF DESCRIPTION OF THE DRAWINGS

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

[0082] Figure 1 It is a schematic diagram of the operation of the DeePMD-kit model in the prior art;

[0083] Figure 2 A schematic diagram showing the architecture comparison between a processor in the prior art and a processor of the Sunway supercomputer;

[0084] Figure 3 Analysis of time-consuming operators in the SW26010-pro processor of the present invention;

[0085] Figure 4 It is a schematic diagram of the operation of the DeePMD-kit model in the present invention on the Sunway supercomputer;

[0086] Figure 5 It is a schematic diagram of the operation of the overlap of computing and memory access on the Sunway supercomputer in the present invention;

[0087] Figure 6 It is a schematic diagram of the operation of the conversion between AoS and SoA in the present invention;

[0088] Figure 7 It is a schematic diagram of the operation of the replication and reduction strategy in the ProdVirial calculation of the present invention;

[0089] Figure 8 It is a schematic diagram of the operation of the mixed precision tabularization solution in the present invention;

[0090] Fig. 9 Schematic diagram of the operation of radial distribution functions gOO(r), gOH(r) and gHH(r) of liquid water in the present invention;

[0091] Fig.10 It is a schematic diagram of the operation of the system energy curve of the water system of the present invention during the 1,000,000-step simulation process;

[0092] Fig.11 It is a schematic diagram of the operation of the gradual performance improvement of the 9,126-atom water system and the 5,184-atom copper system on a single SW26010-pro node in the present invention;

[0093] Fig.12 It is a schematic diagram of the operation of the performance evaluation of the custom operator in the DeePMD-kit model of the present invention;

[0094] Fig.13 It is a schematic diagram of the operation of the strong scalability evaluation of the 100-step MD simulation of the water system in the present invention;

[0095] Fig.14 It is a schematic diagram of the operation of the strong scalability evaluation of the 100-step MD simulation of the copper system in the present invention;

[0096] Fig.15 It is a schematic diagram of the operation of the weak scalability of the water and copper system in the Sunway supercomputer of the present invention; DETAILED DESCRIPTION

[0097] Exemplary embodiments will be described in detail herein, examples of which are shown in the accompanying drawings. When the following description refers to the drawings, the same numbers in different drawings represent the same or similar elements unless otherwise indicated. The embodiments described in the following exemplary embodiments do not represent all embodiments consistent with the present invention. Instead, they are merely examples of devices and methods consistent with some aspects of the present invention as detailed in the appended claims.

[0098] As attached Figure 1-15 As shown in the figure, the operation method of the DeePMD-kit model on the Sunway supercomputer is shown in the figure. The Sunway supercomputer starts 6 MPI processes on the SW26010-pro node. Each core group CG holds an MPI process and is responsible for calculating a sub-region of all atoms. The molecular system is evenly divided into multiple sub-regions according to the active MPI processes. The six core groups CG of the Sunway supercomputer are responsible for the reasoning of the six sub-regions of the MPI process respectively. The programming interface of the Sunway supercomputer is SACA.

[0099] Each core group CG includes a management processing unit MPE connected to the main memory DDR and 64 computing processing units CPE. The management processing unit MPE contains a 256-bit SIMD unit; each computing processing unit CPE contains a 512-bit SIMD unit; each computing processing unit CPE contains a local data memory LDM. In this specific embodiment, the local data memory LDM is manually controlled, and each local data memory LDM cancels the cache configuration. Each local data memory LDM realizes continuous data transmission through the DMA strategy and the main memory DDR; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA;

[0100] Among them, the reasoning method of each core group CG is:

[0101] SACA is used to accelerate time-consuming operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive parts of the time-consuming operators to 64 computation processing units CPE according to the atomic ID. Each computation processing unit CPE calculates the corresponding divided blocks respectively.

[0102] After the calculation processing unit CPE completes the calculation, the management processing unit MPE executes the athread join() instruction to synchronize the calculation results of the 64 calculation processing units CPE, and stores the results in the main memory DDR;

[0103] SIMD parallelism during SACA runtime: The computational processing unit (CPE) issues 512 vectorized instructions to improve the parallel efficiency of custom operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial.

[0104] Through MPI, SACA and SIMD parallelism, the DeePMD-kit model can be applied in the Sunway supercomputer; and by using 512-bit vectorized instructions, the parallel efficiency of the DeePMD-kit operator can be improved.

[0105] like Figure 5 As shown, the Tabulation method uses a fifth-order polynomial approximation to compress the embedded network G i, saving 82% of the total floating point operations (FLOPs). This method is implemented by the Tabulate and TabulateGrad operators in DeePMD-kit. However, the current abulation implementation on the new generation of Shenwei is time-consuming. Analysis of the Tabulate and TabulateGrad operators shows that the memory access efficiency is low when loading polynomial coefficients (that is, the memory bandwidth utilization is less than 1%), which occupies 68.6% and 54.6% of the execution time respectively; in this specific embodiment, the computational efficiency and bandwidth utilization of the Tabulate operator and the TabulateGrad operator are optimized by computing the overlap with the memory access and kernel fusion, and the optimization process is as follows:

[0106] Step 1: Tabulate and TabulateGrad operators use a strategy of overlapping computation and memory access, and their core computation is a three-layer nested loop. The new generation of Shenwei supports high-bandwidth asynchronous DMA, which can efficiently access continuous data blocks in the main memory.

[0107] Each computing processing unit CPE initiates an asynchronous memory access request by issuing a CRTS_dma_iget() or CRTS_dma_iput() instruction. After the DMA request is initiated, the computing processing unit CPE synchronously waits for memory access and checks the DMA status by executing the dma_wait() instruction. Before executing the dma_wait() instruction, the calculations of the Tabulate operator and the TabulateGrad operator will not be blocked, and calculations that have no data dependency on the current calculation pipeline are allowed to be executed before executing the dma_wait() instruction.

[0108] The asynchronous DMA transfer between the local data memory LDM and the main memory DDR is synchronized by executing the dma_wait() instruction during the process of manually controlling the overlap efficiency; and the asynchronous DMA transfer between the LDM and the main memory is manually controlled to maximize the overlap efficiency; each CPE is responsible for executing multiple descriptor vectors D(R i ) to fully utilize the LDM space (for example, 32 in water system and 8 in copper system, using 223KBLDM). With this overlapping strategy, these operators can maximize the 300GB / s DMA bandwidth and improve throughput.

[0109] Each computing processing unit CPE is responsible for executing multiple descriptor vectors D(R i ), input s(r ij ) and its corresponding table coefficients are divided into two buffers, each of which independently processes a portion of the descriptor vector D(R i), the processing of the buffer includes calculation and asynchronous DMA access between the local data memory LDM and the main memory DDR, and the calculation data of the two buffers are synchronized through the dma_wait() instruction;

[0110] When the first buffer completes the DMA transfer from the main memory DDR to the local memory and starts computing, the second buffer initiates a DMA request to load s(r ij ) and table coefficients; when the second buffer completes the DMA request and starts calculation, the first buffer request will be D(R i ) written back to main memory;

[0111] Step 2, Tabulate and TabulateGrad operators both compute the G matrix described in Eq. However, computing G involves frequent memory accesses, including the input s(r ij ) and polynomial coefficients; reuse the G matrix previously calculated by the Tabulate operator through the TabulateGrad operator; this can eliminate redundant calculations in the TabulateGrad operator, and this optimization can save 54.6% of the execution time.

[0112] By overlapping computation with memory access, the Sunway supercomputer supports high-bandwidth asynchronous DMA, which enables it to efficiently access continuous data blocks in the main memory, maximizing the 300GB / s DMA bandwidth of the operator and improving throughput.

[0113] By making the TabulateGrad operator reuse the G matrix previously calculated by the Tabulate operator, frequent memory accesses when calculating the G matrix can be omitted, thereby eliminating redundant calculations of the Tabulate operator, thereby saving 54.6% of the execution time.

[0114] like Figure 6 As shown in Figure 1, DeePMD-kit is implemented as a force field plugin for LAMMPS. Its variable layout must be consistent with LAMMPS and use the array of structures (AoS) format. For example, the r of two atoms ijare stored as (x0, y0, z0), (x1, y1, z1), where the memory interval of these two atoms is 3 (i.e., x0 is loaded in the first loop and x1 is loaded in the second loop. In the ProdEnvMat operator, the rij_a, descrpt_a, and descrpt_a_deriv variables must be stored in AoS format, which results in non-contiguous memory access for each iteration and hinders vectorization optimization. In addition, the ProdEnvMat operator contains many conditional branches, which leads to branch divergence and reduces the computational efficiency of SIMD. Therefore, the computational efficiency of the ProdEnvMat operator is optimized by data layout transformation and branch elimination. The optimization process is as follows:

[0115] Step 1: Convert the AoS variables into the SoA format, perform numerical shuffling via the VSHFW of the Sunway supercomputer, and store them back into the AoS format after vectorization. For example, the r of two atoms ij is stored as {x0,x1,y0,y1,z0,z1};

[0116] The new generation of Shenwei architecture supports the full-word vector shuffle instruction VSHFW, which can efficiently shuffle values. In VSHFW, the source vector is divided into 16 32-bit words, and the shuffle mask includes 16 5-bit components, and each component is responsible for specifying the word position in the source vector so that it can be placed in the target vector. By executing VSHFW, the 16 words of the target vector are allocated according to the mask;

[0117] In this specific embodiment, an example of a 256-bit VSHFW instruction is shown, which shuffles two 256-bit source vectors according to a defined mask. The mask is divided into eight 5-bit binary codes, where the fifth bit indicates which source vector is selected, and the remaining four bits indicate which word in the source vector is selected. Figure 6 As shown in Figure 1, the fourth binary code of the mask is 0x13. This binary code indicates that the corresponding word of the destination vector is taken from the third word of source vector 1. For variables with a memory spacing of 6, we can use 24 VSHFW instructions to convert between AoS and SoA, resulting in a 4x throughput increase with 256-bit VSHFW and an 8x throughput increase with 512-bit VSHFW in just 24 extra instruction cycles.

[0118] Step 2: Calculate a conditional mask for the SIMD vector, multiply the SIMD vector by the conditional mask, and derive the branch output.

[0119] The ProdEnvMat operator uses conditional branches to eliminate redundant computations. For each atom, ProdEnvMat retrieves the global atom index of the neighbor from the neighbor list L(i). If the global atom index is less than 0 (i.e., all neighbors of the current atom have been retrieved), ProdEnvMat skips the current loop and instead retrieves the neighbors of the next atom in L(i). A branch instruction is executed in each loop to check the global atom index, which may cause branch distribution to diverge. To address this issue, in this specific embodiment, a branch elimination strategy is implemented to reduce divergence. Specifically, we calculate a conditional mask for the SIMD vector and derive the branch output by multiplying the SIMD vector by the mask.

[0120] By setting up data layout transformation, the parallel loading, storage, and computing capabilities of the Sunway supercomputer can be fully utilized, thereby efficiently shuffling values, thereby increasing the throughput by 4 times through 256-bit VSHFW and 8 times through 512-bit VSHFW;

[0121] By setting the score elimination, the divergence can be reduced, thereby eliminating redundant calculations, thereby optimizing the computational efficiency of the ProdEnvMat operator in the DeePMD-kit model.

[0122] like Figure 7 As shown, in the CPE array architecture, each cluster includes four computing processing units CPE and 16 slave cluster management SCMs. The 16 slave cluster management SCMs in each CPE array are interconnected through an on-chip network, and each slave cluster management SCM includes an RMA engine and a DMA engine.

[0123] like Figure 7 As shown in Figure 1, when using the intra-operator parallelism scheme on the new Shenwei, there are write conflicts in the ProdVirial operator, which are usually mitigated by atomic operations. The atomic operations on the new Shenwei are implemented by the fetch-and-add mechanism, and each atomic operation requires global synchronization on all CPEs of the CG. However, using these atomic operations will introduce 9*N in the ProdVirial operator. {des} *N {loc} *N {nnei} times synchronization, where N {des} Indicates the number of descriptors, N {loc} Indicates the number of atoms, N {nnei} Indicates the number of neighbor atoms. Frequent synchronization operations reduce parallel efficiency, resulting in performance that is inferior to the serial implementation on MPE. Therefore, in this specific embodiment, the write conflict of the ProdVirial operator is optimized through replication and reduction strategy, and the optimization process is as follows:

[0124] Step 1: Create a virial replica named virial_rep for each computing processing unit CPE and initialize its elements to 0;

[0125] Calculate the increment tmp_v of each computing processing unit CPE on virial and add it to virial_rep;

[0126] After each computing processing unit CPE calculates virial_rep, it synchronizes the calculation results of 64 computing processing units CPE by executing the CRTS_ssync_array() instruction;

[0127] Step 2: The first computing processing unit CPE in each cluster aggregates the partial values ​​through two rounds of operations that are quickly executed by the RMA engine in the cluster;

[0128] The first computational processing unit CPE in the first column of the CPE array derives the intra-row partial reduction by performing the RMA engine reduction on the clusters sharing the same row number;

[0129] The first computing processing unit CPE in the first column performs partial reduction through the RMA engine, and the first computing processing unit CPE of the first cluster obtains the final reduction result and writes it back to the main memory DDR through the DMA strategy.

[0130] By setting the copy and reduce strategy, the integrity of the code before the CRTS_ssync_array() instruction can be guaranteed. Through the three-step RMA-based reduction method, redundant global synchronization can be avoided, thereby improving parallel efficiency.

[0131] TensorFlow is not sufficiently optimized for the new Shenwei architecture, resulting in some official TensorFlow operators being inefficient in DeePMD-kit; therefore, in this specific embodiment, the TensorFlow operators are optimized by fitting network optimization and parallelizing tensor operators. The optimization process of the TensorFlow operators is as follows:

[0132] Step 1: Optimize the fitting network: The fitting network includes multiple fully connected layers and Tanh activation functions, which are used to obtain the corresponding gradients for force calculation during the back propagation process;

[0133] However, the current implementation is not efficient and cannot fully utilize the computing power of the new Shenwei. The original fully connected layer of DeePMD-kit is represented as The fully connected layer of the Sunway supercomputer includes the MatMul operator and the Add operator;

[0134] This introduces redundant data transfers of intermediate tensors and extra kernel launches.

[0135] In this specific embodiment, the MatMul operator and the Add operator are fused into the GEMM in Eigen (calculating $A*B+C$ in one kernel), and swBLAS is used in the Sunway supercomputer to synchronously accelerate the Tanh and TanhGrad operators through SACA to fully utilize the computing power of the CPE in the SW26010-pro processor;

[0136] Step 2: Parallelize tensor operators: Parallelize Pad and Slice operators through SACA, and use DMA strategy to optimize the bandwidth of main memory DDR.

[0137] By fusing the MatMul operator and the Add operator into the GEMM in Eigen, using the swBLAS of the Sunway supercomputer, and accelerating the Tanh and TanhGrad operators through SACA, the computing power of the CPE in the SW26010-pro processor in the Sunway computer can be fully utilized, thereby optimizing the computing efficiency of TensorFlow operators in DeePMD-kit;

[0138] Some specialized tensor operators, such as Pad and Slice, are currently only implemented in MPE versions on the new Shenwei. Unfortunately, these operators do not support parallel execution on CPE and fail to effectively utilize the 300GB / s DMA bandwidth. As a result, the computational performance of DeePMD-kit is significantly affected. To address these issues, the main memory DDR bandwidth utilization can be maximized by parallelizing Pad and Slice operators through SACA, thereby optimizing TensorFlow operators.

[0139] like Figure 8 As shown, for each smoothing threshold value x=s(r {ij} ), the Tabulate operator will look up the corresponding table weight w {ij} To calculate the embedding vector G {ij} , where MMM represents the feature dimension of the embedding matrix. However, each x requires access to M×6×8 bytes of w {ij} , which consumes a lot of memory access and leads to computational inefficiency in large-scale MD simulations. In order to alleviate the bandwidth bottleneck of the Tabulate operator, the bandwidth of the Tabulate operator is optimized by the mixed-precision tabulatetion method. The mixed-precision tabulatetion method is as follows:

[0140] Step 1: Create low-order weights a0, a1, a2, and a3 using bfloat16, and high-order weights a4 and a5 using float32.

[0141] Step 2: Design a mixed-precision Tabulate kernel using SACA and vectorized intrinsics.

[0142] Different from the low-precision truncation method, the proposed mixed-precision method uses low-precision table weights to directly fit the DeePMD-kit descriptor from scratch, which can minimize the compression error. Through statistical analysis of each term in the table weights, we find that the high-order coefficients (i.e., a4 and a5) show a significant range of values ​​during the training phase, indicating that these terms are more sensitive to precision changes. For example, the absolute maximum values ​​of a4 and a5 reach 8.5*10 10 and 1.3*10 11 . Since the values ​​of high-order coefficients are much larger than the representation range of float16, directly using float16 to compress weights will cause overflow problems. Even using bfloat16 to compress high-order weights, including a4 and a5, will result in significant errors; therefore, use bfloat16 to create low-order weights a0, a1, a2, and a3, and use float32 to create high-order weights a4 and a5;

[0143] The reason for choosing bfloat16 instead of float16 to create low-order weights is that the value range of bfloat16 is consistent with float32, which can prevent overflow problems during training. The new Shenwei supercomputer supports low-precision floating-point calculations and 512-bit SIMD support, including float16, bfloat16, and float32. Therefore, we can design mixed-precision Tabulate kernels using SACA and vectorized intrinsics to take advantage of the efficient low-precision computing capabilities of CPE.

[0144] The DeePMD-kit descriptor is directly fitted from scratch by using low-precision table weights in mixed precision, which can minimize the compression error. The overflow problem during training can be effectively prevented by using bfloat16 to create low-order weights. By utilizing the efficient and low-precision computing power of the CPE equipped with the Sunway supercomputer, this can save about 3 times the memory usage during the Tabulate operator calculation process, thereby increasing the throughput by about 23 times.

[0145] The calculation workflow of the optimized Tabulate operator is as follows:

[0146] Step 1. Calculation and with Perform mixed precision vector-vector multiplication to obtain the activation matrix a ij ;

[0147] Step 2: w ij with a ijPerform matrix-vector multiplication to obtain the descriptor vector

[0148] By calculating The time-intensive loop before calculating a ij , and converts the Tabulate operation into two matrix-vector multiplications, which reduces M×38 redundant calculations, thereby optimizing the computational efficiency of the Tabulate operator.

[0149] In this specific embodiment, we selected two typical physical systems—water and copper—as benchmarks for performance evaluation.

[0150] Water is a challenging system, even for AIMD methods, because it requires a balance between weak non-covalent intermolecular interactions, thermal (entropy) effects, and nuclear quantum effects.

[0151] Copper is a representative metal whose properties, such as surface formation energy and stacking fault energy, are difficult to accurately predict using empirical potential fields.

[0152] To simulate and analyze the performance of our work, we use the Tabulate-based DP model as a baseline, where the descriptor is a compressed Tabulate module with an output dimension of 128. The dimension of the fitted network is 224×224×224. The cutoff radius settings are and The maximum number of neighbors is 138 and 512. The MD equations are numerically integrated by the Velocity-Verlet method with 100 steps at time steps of 0.5 fs and 1.0 fs. The velocities of the atoms are randomly initialized according to the Boltzmann distribution at 330 K. The neighbor list is updated every 50 time steps with Thermodynamic data, including kinetic energy, potential energy, temperature and pressure, are collected and recorded every 20 time steps. The following table is obtained: Comparison of energy and force errors between the mixed precision scheme and the baseline method:

[0153]

[0154] A. Accuracy of Mixed Precision Solutions

[0155] 1. Energy and force errors. The accuracy of DeePMD-kit is studied by comparing the force and energy errors with the AIMD method. We use the test dataset calculated by VASP as an example to evaluate the accuracy, which consists of 100 water molecules. First, we predict the energy and force of these 100 water molecules using the mixed precision model, and then calculate the root mean square error (RMSE) of the predicted energy and force relative to the AIMD results. Table I shows the RMSE error of our mixed precision scheme compared with the baseline tabulation scheme. Compared with the baseline scheme, the proposed mixed precision method achieves an energy error and The energy prediction error of the mixed precision scheme is much less than 4×10^-2eV / molecule, and the force error is no less than This indicates that the proposed mixed precision method has a lower error rate.

[0156] Fig. 9 Radial distribution functions of liquid water gOO(r), gOH(r) and gHH(r) are calculated by two DeePMD-kit implementations: the baseline scheme uses float64 precision, and the mixed precision scheme uses bfloat16 and float32.

[0157] 2. Radial distribution function. To further evaluate the accuracy of mixed precision, we calculated the radial distribution function (RDF) of water, which refers to the normalized probability of searching for neighboring atoms at the spherical average distance rrr. The RDFs of oxygen-oxygen (gOO(r)), oxygen-hydrogen (gOH(r)), and hydrogen-hydrogen (gHH(r)) are commonly used to characterize the structural properties of water. Specifically, we selected 192 water atoms and performed 20,000 steps of MD simulations using the proposed mixed precision DeePMD-kit and benchmark. Then, we calculated the RDF of the water system based on the simulation results and used it in Fig. 9 Draw the RDF curve in Fig. 9 , the RDFs of the proposed mixed-precision tabulation scheme and the baseline method completely overlap. The baseline method has been shown in previous studies to achieve comparable accuracy to the AIMD method. Therefore, we can conclude that the proposed mixed-precision scheme is able to accurately predict the physical observables.

[0158] Fig.10 The system energy curve of the water system during 1,000,000 simulation steps. The mixed precision scheme can remain stable in the simulation with a long time step length and has a low error rate.

[0159] 3. Stability evaluation: In order to evaluate the MD stability of the mixed-precision scheme, we selected a water system containing 192 atoms and performed a long-term simulation evaluation of the proposed mixed-precision scheme. We used the same configuration to perform one million steps of NVE simulations for the proposed mixed-precision DeePMD-kit and the benchmark method. Then, we observed the energy fluctuations of the system during the MD simulation and Fig.10 The energy curves calculated by these two DeePMD-kit schemes are plotted in . Fig.10 , we can observe that the system energy is well preserved in long-term MD simulations. In addition, the system energy simulated by the mixed-precision scheme is very close to the benchmark results and follows the same trend, indicating that the mixed-precision method is capable of performing long-term MD simulations with low error rates.

[0160] Fig.11 Incremental performance improvement for a 9,126-atom water system and a 5,184-atom copper system on a single SW26010-pro node. Speedups are normalized relative to the baseline.

[0161] B. Single Node Performance

[0162] Our work on the new Sunway supercomputer includes optimizations of custom operators (i.e., Tabulate, TabulateGrad, ProdEnvMat, ProdVirial, and ProdForce), fitting networks, mixed precision schemes, and tensor operators. Through detailed performance analysis, we found that the most time-consuming parts are ranked as follows in priority: 1 fitting networks (68.1% of the total time), 2 custom operators (16.8%), and 3 tensor operators (13.9%). We perform optimizations step by step in this order and report the corresponding performance gains on a single new Sunway node (one SW26010-pro processor). We perform MD simulations for 100 steps and measure the corresponding cycle times (calculated by LAMMPS) for a water system with 9,126 atoms and a copper system with 5,184 atoms. The benchmark uses DeePMD-kit, a flat MPI version of the compressed model proposed in ppopp in 2022. We perform MD simulations on the same configuration and calculate the speedup of each optimization. Fig.11The performance improvement of DeePMD-kit on a single new Shenwei node is shown. Compared with the baseline, the optimized DeePMD-kit achieves 3.1 times speedup on the water system and 2.9 times speedup on the copper system. Each layer of the fitted network includes MatMul, Add, Tanh, and TanhGrad operators. These results show that the proposed kernel fusion and Tanh optimization methods can effectively improve the computational performance of DeePMD-kit. By performing customized operator optimization, the optimized DeePMD-kit achieves 6.8 times speedup on the water system and 12.1 times speedup on the copper system, which shows that the proposed optimizations including tabulation optimization, data layout optimization, branch optimization, and write conflict avoidance can significantly improve the performance of DeePMD-kit. After mixed-precision optimization on tabulation, our work achieves 8.2 times speedup on the water system and 14.1 times speedup on the copper system, which shows that using low precision can further improve the computational efficiency of CPE by maximizing the bandwidth utilization of DMA. Finally, after optimizing tensor operators, our work achieves a significant speedup of 67.6 times on the water system and 56.5 times on the copper system. This interesting phenomenon indicates that the pre-installed TensorFlow library is not sufficiently optimized for some tensor operators, which seriously affects the performance of DeePMD-kit.

[0163] Fig.12 Performance evaluation of custom operators in DeePMD-kit model. We evaluate the performance improvement of custom operators on 9,216-atom water system and 5,184-atom copper system respectively.

[0164] C. Performance of custom operators

[0165] We further evaluated the performance improvement of each custom operator on water and copper systems. Specifically, we selected Fig.11 We then measured the cycle time of the custom operator for a water system with 9,216 atoms and a copper system with 5,184 atoms. For comparison, we measured the cycle time of the benchmark model using the same configuration and calculated the speedup. Fig.12The speedup of each custom operator compared to the baseline is shown in detail. The evaluation results show that all custom operators achieve significant speedup compared to the baseline, which indicates that the proposed custom operator optimization method significantly improves performance on a single new Shenwei node. However, the performance of ProdEnvMat and ProdForce operators is relatively low, which is mainly due to the non-contiguous memory access in these operators, so we have to use gld and gst instructions for global memory access with a theoretical bandwidth of 51.2GB / s.

[0166] Fig.13 : Strong scalability evaluation of 100-step MD simulations of water systems. We evaluate the simulation efficiency on the new Shenwei with 14,910,336 atoms and compare the scalability on Fugaku with 8,294,400 atoms and Summit with 41,472,000 atoms.

[0167] Fig.14 : Strong scalability evaluation of 100-step MD simulations of the copper system. We evaluate the simulation efficiency on the new Shenwei with 2,540,160 atoms and compare the scalability on Fugaku with 2,177,280 atoms and on Summit with 13,500,000 atoms.

[0168] D. Strong scalability

[0169] We evaluate the strong scalability of the optimized DeePMD-kit on the new Sunway supercomputer. To compare with the existing DeePMD-kit implementations on Summit and Fugaku, we adopt the same node configuration, and the experiment starts with 20 nodes and gradually expands to 4,560 nodes on the new Sunway. Specifically, we choose a water system with 14,910,336 atoms and a copper system with 2,540,160 atoms as the benchmark on the new Sunway to fully utilize the memory capacity of each node (96GB). Fig.13 and 14 The parallel efficiency of the water system and the copper system are shown, which are normalized to the efficiency of 50 nodes and 20 nodes, respectively. For the water system / copper system, the parallel efficiency of 4,560 nodes on the new Shenwei reaches 41.2% and 23.1%, respectively. Compared with the implementation on Summit and Fugaku, we found that all machines showed good scalability up to 570 nodes with a parallel efficiency of more than 60%. In addition, all machines showed the potential to scale to 4,560 nodes.

[0170] Fig.15: Weak scalability of the new Shenwei water and copper systems. The redesigned DeePMD-kit can achieve simulations of 3.5 billion atoms on the entire new Shenwei machine, reaching a peak performance of 68.9PFLOPS, and for the copper system, it can achieve simulations of 1.6 billion atoms, reaching a peak performance of 49.5PFLOPS.

[0171] E. Weak scalability

[0172] We evaluate the weak scalability of the optimized DeePMD-kit and report the maximum simulation scale (i.e., number of atoms) and achievable floating-point operation performance (FLOPS) for water and copper systems at different node numbers, e.g. Fig.15 As shown. According to the weak scalability results, the optimized DeePMD-kit can achieve 41.4PFLOPS on a copper system containing 13 billion atoms, and further expand to 57.1PFLOPS on a water system containing 29 billion atoms. Through further estimation of the entire machine (i.e., 107,520 nodes), the redesigned DeePMD-kit implementation can perform ultra-large-scale simulations of 35 billion atoms on the new Shenwei, reaching 68.9PFLOPS (about 5% of the theoretical peak performance). It should be noted that the main reason for the performance utilization is the bandwidth limitation of the custom operators in DeePMD-kit (including Tabulate, TabulateGrad, ProdEnvMat, ProdVirial, and ProdForce). Through further performance analysis of a single new Shenwei node, these DeePMD-kit operators can achieve an average bandwidth of 262.9GB / s (about 86% DMA bandwidth) for the water system and 257.4GB / s (about 84% DMA bandwidth) for the copper system. We point out that the SW26010-pro processor is only equipped with 300GB / s DDR4 memory, which is about a quarter of the current HBM memory bandwidth. If equipped with HBM, our work would achieve 263.6PFLOPS (about 20% peak) without any code modification, which would be comparable to our previous work on Fugaku (124.8PFLOPS from 442PFLOPS, 28% peak).

[0173] Those skilled in the art will readily appreciate other embodiments of the present invention after considering the specification and practicing the invention herein. This application is intended to cover any variations, uses or adaptations of the present invention that follow the general principles of the present invention and include common knowledge or customary techniques in the art that the present invention did not invent. The specification and examples are intended to be exemplary only, and the true scope and spirit of the present invention are indicated by the appended claims.

[0174] It should be understood that the present invention is not limited to the exact construction that has been described above and shown in the drawings and that various modifications and changes may be made without departing from the scope thereof. The scope of the present invention is limited only by the appended claims.

Claims

1. The operation method of the DeePMD-kit model on the Shenwei supercomputer is characterized in that: The six core groups CG of the Shenwei supercomputer are responsible for the reasoning of the six sub-areas of the MPI process respectively; the programming interface of the Shenwei supercomputer is SACA; Each of the core groups CG includes a management processing unit MPE connected to the main memory DDR and 64 computing processing units CPE, the management processing unit MPE contains a 256-bit SIMD unit; each computing processing unit CPE contains a 512-bit SIMD unit; each computing processing unit CPE contains a local data memory LDM, each local data memory LDM cancels the cache configuration, and each local data memory LDM realizes continuous data transmission through the DMA strategy and the main memory DDR; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA; Among them, the reasoning method of each core group CG is: SACA is used to accelerate the time-consuming operators in the DeePMD-kit model operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive parts of the time-consuming operators to 64 computation processing units CPE according to the atomic ID. Each computation processing unit CPE calculates the corresponding divided blocks respectively. After the calculation processing unit CPE completes the calculation, the management processing unit MPE executes the athread join() instruction to synchronize the calculation results of the 64 calculation processing units CPE, and stores the results in the main memory DDR; SIMD parallelism during SACA runtime: 512-bit vectorized instructions are issued by the computing processing unit CPE to improve the parallel efficiency of custom operators in the DeePMD-kit model operators, including Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial.

2. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: The molecular system is evenly divided into multiple sub-regions according to the active MPI processes; the Sunway supercomputer starts 6 MPI processes on the SW26010-pro node, each core group CG holds an MPI process and is responsible for calculating one sub-region of all atoms.

3. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: The computational efficiency and bandwidth utilization of the Tabulate operator and TabulateGrad operator are optimized by overlapping computation with memory access and kernel fusion. The optimization process is as follows: Step 1, each computing processing unit CPE issues a CRTS_dma_iget() or CRTS_dma_iput() instruction to initiate an asynchronous memory access request; after the DMA request is initiated, the computing processing unit CPE synchronously waits for memory access and checks the DMA status by executing the dma_wait() instruction. Before executing the dma_wait() instruction, the calculations of the Tabulate operator and the TabulateGrad operator will not be blocked, and calculations that have no data dependency on the current calculation pipeline are allowed to be executed before executing the dma_wait() instruction; The asynchronous DMA transfer between the local data memory LDM and the main memory DDR is synchronized by executing the dma_wait() instruction during the manual control of the overlapping efficiency; Each computing processing unit CPE is responsible for executing multiple descriptor vectors D(R i ), input s(r ij ) and its corresponding table coefficients are divided into two buffers, each of which independently processes a portion of the descriptor vector D(R i ), the processing of the buffer includes calculation and asynchronous DMA access between the local data memory LDM and the main memory DDR, and the calculation data of the two buffers are synchronized through the dma_wait() instruction; When the first buffer completes the DMA transfer from the main memory DDR to the local memory and starts computing, the second buffer initiates a DMA request to load s(r ij ) and table coefficients; when the second buffer completes the DMA request and starts calculation, the first buffer request will be D(R i ) written back to main memory; Step 2: The TabulateGrad operator uses the G matrix previously calculated by the Tabulate operator.

4. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: The computational efficiency of the ProdEnvMat operator is optimized through data layout transformation and branch elimination. The optimization process is as follows: Step 1: Convert the AoS variable into the SoA format, perform numerical shuffling through the VSHFW of the Sunway supercomputer, and store it back into the AoS format after vectorization. In VSHFW, the source vector is divided into 16 32-bit words, the shuffle mask includes 16 5-bit components, and each component is responsible for specifying the word position in the source vector. By executing VSHFW, the 16 words of the target vector are allocated according to the mask; Step 2: Calculate a conditional mask for the SIMD vector, multiply the SIMD vector by the conditional mask, and derive the branch output.

5. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: In the CPE array architecture, each cluster includes four computing processing units CPE and 16 slave cluster management SCMs. The 16 slave cluster management SCMs in each CPE array are interconnected through an on-chip network. Each slave cluster management SCM includes an RMA engine and a DMA engine.

6. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 5, characterized in that: The write conflict of the ProdVirial operator is optimized through replication and reduction strategies. The optimization process is as follows: Step 1: Create a virial replica named virial_rep for each computing processing unit CPE and initialize its elements to 0; Calculate the increment tmp_v of each computing processing unit CPE on virial and add it to virial_rep; After each computing processing unit CPE calculates virial_rep, it synchronizes the calculation results of 64 computing processing units CPE by executing the CRTS_ssync_array() instruction; Step 2: The first computing processing unit CPE in each cluster aggregates the partial values ​​through two rounds of operations that are quickly executed by the RMA engine in the cluster; The first computational processing unit CPE in the first column of the CPE array derives the intra-row partial reduction by performing the RMA engine reduction on the clusters sharing the same row number; The first computing processing unit CPE in the first column performs partial reduction through the RMA engine, and the first computing processing unit CPE of the first cluster obtains the final reduction result and writes it back to the main memory DDR through the DMA strategy.

7. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: The TensorFlow operators are optimized by fitting network optimization and parallelizing tensor operators. The optimization process of TensorFlow operators is as follows: Step 1, optimizing the fitting network: the fitting network includes a plurality of fully connected layers and a Tanh activation function, which are used to obtain corresponding gradients for force calculation during the back propagation process; the fully connected layer of the Shenwei supercomputer includes a MatMul operator and an Add operator; Fuse the MatMul operator and the Add operator into the GEMM in Eigen, and use swBLAS in the Shenwei supercomputer to synchronously accelerate the Tanh and TanhGrad operators through SACA; Step 2: Parallelize tensor operators: Parallelize Pad and Slice operators through SACA, and use DMA strategy to optimize the bandwidth of main memory DDR.

8. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: The bandwidth of the Tabulate operator is optimized through the mixed-precision tabulatetion method. The mixed-precision tabulatetion method is as follows: Step 1: Create low-order weights a0, a1, a2, and a3 using bfloat16, and high-order weights a4 and a5 using float32. Step 2: Design a mixed-precision Tabulate kernel using SACA and vectorized intrinsics.

9. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 8, characterized in that: The calculation workflow of the optimized Tabulate operator is as follows: Step 1. Calculation and with Perform mixed precision vector-vector multiplication to obtain the activation matrix a ij ; Step 2: w ij with a ij Perform matrix-vector multiplication to obtain the descriptor vector

Images