Machine learning force field molecular dynamics efficient parallel acceleration method, system and device based on taftogen S60 chip architecture and storage medium

Abstract

The invention discloses a machine learning force field molecular dynamics efficient parallel acceleration method, system and device based on a taftogen S60 chip architecture and a storage medium. The method comprises the following steps of: firstly, acquiring initial data of a molecular system, and converting the initial data into a structure array (SoA) format to be stored in a high-bandwidth memory (HBM) of a chip; then, the simulation space is divided into a plurality of local sub-domains based on a space decomposition strategy, and atomic data are efficiently transmitted to an on-chip temporary storage memory (SRAM) through a direct memory access (DMA) engine; then, constructing a neighbor list in an SRAM (Static Random Access Memory), generating a local environment descriptor, and executing forward reasoning of a machine learning force field model by calling a tensor calculation unit of a chip to obtain atomic energy or atomic force; and finally, finishing stress reduction summation on the chip, executing numerical integration by using a vector processing unit, and updating the speed and the position of the atom. According to the method, through SoA data layout and a DMA strategy, the memory bandwidth utilization rate is improved, non-bonding force calculation is accelerated, and the expandability of parallel simulation is improved; and through a mixed precision strategy, the calculation efficiency is improved while the precision is ensured, and the method is suitable for long-time-scale high-precision simulation.

Description

Technical Field

[0001] This invention relates to a system, device, and storage medium for a highly efficient parallel acceleration method for machine learning force field molecular dynamics based on the Frostbite S60 chip architecture. Background Technology

[0002] Modern molecular dynamics (MD) simulations trace atomic trajectories by numerically integrating the Newtonian equations of motion for each atom in a many-body system. Specifically, the simulations take the following form: , in Let be the mass of the i-th atom. Let its position vector be, The net force acting on a molecular force field is an important tool for studying the evolution of the microstructure of materials, conformational changes in biomolecules, and the dynamics of chemical reactions. Traditional molecular dynamics (MD) simulations typically employ analytical forms of classical force fields, such as the Lennard–Jones potential, Coulomb interactions, and molecular force fields involving bonds, angles, and dihedral angles. The parameters of these potential functions are mostly obtained through experimental fitting or quantum chemical calculations. However, with the increasing scale of molecular systems and the growing demand for precision, the applicability of classical force fields in complex chemical environments, strongly electron-correlated systems, and bonding / breaking processes is significantly limited.

[0003] In recent years, Machine Learning Force Field (MLFF) technology based on deep learning has developed rapidly. This type of method utilizes large amounts of ab initio or high-precision quantum chemical data to train a parameterized potential surface E(θ), enabling the description of the local environment. nonlinear mapping By approximating the high-precision energy surface on the training data distribution, energy with near-abstract accuracy can be obtained. With atomic force The predictions. Typical machine learning force fields include Behler–Parrinello-type symmetric function neural networks, neural networks based on atomic environment vectors (AEV), graph neural networks (GNN), and message passing neural networks (MPNN). At each MD time step, these models need to construct high-dimensional feature vectors or graph structures based on the local atomic environment, perform multi-layer matrix multiplication and nonlinear activation to obtain atomic energy and / or atomic force; their computational mode is essentially similar to deep learning inference.

[0004] Existing implementations of machine learning force fields are primarily based on foreign GPU platforms and their software ecosystems (such as CUDA and cuDNN). These implementations typically only optimize forward inference of the neural network, while MD-specific aspects such as neighbor list construction, local environment descriptor computation, data rearrangement, and cross-subdomain communication are still handled by CPUs or general-purpose GPUs in a more general manner. In domestically produced AI acceleration chip platforms, especially those employing architectures combining high-bandwidth memory (HBM), large-capacity on-chip temporary RAM (SRAM), tensor computation units, and on-chip interconnect networks (NoC), there is currently a lack of systematic hardware and software co-design schemes for molecular dynamics simulations of machine learning force fields. This results in the chips' tensor computation capabilities and on-chip storage advantages not being fully utilized in scientific computing scenarios.

