Running Method of DeePMD-kit Model on Sunway Supercomputer

Abstract

The present invention relates to the field of computer technology, and particularly to a method for running the DeePMD-kit model on the Shenwei supercomputer. Six core groups CG of the Shenwei supercomputer are respectively responsible for the inference of six sub-regions of MPI processes. Each core group CG includes an MPE and 64 CPEs. The MPE contains a 256-bit SIMD unit; each CPE contains a 512-bit SIMD unit; each CPE contains a local data memory, and each local data memory LDM realizes continuous data transmission with the main memory through the DMA strategy; data transmission between CPEs in the same core group CG is realized through remote memory access RMA; the inference method of each core group CG is as follows: time-consuming operators are accelerated by SACA, and the MPE distributes the computationally intensive part of the time-consuming operators to 64 CPEs; after the calculation is completed, the MPE synchronizes the calculation results of the 64 CPEs and stores the results in the main memory; SIMD parallelism during the operation of SACA: 512-bit vectorized instructions are issued by the CPE. Through the above technical solution, the present invention has the advantage of enabling the DeePMD-kit model to be applied on 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. The matrix is a sub-matrix of < , representing the vector of the first M rows. The environmental matrix

[0007]

[0008] where r ij = r j - r i , is the relative distance between atom i and neighbor atom j, s(r ij ) = w(r ij ) / |r ij |, where w(r ij ) is expressed as a gate function that smoothly decays from 1 to 0 when |r ij | < r c . The matrix is generated by fitting s(r ij ) with an embedding network and is 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 high-dimensional functional relationship between the local environmental features obtained from D and the atomic energy contribution E 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 fully connected networks and uses Tanh as the activation function of the network. By combining D and N, the DeePMD-kit model can achieve ab initio accuracy by fitting ab initio data generated by scientific computing software such as VASP, Quantum Espresso, and PWmat. The force operation of atom i is obtained during the backpropagation process (as shown in Figure 1 (f)) and is expressed as:

[0014]

[0015] The virial operation of the system, as shown in Figure 1 (g), is achieved through the calculation of the virial operator in the model and is expressed as:

[0016] Ξ = ∑ i Ri F i ;

[0017] The floating-point operations embedded in the network account for more than 95% of the overall FLOPs. A model compression algorithm is designed to reduce the overall computational amount, and this module is expressed as:

[0018]

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

[0020] However, for the computers used in the existing DeePMD-kit models, the processors used are usually CPU processors or GPU processors.