[0005] The Suiyuan S60 chip, a domestically developed high-performance AI accelerator for deep learning training and inference, integrates multi-cluster computing cores, high-bandwidth memory (HBM), large-capacity SRAM, vector processing units, tensor computing units, and programmable on-chip interconnects. Existing technologies often treat the Suiyuan S60 chip as a general-purpose matrix multiply-accumulate engine, deploying machine learning force field models on it only through deep learning frameworks. The MD program itself still primarily runs on CPU / GPU platforms, leading to frequent data transfer between the CPU, GPU, and AI chip. This makes it difficult to maintain a high hit rate for neighbor lists and local environment descriptors on-chip, hindering the realization of overall simulation efficiency and energy efficiency advantages. Furthermore, machine learning force fields typically use low-precision formats such as BF16 / FP16 for training and inference. How to combine mixed-precision strategies on domestic AI chips to improve throughput while ensuring energy conservation and structural stability in long-term MD simulations remains a problem that has not yet been fully solved in existing technologies.

[0006] Therefore, there is an urgent need for a highly efficient parallel acceleration system and method for machine learning force field molecular dynamics that deeply matches the characteristics of the storage level, tensor computing unit and on-chip interconnection network of the Suiyuan S60 chip. This would enable machine learning force field forward inference, neighbor search, local environment construction, data transfer and cross-subdomain communication to be highly coordinated within the chip, significantly improving the performance and scalability of domestic artificial intelligence acceleration chips in scientific computing scenarios while ensuring physical accuracy. Summary of the Invention

[0007] The purpose of this invention is to address the problems existing in the migration of machine learning force field molecular dynamics simulation to domestic artificial intelligence acceleration chips, such as insufficient architecture adaptation, low memory access efficiency, low coupling between machine learning model inference and MD algorithm, and high communication overhead. To address these issues, this invention proposes a highly efficient parallel acceleration system and method for machine learning force field molecular dynamics that deeply integrates the storage layer, tensor computing unit, and on-chip interconnect network characteristics of the Suiyuan S60 chip.

[0008] To address the above problems, this invention provides an efficient parallel acceleration method for machine learning force field molecular dynamics based on the FlintS60 chip architecture, characterized by comprising: Step S101: Obtain the topology, atom type, and coordinates of the molecular system to be simulated. ,speed ,quality , charge And the parameters of the machine learning force field model The initial data is collected, and the atomic data is rearranged according to spatial position and atomic type, converted into Structure of Arrays (SoA) format, and then stored in the high bandwidth memory (HBM) of the Suiyuan S60 chip; Step S102: Based on the spatial decomposition strategy, the entire simulation space is divided into multiple local spatial subdomains. And by utilizing the Direct Memory Access (DMA) engine on the Flint S60 chip, the local spatial subdomains to be computed within the current simulation time step are... The atomic data is transmitted in bursts to the on-chip temporary RAM (SRAM) of the corresponding computing core; Step S103, based on the preset truncation radius in the on-chip temporary storage memory. Build a neighbor list such that if and only if When atom j belongs to the neighbor set N(i) of atom i; a local environment descriptor vector is generated for each atom i based on the neighbor atom information. The multiple local environment descriptors are batch-concatenated to form a first feature matrix X. The tensor computing unit of the Suiyuan S60 chip is then invoked to perform forward inference of a machine learning force field model consisting of a series of generalized matrix multiplications (GEMM) and nonlinear activations to obtain the atomic energy of each atom in the molecular system. or atomic force ; Step S104: On-chip, the atomic forces and bonding interactions output by the machine learning force field model are reduced and summed to obtain the total force. Then, a numerical integration algorithm is executed using the vector processing unit to update the time step of each atom in the molecular system. The speed after With position Complete the integral evolution for the current time step.