[0021] The new-generation Sunway supercomputer is equipped with a new-generation heterogeneous multi-core processor (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 an SW26010-pro processor, which contains 6 core groups (CGs), and the theoretical peak performance under double-precision computing is about 14 TFLOPS. Each core group contains a management processing unit (MPE) and an 8×8 computing processing unit (CPE) grid, connected to 16 GB DDR4 main memory with a bandwidth of 51.2 GB / s. Each CPE has a 256 KB manually controlled local data memory (LDM). The MPE and CPE respectively have 256-bit and 512-bit SIMD units.

[0022] The SW26010-pro processor supports direct memory access (DMA), and can transfer continuous data between the LDM and the main memory, with a theoretical bandwidth of 307 GB / s. In addition, data transfer between CPEs within the same core group is carried out through remote memory access (RMA), with a bandwidth of about 400 GB / s. The programming interface of the new-generation Sunway 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, migrating the DeePMD-kit model to the new Sunway supercomputer still faces many challenges.

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

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

[0026] Instead of optimizing main memory access through multi-level caches like CPUs, the SW26010-pro equips each CPE with a programmable fast LDM, similar to a manually controlled L1 cache, which makes the cache policy more flexible and memory access more configurable. In addition, different from GPUs that use shared memory for thread data sharing, or CPUs that rely on caches for inter-thread communication, the SW26010-pro does not provide physically shared memory for CPEs. Instead, data exchange between different CPEs is managed through the RMA mechanism, and this uniqueness creates a unique programming framework. We should make full use of these characteristics of the SW26010-pro architecture to obtain excellent performance.

[0027] Currently, the resource utilization rate of DeePMD-kit operators is quite low. Different from the implementations using OpenMP (GPU) and CUDA (GPU), the design of MPE and CPE codes requires careful consideration. The MPE calls 64 CPEs for calculation through the \textit{athread_spawn()} instruction, and each CPE performs calculations and memory access simultaneously. However, previous work did not optimize the DeePMD-kit operators for the SW26010-pro architecture, resulting in all operators being executed only on the MPE. In addition, some new features of CPEs have not been effectively utilized, including memory access through DMA, on-chip data sharing through RMA, and 512-bit vectorization. As Figure 3 shown in the performance analysis of the DeePMD-kit time-consuming operators, the peak performance and memory bandwidth of each implemented 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 computationally intensive custom TensorFlow operators and official operators; the custom operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial; the 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 Shenwei supercomputer to ensure high resource utilization and computational efficiency. Summary of the Invention

[0030] To solve the problems in the related art, this application provides a DeePMD-kit model applied to the SW26010-pro processor of the Shenwei supercomputer, which solves the problem that the existing DeePMD-kit model cannot be applied to the SW26010-pro processor in the new generation of Shenwei supercomputers.

[0031] The technical solution is as follows:

[0032] A method for running the DeePMD-kit model on the Shenwei supercomputer, characterized in that six core groups CG of the Shenwei supercomputer are respectively responsible for the inference of six sub-regions 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 with the main memory DDR through the DMA strategy; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA;

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

[0035] Accelerate time-consuming operators through SACA. The time-consuming operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive part of the time-consuming operators to 64 computing processing units CPE according to the atomic ID, and each computing processing unit CPE calculates the corresponding divided blocks;

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

[0037] SIMD parallelism during SACA operation: The parallel efficiency of custom operators is improved by the computing processing unit CPE issuing 512-bit vectorized instructions. The custom operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, and ProdVirial.

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

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

[0040] Preferably, the calculation efficiency and bandwidth utilization rate of the Tabulate operator and the TabulateGrad operator are optimized by overlapping the calculation and 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 the 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 are not blocked, and calculations that have no data dependence on the current calculation stream are allowed to be executed before executing the dma_wait() instruction;

[0042] During the asynchronous DMA transfer between the local data memory LDM and the main memory DDR, synchronization is achieved by manually controlling the overlapping efficiency through executing the dma_wait() instruction;

[0043] Each computing processing unit CPE is responsible for executing multiple descriptor vectors D(R i ) and input s(r ij) and its corresponding table coefficients are divided into two buffers, and each buffer independently processes and calculates part of the descriptor vector D(R in the compute processing element CPE i ) The processing of the buffer includes calculation and asynchronous DMA access between the local data memory LDM and the main memory DDR. The calculation data of the two buffers is 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 to calculate, the second buffer initiates a DMA request to load s(r ij ) and the table coefficients; when the second buffer completes the DMA request and starts to calculate, the first buffer requests to write D(R i ) back to the main memory;

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

[0046] Through the above technical solution, through the setting of overlapping calculation and memory access, because the Sunway supercomputer supports high-bandwidth asynchronous DMA, it can efficiently access continuous data blocks in the main memory, enabling the operator to maximize the DMA bandwidth of 300GB / s and improving the throughput;

[0047] By enabling the TabulateGrad operator to reuse the G matrix previously calculated by the Tabulate operator, frequent memory access during the calculation of the G matrix can be omitted, thereby eliminating the redundant calculation of the Tabulate operator, thus saving 54.6% of the execution time.

[0048] Preferably, the calculation efficiency of the ProdEnvMat operator is optimized through data layout conversion and branch elimination. The optimization process is as follows:

[0049] Step 1, Convert the AoS variable to the structure array SoA format, perform numerical shuffling through the VSHFW of the Sunway supercomputer, and store it back in 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 the conditional mask for the SIMD vector, multiply the SIMD vector by the conditional mask, and then 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 replication and reduction strategies, the integrity of the code before the CRTS_ssync_array() instruction can be ensured. Through the three-step RMA-based reduction method, redundant global synchronization can be avoided, thus improving the parallel efficiency.

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

[0064] 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 backpropagation process; the fully connected layers of the Shenwei supercomputer include MatMul operators and Add operators;

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

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

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

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

[0069] Preferably, the bandwidth of the Tabulate operator is optimized by the tabulation method of mixed precision. The tabulation method of mixed precision is as follows:

[0070] Step 1, create low-order weights a0, a1, a2, and a3 through bfloat16, and create high-order weights a4 and a5 through float32;

[0071] Step 2, design a mixed-precision Tabulate kernel through SACA and vectorized intrinsics.

[0072] Through the above technical solution, directly fitting the DeePMD-kit descriptor from scratch using low-precision table weights in mixed precision can minimize the compression error, and creating low-order weights using bfloat16 can effectively prevent overflow problems during the training process; by leveraging the highly efficient low-precision computing power of the CPE equipped in the Sunway supercomputer, this can save approximately 3 times the memory occupancy during the calculation of the Tabulate operator, thereby increasing the throughput by approximately 23 times.

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

[0074] Step 1, calculate and perform a vector-vector multiplication in mixed precision with to obtain the activation matrix a ij ;

[0075] Step 2, w ij and a ij perform a matrix-vector multiplication to obtain the descriptor vector

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

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

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

[0079] 2. Through performance optimization of operators such as Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, MatMul, Tanh, TanhGrad, Slice, and Pad, the utilization rate of the new generation of Sunway resources can be maximized.

[0080] It should be understood that the above general description and the following detailed description are only exemplary and do not limit the present invention. Description of the Drawings

[0081] The accompanying drawings here are incorporated into the specification and form a part of this specification, showing embodiments in accordance with the present invention, and are used together with the specification to explain the principles of the present invention.

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

[0083] Figure 2 It is a schematic diagram of the architecture comparison between the processor in the prior art and the processor of the Shenwei supercomputer;

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

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

[0086] Figure 5 It is a schematic diagram of the operation of the overlap of computing and memory access on the Shenwei 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 in the present invention;

[0089] Figure 8 It is a schematic diagram of the operation of the mixed-precision tabulation scheme in the present invention;

[0090] Figure 9 It is a schematic diagram of the operation of the radial distribution functions gOO(r), gOH(r) and gHH(r) of liquid water in the present invention;

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

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

[0093] Figure 12 It is a schematic diagram of the performance evaluation of custom operators in the DeePMD-kit model in the present invention;

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

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

[0096] Figure 15 It is a schematic diagram of the running of the weak scalability of the water and copper systems in the Shenwei supercomputer in the present invention; Detailed implementation manners

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

[0098] As shown in the appended Figures 1-15 figures, the running method of the DeePMD-kit model on the Shenwei supercomputer. The Shenwei supercomputer starts 6 MPI processes on the SW26010-pro nodes. Each core group CG holds one MPI process and is responsible for calculating a sub-region of one of all the atoms; the molecular system is evenly divided into multiple sub-regions according to the active MPI processes; the six core groups CG of the Shenwei supercomputer are respectively responsible for the inference of the six sub-regions of the MPI processes; the programming interface of the Shenwei 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. Each local data memory LDM cancels the cache configuration, and each local data memory LDM realizes continuous data transmission with the main memory DDR through the DMA strategy; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA;

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

[0101] Accelerate time-consuming operators through SACA. The time-consuming operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive part of the time-consuming operators to 64 computing processing units CPE according to the atomic ID, and each computing processing unit CPE calculates the corresponding divided blocks respectively;

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

[0103] SIMD parallelism during SACA operation: Improve the parallel efficiency of custom operators by the computing processing unit CPE issuing 512-bit vectorized instructions. The custom operators include 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 operators can be improved.

[0105] As Figure 5 shown, the Tabulation method uses a fifth-order polynomial approximation to compress the embedding 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 Sunway is time-consuming. Analysis of the Tabulate and TabulateGrad operators shows that the memory access efficiency is low when loading polynomial coefficients (i.e., the memory bandwidth utilization rate is less than 1%), occupying 68.6% and 54.6% of the execution time respectively; in this specific embodiment, optimize the calculation efficiency and bandwidth utilization rate of the Tabulate operator and the TabulateGrad operator by overlapping calculation and memory access and kernel fusion. The optimization process is as follows:

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

[0107] Each computing processing element 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 element CPE synchronously waits for the 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 are not blocked, and it is allowed to execute calculations that have no data dependencies with the current computing pipeline before executing the dma_wait() instruction;

[0108] During the process of manually controlling the overlap efficiency of the asynchronous DMA transfer between the local data memory LDM and the main memory DDR, synchronization is achieved by executing the dma_wait() instruction; 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 make full use of the LDM space (for example, 32 in the water system, 8 in the copper system, using 223KB LDM). Through this overlap strategy, these operators can maximize the DMA bandwidth of 300GB / s and improve the throughput.

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

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

[0111] Step 2: Both the Tabulate and TabulateGrad operators calculate the G matrix described in the equation. However, calculating 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] Through the setting of overlapping computation and memory access, because the Shenwei supercomputer supports high-bandwidth asynchronous DMA, it can efficiently access continuous data blocks in the main memory, enabling the operator to maximize the DMA bandwidth of 300 GB / s and improving throughput.