[0009] As a further improvement of the present invention, step S103 includes: Construct a code containing the atom type of the target atom i within the current subdomain. Neighboring atoms Distance between Angle and atom type coding Local environment feature vector ; Multiple local environment feature vectors The atoms are stacked row-wise to form the first feature matrix X, where the number of rows in X is the number of atoms in the current batch. The number of columns is the feature dimension d; The first feature matrix X is sequentially input into at least one fully connected neural network, graph neural network, or message-passing neural network. Each layer of the network performs matrix operations and nonlinear activations through the tensor computation unit to form inter-layer mappings. Where W(l) and b(l) are the weight matrix and bias vector of the l-th layer, respectively. Using a nonlinear activation function, the corresponding atom embedding vector and atom energy are ultimately obtained. , and by The total energy of the system; Based on the total energy E of the system in atomic coordinates Automatic differentiation is used to obtain the forces acting on atoms: , Alternatively, a pre-trained force prediction head can be used. Based on the last layer of embedded vectors, the atomic force vectors are directly output. ; Combined with preset cutoff radius The generated mask matrix The energy or force contribution of the atom pairs (i,j) marked as invalid in the mask matrix is ​​masked, so that when At that time, the corresponding energy and force terms do not participate in the accumulation.

[0010] As a further improvement of the present invention, step S102 employs a double-buffered data prefetching mechanism, including: In the on-chip temporary storage memory of each computing core of the Frostbite S60 chip, a first buffer B1 and a second buffer B2 are divided. When the computing core uses the atomic data in the first buffer B1 and the corresponding local environment descriptor to execute the current Nth spatial subdomain During neighbor list construction, machine learning force field forward inference, and interaction calculation, the DMA engine is used to... Subdomains of space The atomic data is prefetched into the second buffer B2; After completing the Nth spatial subdomain After the calculation, the roles of the first buffer B1 and the second buffer B2 are swapped, so that the computational core can begin processing the (N+1)th spatial subdomain. Simultaneously prefetch the N+2th spatial subfield The data allows for the overlap of machine learning force field inference and memory access.

[0011] As a further improvement of the present invention, it also includes a boundary atom exchange step based on the on-chip interconnect network NoC, including: After completing the integral update in step S104, for each spatial subdomain The migration atoms that cross the subdomain boundary are detected to obtain the set of migration atoms. ; For sets Each migrating atom in the process is constructed by the current computing core, containing the target computing core identifier and atom coordinates. ,speed ,type , charge and the current forces The data packets are sent directly to the adjacent computing cores via the NoC on-chip interconnect network of the S60 chip. The adjacent computing cores receive and merge the migrated atomic data in their on-chip temporary storage buffers without passing through high-bandwidth memory, thereby completing the atomic data synchronization across subdomain boundaries and maintaining the local consistency required for neighbor list construction and machine learning force field calculations.

[0012] As a further improvement of the present invention, steps S103 and S104 employ a mixed-precision calculation strategy, including: Local environment descriptor construction, machine learning force field forward inference and intermediate embedding vectors, atomic energy During storage and computation, high-throughput computation is performed using the BF16 or FP16 computing units of the Frostbite S60 chip. In the calculation of the cumulative atomic force F, the total system energy E, the statistical properties of the virial tensor W, and the atomic positions... and speed During the integral update process, FP32 or FP64 precision simulated by software is used for computation to suppress energy drift and numerical instability in long-term simulations.

[0013] Another subject of this invention is a parallel acceleration system for machine learning force field molecular dynamics based on the Flint S60 chip architecture, comprising: At least one Flint S60 chip and its connected high-bandwidth memory HBM; A non-volatile storage medium storing program instructions executable on the FlintS60 chip, wherein, when loaded and executed, the program instructions configure the FlintS60 chip to perform the steps of the method according to any one of claims 1 to 5; wherein, The program instructions include at least: The data rearrangement and mapping module is used to rearrange the atomic coordinates in a molecular system. ,speed ,quality , charge Atom type The sparse atomic state data is converted into SoA format adapted to the DMA access of the Flint S60 chip, and the converted data is written into the high-bandwidth memory HBM. The local environment and tensor mapping module is used to construct the truncation radius in the on-chip temporary storage memory (SRAM). Neighbor list and local environment descriptors The descriptor tensor is mapped to the input feature matrix X of the tensor computation unit, and then the forward inference of the machine learning force field model is realized through multi-layer GEMM operation. The SRAM management module is used to explicitly manage the space allocation and data swapping of the on-chip temporary storage of the S60 chip, maintain the locality of atomic data, neighbor lists and local environment descriptors, and control the switching of each buffer under the double buffer prefetch mechanism. It also includes a communication management module based on an on-chip interconnect network, used to control the direct routing and merging of migrated atomic data across spatial subdomain boundaries between different computing cores, enabling on-chip data synchronization without the need for HBM relay.