[0113] By enabling the TabulateGrad operator to reuse the G matrix previously calculated by the Tabulate operator, frequent memory access during the calculation of the G matrix can be omitted, thereby eliminating redundant calculations in the Tabulate operator, and saving 54.6% of the execution time.

[0114] As Figure 6 shown, DeePMD-kit is implemented as a force field plugin for LAMMPS. Its variable layout must be unified with LAMMPS and adopt the structure-of-arrays (AoS) format. For example, the r of two atoms ij is stored as (x0,y0,z0),(x1,y1,z1), where the memory interval between the 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 the AoS format, which results in non-contiguous memory access for each iteration, hindering vectorization optimization. In addition, the ProdEnvMat operator contains many conditional branches, resulting in branch divergence and reducing the computational efficiency of SIMD. Therefore, optimize the computational efficiency of the ProdEnvMat operator through data layout transformation and branch elimination. The optimization process is as follows:

[0115] Step 1: Convert the AoS variables to the structure-of-arrays of structures (SoA) format, perform numerical shuffling through the VSHFW of the Shenwei supercomputer, and store them back in 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 perform numerical shuffling. 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 to place it in the target vector. By executing VSHFW, the 16 words of the target vector are assigned according to the mask;

[0117] In this specific embodiment, an example of a 256-bit VSHFW instruction is shown. This instruction 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 to select, and the remaining four bits indicate which word in the source vector to select. As Figure 6 shown, the fourth binary code of the mask is 0x13. This binary code indicates that the corresponding word of the target vector is obtained from the third word of source vector 1. For variables with a memory interval of 6, we can use 24 VSHFW instructions to perform conversions between AoS and SoA, such that only 24 additional instruction cycles can increase the throughput by 4 times with 256-bit VSHFW and by 8 times with 512-bit VSHFW.

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