[0014] Another subject of the invention is a computer-readable storage medium having program instructions stored thereon that, when executed by a processor, cause the processor to perform the method according to any one of claims 1 to 5.

[0015] Another subject of the invention is an electronic device comprising at least one processor and a memory, the memory storing a program executable on the processor, the program, when executed, causing the processor to perform the method according to any one of claims 1 to 5.

[0016] Compared with the prior art, the present invention has the following beneficial effects: (1) By introducing SoA data layout and DMA-based HBM to SRAM data migration strategy, the residence time and access granularity of atomic data in the storage level are explicitly controlled, avoiding the jitter problem in the traditional caching mechanism and improving the effective memory bandwidth utilization. (2) By tensorizing the local atomic environment and mapping it to multi-layer GEMM-type matrix operations, this invention can fully utilize the large-scale multiply-accumulate computing power of the tensor computing unit of the Suiyuan S60 chip in deep neural network inference, and significantly accelerate the non-bonded force calculation kernel of the machine learning force field. (3) By employing double buffered data prefetching in the on-chip temporary storage and performing migration atom-free switching on the on-chip interconnect network, this invention significantly reduces global memory access and cross-chip communication overhead while achieving a high degree of overlap between machine learning force field inference and memory access and communication, thereby improving the scalability of multi-chip parallel simulation. (4) By designing a hybrid precision strategy for BF16 / FP16 and FP32 / FP64, the invention effectively controls numerical errors such as energy drift while making full use of low precision and high throughput computing power, making it suitable for long-scale machine learning force field molecular dynamics simulation with high precision requirements. Attached Figure Description

[0017] Figure 1 This is a flowchart of an efficient parallel acceleration method for machine learning force field molecular dynamics based on the Frostbite S60 chip architecture according to some embodiments of the present invention. Figure 2 This is a block diagram of a parallel acceleration system for machine learning force field molecular dynamics based on the Flint S60 chip according to some embodiments of the present invention. Figure 3 This is a schematic diagram of the spatial division of a molecular system and the construction of a local environment according to some embodiments of the present invention. Detailed Implementation

[0018] To make the objectives, technical solutions, and beneficial effects of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be understood that the following embodiments are for illustrative purposes only and are not intended to limit the scope of protection of this invention.

[0019] Example 1: Overall Process refer to Figure 1 In one embodiment of the present invention, the efficient parallel acceleration method for machine learning force field molecular dynamics based on the Frostbloon S60 chip architecture includes the following steps.

[0020] First, the data acquisition module reads the topological structure and atom types of the molecular system to be simulated. Initial coordinates ,speed ,quality , charge Temperature and pressure control parameters and parameters of the pre-trained machine learning force field model The machine learning force field model can be an atom-centric neural network based on symmetry functions, a neural network based on atom-environment vectors, a graph neural network, or a message-passing neural network, etc. To avoid frequent data format conversions during the simulation, this embodiment rearranges physical quantities such as atomic coordinates, velocities, forces, atom type codes, and charges into SoA format during the initialization phase. That is, all atomic values ​​of the same physical quantity are stored continuously in memory, such as the X-coordinate array, Y-coordinate array, and Z-coordinate array being arranged continuously, thereby adapting to the burst transmission characteristics of the Suiyuan S60 chip's DMA engine, and writing the above data into HBM.

[0021] Then, the task partitioning module divides the entire simulation space into multiple three-dimensional mesh subdomains based on a spatial decomposition strategy. The computational tasks of each subdomain are mapped to different computing cores on the Flint S60 chip. For each time step... The DMA engine loads data such as corresponding atom coordinates, atom type, charge, and force generated in the previous step from HBM to on-chip temporary storage memory (SRAM) in burst mode, based on subdomains. To improve SRAM space utilization, this embodiment divides the SRAM logic of each computing core into multiple regions, including its own atom region, neighbor cache region, local environment descriptor region, force accumulation region, and double-buffered prefetch region, and explicitly manages the lifecycle of each region at the software layer.

[0022] Next, the neighbor list and local environment generation module constructs the cutoff radius for each subdomain in SRAM. Neighbor list Specifically, this embodiment employs a Cell-Linked List-based spatial bucketing method to divide the simulated space into areas equal to or less than the cutoff radius. Each grid cell, of the same order of magnitude, maintains a linked list pointing to the indices of its internal atoms. By traversing the current subdomain and several of its neighboring grid cells, a set of neighbor indices for each atom within the current subdomain is constructed. Based on this, the local environment generation module generates each atom according to the input format of the machine learning force field model. Constructing a local environment descriptor Examples of descriptors include radial distribution functions, angular distribution functions, multibody symmetric functions, adjacency matrix features based on embedding layers, or other high-dimensional feature vectors. These descriptors are then packaged into tensors and stored in the local environment descriptor region.

[0023] Subsequently, the machine learning force field tensor calculation module calls the tensor calculation unit of the Suiyuan S60 chip to perform forward inference on the aforementioned local environment descriptor tensor. Taking a fully connected neural network as an example, in this embodiment, the weight matrix of each layer is denoted as W(l), the bias vector as b(l), and the feature tensor as X(l). During the forward inference process, the local environment feature matrix X(l) is... Stacked, executed layer by layer ,in It is a nonlinear activation function, until the atomic energy is output. or atomic force For graph neural networks or message-passing neural networks, this embodiment adopts a similar approach, encoding adjacency relationships and node features into dense or sparse tensor representations adapted to GEMM operations, and obtaining node embedding vectors and energy predictions through multiple rounds of message passing and aggregation operations.

[0024] When the machine learning force field model outputs atomic energy In this embodiment, atomic forces are calculated using automatic differentiation or analytical derivation; when the model directly outputs atomic forces... In this case, the differentiation steps described above can be omitted. To ensure the cutoff radius... Non-atomic pairs do not produce spurious interactions. In this embodiment, a Boolean mask matrix is ​​generated based on a neighbor list. The energy or force contribution of atomic pairs not within the cutoff radius is uniformly set to zero, thus balancing computational efficiency and physical rationality.

[0025] Next, the update module uses the vector processing unit to perform the calculation of bonding interactions in SRAM, including bond stretching, bond angle bending, dihedral torsion, and constraint potentials, and then uses the machine learning force field to calculate the nonbonded forces. The total force F on the atom is obtained by vector summation with the bonding force. Then, the atom's velocity and position are updated using Velocity–Verlet or other time-reversible integration algorithms, for example: , , .

[0026] In this embodiment, FP32 precision is used in the integral update stage, and numerical compensation algorithms such as Kahan summation can be selected to further reduce rounding errors.

[0027] Finally, the communication management module performs boundary detection of atoms within the subdomain locally, marking migrated atoms that cross the subdomain boundary. For each migrated atom, the current computing core constructs a data packet containing the target subdomain identifier, atom coordinates, velocity, type, and related state variables, and routes it directly to the target computing core's receive buffer via the on-chip interconnect network (NoC). After receiving the data packet, the target computing core inserts the migrated atom into its own atom list and incorporates it into the local environment construction and machine learning force field inference in the next neighbor list reconstruction cycle. Since the above process is completed entirely within the chip, without needing to write data back to the HBM and forward it via an external bus, cross-subdomain communication latency and power consumption are significantly reduced.

[0028] Example 2: Double-buffered prefetching and mixed precision strategy In another embodiment, the present invention addresses the computationally intensive and memory access-regular characteristics of machine learning force field inference by employing a double-buffered prefetch mechanism in SRAM and combining it with a mixed-precision computing strategy to further improve overall performance and energy efficiency.