[0119] The ProdEnvMat operator uses conditional branches to eliminate redundant calculations. For each atom, ProdEnvMat retrieves the global atom indices of its neighbors 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 proceeds to retrieve 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 lead to branch divergence. To solve this problem, in this specific embodiment, a branch elimination strategy is implemented to reduce the divergence. Specifically, we calculate a conditional mask for the SIMD vector and derive the branch output by multiplying the SIMD vector with this mask.

[0120] Through the setting of data layout conversion, the parallel loading, storage, and computing capabilities of the Sunway supercomputer can be fully utilized, and thus efficient numerical shuffling can be achieved, so that the throughput can be increased by 4 times with 256-bit VSHFW and by 8 times with 512-bit VSHFW;

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

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

[0123] As Figure 7As shown, when adopting the in-operator parallelism scheme on the new Sunway, write conflicts exist in the ProdVirial operator, and these conflicts are usually alleviated through atomic operations. The atomic operations on the new Sunway are implemented through the fetch-and-add mechanism, and each atomic operation requires global synchronization across all CPEs of the CG. However, using these atomic operations introduces 9*N {des} *N {loc} *N {nnei} synchronizations in the ProdVirial operator, where N {des} represents the number of descriptors, N {loc} represents the number of atoms, and N {nnei} represents the number of neighbor atoms. Frequent synchronization operations reduce the parallel efficiency, resulting in its performance being inferior to the serial implementation on the MPE. Therefore, in this specific embodiment, the write conflicts of the ProdVirial operator are optimized through the replication and reduction strategy, and the optimization process is as follows:

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

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

[0126] After each computing processing element CPE calculates virial_rep, synchronize the calculation results of 64 computing processing elements CPE by executing the CRTS_ssync_array() instruction;

[0127] Step 2: Through operations quickly executed within the cluster by the two-round RMA engine, the first computing processing element CPE in each cluster aggregates partial values;

[0128] Through the reduction of the RMA engine for clusters sharing the same row number, the first computing processing element CPE in the first column of the CPE array exports the in-row partial reduction;

[0129] The first computing processing element CPE in the first column performs partial reduction through the RMA engine, and the first computing processing element CPE in 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 replication and reduction strategy, the integrity of the code before the CRTS_ssync_array() instruction can be ensured. Through the three-step RMA-based reduction method, redundant global synchronization can be avoided, thereby improving the parallel efficiency.

[0131] The optimization of TensorFlow for the new Sunway architecture is insufficient, resulting in low efficiency of some official TensorFlow operators 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 backpropagation process.

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

[0134] This introduces redundant data transfer of intermediate tensors and additional kernel launches.

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

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

[0137] By fusing the MatMul operator and the Add operator into GEMM in Eigen, using 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 operators in DeePMD-kit;

[0138] Some specialized tensor operators, such as Pad and Slice, are currently only implemented in the MPE version on the new Sunway. Unfortunately, these operators do not support parallel execution on the CPE and fail to effectively utilize the 300GB / s DMA bandwidth. Therefore, the computing performance of DeePMD-kit is significantly affected. To solve these problems, parallelizing the Pad and Slice operators through SACA can maximize the utilization of the main memory DDR bandwidth, thereby optimizing the TensorFlow operators.

[0139] As Figure 8 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 characteristic dimension of the embedding matrix. However, each x needs to access M×6×8 bytes of w {ij} , which consumes a large amount of memory access and results in low computational efficiency in large-scale MD simulations. To alleviate the bandwidth bottleneck of the Tabulate operator, the bandwidth of the Tabulate operator is optimized by a mixed-precision tabulation method, and the mixed-precision tabulation method is as follows:

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

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

[0142] Different from the low-precision truncation method, the proposed mixed-precision method directly fits the DeePMD-kit descriptor from scratch using low-precision table weights, which can minimize the compression error. Through the statistical analysis of each item in the table weights, we found that the high-order coefficients (i.e., a4 and a5) show significant value ranges during the training phase, indicating that these items 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 the high-order coefficients are much larger than the representable range of float16, directly using float16 to compress the weights will cause overflow problems. Even using bfloat16 to compress the 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 to create low-order weights instead of float16 is that the value range of bfloat16 is consistent with that of float32, which can prevent overflow problems during the training process. The new Sunway supercomputer supports low-precision floating-point calculations and 512-bit SIMD support, including float16, bfloat16, and float32. Therefore, we can use SACA and vectorized intrinsics to design a mixed-precision Tabulate kernel to utilize the efficient low-precision computing power of the CPE.

[0144] Directly fitting the DeePMD-kit descriptor from scratch by using low-precision table weights in mixed precision can minimize the compression error. Creating low-order weights using bfloat16 can effectively prevent overflow problems during the training process. By leveraging the highly efficient low-precision computing power of the CPE equipped on the Sunway supercomputer, this can save approximately 3 times the memory occupancy during the calculation of the Tabulate operator, thereby increasing the throughput by about 23 times.

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

[0146] Step 1: Calculate and perform a vector-vector multiplication in mixed precision with to obtain the activation matrix a ij ;

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

[0148] By calculating a before the time-intensive loop of calculating ij and converting the Tabulate operation into two matrix-vector multiplications, this reduces M×38 redundant calculations, thereby optimizing the calculation 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 the AIMD method, because it requires a balance between weak non-covalent intermolecular interactions, thermal (entropy) effects, and nuclear quantum effects.

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

[0152] To simulate and analyze the performance of our work, we used a Tabulate-based DP model as a benchmark, where the descriptor is a compressed Tabulate module with an output dimension of 128. The dimension of the fitting network is 224×224×224. The cutoff radii are set to and The maximum number of neighbors is 138 and 512. The MD equations are numerically integrated for 100 steps using the Velocity-Verlet method with 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 a buffer region. 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 the mixed-precision scheme