[0029] Specifically, for the SRAM of each computing core, this embodiment divides the data region used to store atomic coordinates, atomic types, and local environment descriptors into two buffers of equal size, B1 and B2. When the computing core uses the data in B1 to perform neighbor list construction and machine learning force field forward inference for the current Nth subdomain Ω(N), the DMA engine prefetches the data of the (N+1)th subdomain Ω(N+1) in parallel into B2. After the Nth subdomain computation is completed, the roles of the two buffers are swapped; B2 becomes the working area for the new round of computation, while B1 is used to prefetch the data of the (N+2)th subdomain. In this way, this embodiment achieves pipelining overlap of machine learning force field inference and HBM memory access within a single computing core, significantly reducing the impact of memory access latency on overall performance.

[0030] Regarding the mixed-precision strategy, this embodiment employs BF16 or FP16 data formats in the local environment descriptor construction, embedded vector computation, and machine learning force field forward inference stages to fully utilize the high throughput of the tensor computation unit of the Suiyuan S60 chip in low-precision multiply-accumulate operations. In the atomic force accumulation, total system energy E and virial tensor W calculations, and integral update stages, FP32 or FP64 precision simulated via double-single methods is used for computation and storage to suppress energy drift and structural shifts in long-scale simulations. Experiments show that in typical biomacromolecules and solid material systems, this mixed-precision strategy can significantly reduce computational and storage overhead without significantly affecting the statistical accuracy of physical quantities such as temperature, pressure, and radial distribution functions.

[0031] Example 3: System Structure and Module Division refer to Figure 2 In yet another embodiment, the parallel acceleration system for machine learning force field molecular dynamics based on the Flint S60 chip provided by the present invention includes: At least one S60 chip and a high-bandwidth memory (HBM) connected thereto, wherein the S60 chip includes a multi-cluster computing core, a tensor computing unit, a vector processing unit, an on-chip temporary storage memory (SRAM), a direct memory access (DMA) engine, and an on-chip interconnect network (NoC). A non-volatile storage medium stores program instructions that can run on the Suiyuan S60 chip. When these program instructions are loaded and executed, the Suiyuan S60 chip is configured to perform the steps of the method of the present invention, and logically divided into the following functional modules: The data acquisition and rearrangement module is used to read the topology, initial atomic state, and machine learning force field model parameters of the molecular system from external storage or host system, and to complete the SoA data layout transformation and write the processed data into HBM. The spatial subdomain partitioning and task scheduling module is used to divide the global simulation space into multiple subdomains according to the spatial size and number of particles of the simulation system, and map each subdomain to different computing cores or different chips for execution, while managing task distribution and load balancing in the time step loop. The neighbor list and local environment generation module is used to construct a neighbor list N(i) based on the cutoff radius in SRAM and generate the corresponding local environment descriptor tensor according to the input requirements of the machine learning force field model. The Machine Learning Force Field Tensor Calculation Module is used to call the tensor calculation unit to perform multi-layer neural network forward inference on the local environment descriptor tensor, output atomic energy and atomic force, and perform automatic differentiation or force prediction head calculation when needed. The update and statistics module is used to perform bond interaction calculations, force accumulation, numerical integration updates, and calculations of statistical quantities such as system energy, temperature, and pressure using the vector processing unit; The communication management module is used to realize the seamless switching of atoms across subdomain boundaries using the on-chip interconnect network (NoC), and to collaborate with external communication libraries in multi-chip or multi-card cluster scenarios to achieve large-scale parallel simulation.

[0032] Example 4: Management and Online Update of Machine Learning Force Field Model In another embodiment, the present invention further incorporates the model update requirements of machine learning force fields by introducing a model management and online update module into the system. This module can be deployed on the node where the Flint S60 chip is located or on an external training node, and its main functions include: Monitor model uncertainty indicators in current machine learning force field molecular dynamics simulations, such as prediction discrepancies, force error estimation, or energy conservation deviations based on ensemble models. When the uncertainty or error exceeds a preset threshold, a representative configuration is sampled from the simulation process, and high-precision quantum chemical calculations are triggered to obtain reference energy and force data. Based on the newly sampled data, the parameters θ of the machine learning force field model are retrained or fine-tuned externally or locally, and the updated model parameters are written back to the HBM or on-chip memory of the Suiyuan S60 chip through the parameter synchronization mechanism. Without disrupting the continuity of MD time progression, the machine learning force field model version is switched at the boundaries of several time steps to achieve adaptive model iteration based on live learning.

[0033] Those skilled in the art will understand that, without departing from the spirit and essence of the present invention, various modifications and substitutions can be made to the parameter settings, data structure organization methods, and specific hardware resource allocation in the above embodiments. For example, adjusting the subdomain partitioning granularity for different scale systems, adopting other forms of local environment descriptors and neural network structures, superimposing cross-node MPI communication mechanisms in multi-card environments, or offloading some numerical computation tasks to other types of accelerators, etc. All these equivalent substitutions should fall within the protection scope of the present invention.

[0034] The technical principles of the present invention have been described above with reference to specific embodiments. These descriptions are merely for explaining the principles of the invention and should not be construed as limiting the scope of protection of the invention in any way. Based on this explanation, those skilled in the art can readily conceive of other specific embodiments of the invention without inventive effort, and these embodiments will all fall within the scope of protection of the present invention.

Claims

1. A machine learning force field molecular dynamics high efficient parallel acceleration method based on the S60 chip architecture, characterized in that, Comprising: Step S101: Obtain the topology, atom type, and coordinates of the molecular system to be simulated. ,speed ,quality , charge And the parameters of the machine learning force field model The initial data is collected, and the atomic data is rearranged according to spatial position and atomic type, converted into Structure of Arrays (SoA) format, and then stored in the high bandwidth memory (HBM) of the Suiyuan S60 chip; Step S102: Based on the spatial decomposition strategy, the entire simulation space is divided into multiple local spatial subdomains. And by utilizing the Direct Memory Access (DMA) engine on the Flint S60 chip, the local spatial subdomain to be computed within the current simulation time step is... The atomic data is transmitted in bursts to the on-chip temporary RAM (SRAM) of the corresponding computing core; Step S103, based on a preset cutoff radius, in the on-chip temporary storage Constructing a neighbor list , so that when and only when , atom j belongs to the neighbor set N(i) of atom i; generating a local environment descriptor vector for each atom i according to the neighbor atom information , splicing a plurality of said local environment descriptors into batches to form a first feature matrix X, calling the tensor computing unit of the projectile original S60 chip to perform forward inference of a machine learning force field model composed of a series of general matrix multiplication (GEMM) and nonlinear activation, to obtain the atomic energy or atomic force of each atom in the molecular system ;​ Step S104, the atomic force and the bond interaction force output by the machine learning force field model are reduced and summed on the chip to obtain a total force, and a numerical integration algorithm is executed by using a vector processing unit to update the velocity of each atom in the molecular system after a time step and position , completing the integration evolution of the current time step.​ 2. The method of claim 1, wherein the method is based on the S60 chip architecture of Blaze. The step S103 comprises: Construct a code containing the atom type of the target atom i within the current subdomain. Neighboring atoms Distance between Angle and atom type coding Local environment feature vector ; a plurality of said local environment feature vectors The first feature matrix X is formed by stacking the rows, the number of rows of X is the number of atoms in the current batch , and the number of columns is the feature dimension d. The first feature matrix X is sequentially input into at least one layer of a fully connected neural network, a graph neural network or a message passing neural network, and each layer of the network performs matrix operation and nonlinear activation through the tensor calculation unit to form interlayer mapping: where W(l), b(l) are the weight matrix and bias vector of the lth layer, respectively, is a nonlinear activation function, and finally the corresponding atomic embedding vector and atomic energy are obtained and the total energy of the system is composed of Based on the total energy E of the system on atomic coordinates Automatic differentiation is performed to obtain atomic forces: , Or use a pre-trained force prediction head On the basis of the last layer embedding vector, the atomic force vector is directly output ; Combining preset cut-off radii Generated mask matrix The energy or force contribution of an atom pair (i,j) marked as invalid in the mask matrix is shielded such that when the corresponding energy term and force term do not participate in the accumulation.