[0155] 1. Energy and force errors. The accuracy of DeePMD-kit is investigated 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 precision, which consists of 100 water molecules. First, we predict the energy and forces of these 100 water molecules using the mixed-precision model, and then calculate the root-mean-square error (RMSE) of the predicted energy and forces with respect to the AIMD results. Table I shows the RMSE errors 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 of 2.43×10^-3 eV / molecule and a force error. The energy prediction error of the mixed-precision scheme is much less than 4×10^-2 eV / molecule, and the force error is no less than This indicates that the proposed mixed-precision method has a lower error rate.

[0156] Figure 9 For the radial distribution functions gOO(r), gOH(r), and gHH(r) of liquid water; calculated by two DeePMD-kit implementations respectively: 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 the mixed-precision, we calculate the radial distribution function (RDF) of water, which is the normalized probability of searching for neighboring atoms at a 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 select 192 water atoms and perform 20,000-step MD simulations using the proposed mixed-precision DeePMD-kit and the baseline. Then, we calculate the RDF of the water system based on the simulation results and plot the RDF curves in Figure 9 According toFigure 9 , the RDF of the proposed mixed-precision tabulation scheme and the baseline method completely overlaps. The baseline method has been proven in previous studies to achieve accuracy comparable to that of the AIMD method. Therefore, we can conclude that the proposed mixed-precision scheme can accurately predict physical observables.

[0158] Figure 10 is the system energy curve for the water system during a 1,000,000-step simulation. The mixed-precision scheme can remain stable during long-time-step simulations and has a low error rate.

[0159] 3. Stability assessment: To evaluate the MD stability of the mixed-precision scheme, we selected a water system containing 192 atoms and performed long-term simulation evaluations on the proposed mixed-precision scheme. We performed a million-step NVE simulation on the proposed mixed-precision DeePMD-kit and the baseline method using the same configuration. Then, we observed the energy fluctuations of the system during the MD simulation and plotted the energy curves calculated by these two DeePMD-kit schemes in Figure 10 . According to Figure 10 , we can observe that the system energy is well maintained during long-term MD simulations. In addition, the system energy simulated by the mixed-precision scheme is very close to the baseline results and follows the same trend, indicating that the mixed-precision method can perform long-term MD simulations with a low error rate.

[0160] Figure 11 is the step-by-step performance improvement for a 9,126-atom water system and a 5,184-atom copper system on a single SW26010-pro node. The performance speedup is normalized with respect to the baseline.

[0161] B. Single-node performance

[0162] Our work on the new Sunway supercomputer includes the optimization 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 by priority as follows: 1 fitting network (68.1% of the total time), 2 custom operators (16.8%), and 3 tensor operators (13.9%). We optimized step by step in this order and reported the corresponding performance improvements on a single new Sunway node (one SW26010-pro processor). We performed 100-step MD simulations and measured the corresponding cycle times (calculated by LAMMPS) of a water system with 9,126 atoms and a copper system with 5,184 atoms. The benchmark used is the planar MPI version DeePMD-kit of the compression model proposed in ppopp 22. We performed MD simulations under the same configuration and calculated the speedup of each optimization. Figure 11 Shows the performance improvements of DeePMD-kit on a single new Sunway node. Compared with the benchmark, the optimized DeePMD-kit achieved a 3.1-fold speedup on the water system and a 2.9-fold speedup on the copper system. Each layer of the fitting network includes MatMul, Add, Tanh, and TanhGrad operators. These results indicate that the proposed kernel fusion and Tanh optimization methods can effectively improve the computational performance of DeePMD-kit. By performing custom operator optimization, the optimized DeePMD-kit achieved a 6.8-fold speedup on the water system and a 12.1-fold speedup on the copper system, indicating 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 performing mixed-precision optimization on tabulation, our work achieved an 8.2-fold speedup on the water system and a 14.1-fold speedup on the copper system, indicating that using low precision can further improve the computational efficiency of CPE by maximizing the bandwidth utilization of DMA. Finally, after optimizing the tensor operators, our work achieved a significant 67.6-fold speedup on the water system and a 56.5-fold speedup on the copper system. This interesting phenomenon indicates that the pre-installed TensorFlow library is insufficiently optimized for some tensor operators, seriously affecting the performance of DeePMD-kit.