3. The method of claim 1, wherein the method is based on a blaze S60 chip architecture. The step S102 employs a double-buffering data prefetching mechanism, comprising: Dividing a first buffer B1 and a second buffer B2 in the on-chip scratchpad memory of each computing core of the Spark S60 chip; When the computing core executes neighbor list construction, machine learning force field forward inference and interaction calculation for the current Nth spatial subdomain with atomic data in the first buffer B1 and the corresponding local environment descriptor , the DMA engine is utilized to prefetch atomic data of the Nth spatial subdomain to the second buffer B2.​ Upon completion of the computation of the Nth spatial subdomain , the roles of the first buffer B1 and the second buffer B2 are exchanged, so that the computing core prefetches data of the N+2th spatial subdomain while starting to process the N+1th spatial subdomain , thereby achieving overlap of machine learning force field inference and memory access.

4. The method of claim 1, wherein the method is based on the S60 chip architecture of Blaze Origin, and Further comprising a boundary atomic exchange step based on a network-on-chip NoC, comprising: After the integral update of the step S104 is completed, for each spatial subdomain The migrating atoms crossing the subdomain boundary are detected to obtain a migrating atom set ; For each migration atom in the set , a data packet containing the target computing core identity, atom coordinates , velocity , type , charge , and current force is built by the current computing core and sent directly to the neighboring computing core through the on-chip interconnection network NoC of the projectile S60 chip. Receiving and merging the migrated atomic data in the receiving buffer of the on-chip scratchpad memory of the adjacent computing core without going through the high-bandwidth memory mediation to complete the atomic data synchronization across the sub-domain boundary, maintaining the local consistency required for neighbor list construction and machine learning force field calculation.

5. The machine learning force field molecular dynamics high efficient parallel acceleration method based on the S60 chip architecture according to claim 1, characterized in that, The step S103 and the step S104 employ a mixed-precision calculation strategy, comprising: In the local environment descriptor construction, machine learning force field forward inference, and intermediate embedding vector, atomic energy storage and operation process, the BF16 or FP16 calculation unit of the projectile original S60 chip is used for high throughput operation; In the integration of atomic forces F, the calculation of the statistical quantities of the total energy E and the Virial tensor W, and the integration update of atomic positions and velocities , FP32 or FP64 precision simulated by software is used to suppress energy drift and numerical instability in long-time scale simulations.

6. A machine learning force field molecular dynamics parallel acceleration system based on the S60 chip architecture of Blaze, characterized in that, Comprising: At least one Spark S60 chip and a high-bandwidth memory HBM connected thereto; A non-volatile storage medium storing program instructions executable on the Spark S60 chip, the program instructions, when loaded and executed, configure the Spark S60 chip to perform the steps of the method of any one of claims 1 to 5; wherein, The program instructions at least comprise: a data rearrangement and mapping module configured to convert sparse atom state data of coordinates , velocities , masses , charges , atom types , etc. of atoms in a molecular system into a SoA format suitable for DMA access by the ray tracing S60 chip, and to write the converted data to the high bandwidth memory HBM; A local environment and tensor mapping module is configured to build a truncated radius in an on-chip scratchpad memory (SRAM) a neighbor list within the truncated radius and a local environment descriptor The descriptor tensor is mapped as an input feature matrix X of a tensor computing unit, and a forward inference of a machine learning force field model is realized through a multi-layer GEMM operation. An SRAM management module for explicitly managing the space allocation and data swapping of the on-chip scratchpad memory of the Spark S60 chip, maintaining the locality of atomic data, neighbor lists, and local environment descriptors, and controlling the switching of each buffer under the double-buffering prefetching mechanism; And a communication management module based on a network-on-chip for controlling the direct routing and merging of migrated atomic data across the spatial sub-domain boundary between different computing cores, realizing on-chip data synchronization without going through the HBM mediation.

7. A computer-readable storage medium having program instructions stored thereon, the program instructions, when executed by a processor, causing the processor to perform the method of any one of claims 1 to 5.

8. An electronic device comprising at least one processor and a memory, the memory storing a program executable on the processor, the program, when executed, causing the processor to perform the method of any one of claims 1 to 5.

Images