[0163] Figure 12 Performance evaluation of custom operators in the DeePMD-kit model. We evaluated the performance improvements of custom operators on a 9,216-atom water system and a 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 Figure 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. Figure 12 The 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] Figure 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] Figure 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). Figure 13 and 14Shows the parallel efficiency of the water system and the copper system, which are normalized to the efficiency of 50 nodes and 20 nodes respectively. For the water system / copper system, on the new Sunway, the parallel efficiency of 4,560 nodes reaches 41.2% and 23.1% respectively. Compared with the implementations on Summit and Fugaku, we find that all machines show good scalability with up to 570 nodes and the parallel efficiency exceeds 60%. In addition, all machines show the potential to scale to 4,560 nodes.

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

[0171] E. Weak scalability

[0172] We evaluated the weak scalability of the optimized DeePMD-kit and reported the maximum simulation scale (i.e., the number of atoms) and the achievable floating-point operation performance (FLOPS) of the water system and the copper system with different numbers of nodes, as Figure 15 shown. According to the weak scalability results, the optimized DeePMD-kit can reach 41.4 PFLOPS on the copper system with 13 billion atoms and further scale to 57.1 PFLOPS on the water system with 29 billion atoms. Through further estimation of the whole machine (i.e., 107,520 nodes), the redesigned DeePMD-kit implementation can perform ultra-large-scale simulations of 35 billion atoms on the new Sunway, reaching 68.9 PFLOPS (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 Sunway node, the average achievable bandwidth of these DeePMD-kit operators is 262.9 GB / s (about 86% of the DMA bandwidth) for the water system and 257.4 GB / s (about 84% of the DMA bandwidth) for the copper system. We point out that the SW26010-pro processor is only equipped with 300 GB / s of DDR4 memory, which is about a quarter of the current HBM memory bandwidth. If equipped with HBM, our work will reach 263.6 PFLOPS (about 20% of the peak) without modifying any code, which will be comparable to our previous work on Fugaku (124.8 PFLOPS out of 442 PFLOPS, 28% of the peak).

[0173] Other embodiments of the present invention will be readily apparent to those skilled in the art upon consideration of the specification and practice of the invention herein. This application is intended to cover any variations, uses, or adaptations of the invention following the general principles of the invention and including known or customary techniques in the art not invented by the present invention. The specification and examples are to be considered exemplary only, and the true scope and spirit of the invention are pointed out by the appended claims.

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

Claims

1. Method for running the DeePMD-kit model on the Shenwei supercomputer, characterized in that, Six core groups CG of the Shenwei supercomputer are respectively responsible for the inference of six sub-regions of the MPI process; the programming interface of the Shenwei supercomputer is SACA; 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. Each local data memory LDM cancels the cache configuration, and each local data memory LDM realizes continuous data transmission with the main memory DDR through the DMA strategy; the computing processing units CPE in the same core group CG realize data transmission through remote memory access RMA; Among them, the inference method of each core group CG is as follows: Accelerate the time-consuming operators in the DeePMD-kit model operator through SACA. The time-consuming operators include Tabulate, TabulateGrad, ProdEnvMat, ProdForce, ProdVirial, Slice, and Pad. The management processing unit MPE distributes the computationally intensive part of the time-consuming operator to 64 computing processing units CPE according to the atomic ID, and each computing processing unit CPE calculates the corresponding divided block; After the computing processing unit CPE finishes the calculation, the management processing unit MPE executes the athreadjoin() instruction to synchronize the calculation results of the 64 computing processing units CPE and stores the results in the main memory DDR; SIMD parallelism during the operation of SACA: Improve the parallel efficiency of the customized operators in the DeePMD-kit model operator by issuing 512-bit vectorized instructions through the computing processing unit CPE. The customized operators include 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 Shenwei supercomputer starts 6 MPI processes on the SW26010-pro node, and each core group CG holds one MPI process and is responsible for calculating a sub-region of all atoms.

3. The operating method of the DeePMD-kit model according to claim 1 on the Sunway supercomputer, characterized in that, Optimize the calculation efficiency and bandwidth utilization rate of the Tabulate operator and the TabulateGrad operator by overlapping calculation and 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 the 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 are not blocked, and calculations that have no data dependencies with the current calculation pipeline are allowed to be executed before executing the dma_wait() instruction. During the asynchronous DMA transfer between the local data memory LDM and the main memory DDR, synchronization is achieved by executing the dma_wait() instruction during the process of manually controlling the overlapping efficiency. Each computing processing unit CPE is responsible for executing multiple descriptor vectors D(R i ), dividing the input s(r ij ) and its corresponding table coefficients into two buffers, and each buffer independently processes part of the descriptor vector D(R i ) in the computing processing unit CPE. The processing of the buffer includes calculations and asynchronous DMA access between the local data memory LDM and the main memory DDR. The calculation data of the two buffers is synchronized through the dma_wait() instruction; When the first buffer finishes 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 the table coefficients; when the second buffer finishes the DMA request and starts computing, the first buffer requests to write D(R j ) back to the main memory; Step 2: The TabulateGrad operator uses the G matrix previously calculated by the Tabulate operator.

4. The operating method of the DeePMD-kit model according to claim 1 on the Sunway supercomputer, characterized in that, Optimize the calculation efficiency of the ProdEnvMat operator through data layout transformation and branch elimination. The optimization process is as follows: Step 1: Convert the AoS variable to the structure of array SoA format, perform numerical shuffling through the VSHFW of the Sunway supercomputer, and store it back in 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 the conditional mask for the SIMD vector, and multiply the SIMD vector by this conditional mask to deduce the branch output.

5. The operating method of 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 subordinate cluster managers SCM. The 16 subordinate cluster managers SCM in each CPE array are interconnected through the on-chip network, and each subordinate cluster manager SCM includes an RMA engine and a DMA engine.

6. The operating method of the DeePMD-kit model on the Sunway supercomputer according to claim 5, characterized in that, Optimize the write conflict of the ProdVirial operator through the copy and reduction strategy. The optimization process is as follows: Step 1: Create a virial copy 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 the virial and add it to virial_rep. After each computing processing unit CPE calculates virial_rep, synchronize the calculation results of 64 computing processing units CPE by executing the CRTS_ssync_array() instruction. Step 2: Through two rounds of operations quickly executed by the RMA engine within the cluster, the first computing processing unit CPE in each cluster aggregates partial values. Through the reduction of the RMA engine executed on the clusters sharing the same row number, the first computing processing unit CPE in the first column of the CPE array derives the partial reduction within the row. The first computing processing element CPE in the first column performs partial reduction through the RMA engine. The first computing processing element CPE in the first cluster obtains the final reduction result and writes it back to the main memory DDR through the DMA strategy.

7. The operating method of the DeePMD-kit model according to claim 1 on the Sunway supercomputer, characterized in that, Optimize TensorFlow operators by fitting network optimization and parallelizing tensor operators. The optimization process of TensorFlow operators is as follows: 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 backpropagation. The fully connected layers of the Shenwei supercomputer include MatMul operators and Add operators. Fuse the MatMul operator and the Add operator into GEMM in Eigen, and use swBLAS in the Shenwei supercomputer to accelerate the Tanh and TanhGrad operators synchronously through SACA. Step 2: Parallelize tensor operators. Parallelize the Pad and Slice operators through SACA and optimize the bandwidth of the main memory DDR using the DMA strategy.

8. The method for running the DeePMD-kit model on the Shenwei supercomputer according to claim 1, characterized in that: Optimize the bandwidth of the Tabulate operator through the tabulation method of mixed precision. The tabulation method of mixed precision is as follows: Step 1: Create low-order weights a0, a1, a2, and a3 through bfloat16, and create high-order weights a4 and a5 through float32. Step 2: Design a mixed-precision Tabulate kernel through 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 optimization of the computing workflow of the Tabulate operator is as follows: Step 1, calculate and perform a mixed-precision vector-vector multiplication with to obtain the activation matrix a ij ; Step 2, w ij Perform a matrix-vector multiplication with o ij to obtain a descriptor vector

Images