diff --git a/CMakeLists.txt b/CMakeLists.txt index bb93ecfce..0fdb4d6d3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -30,7 +30,7 @@ file(GLOB GPUMD_SOURCES src/utilities/*.cu) add_executable(gpumd ${GPUMD_SOURCES}) target_include_directories(gpumd PRIVATE src) -target_link_libraries(gpumd PRIVATE CUDA::cublas CUDA::cusolver CUDA::cufft) +target_link_libraries(gpumd PRIVATE CUDA::cublas CUDA::cusolver CUDA::cufft Threads::Threads) # --- nep --- file(GLOB NEP_SOURCES diff --git a/developer/README.md b/developer/README.md index 8fb36ab14..eca2e5b09 100644 --- a/developer/README.md +++ b/developer/README.md @@ -55,8 +55,8 @@ The main source directories are: Files added under an existing source directory are normally collected automatically by the makefiles. The CUDA and HIP makefiles also collect `.cu` files one directory below `src/force/` for `gpumd`, and `src/main_nep/` for `nep`. Training in `src/main_nep/nep_spin/` uses the shared magnetic descriptor and -force headers in `src/force/nep_spin/`. GPUMD Spin force dispatch is added -separately; the headers do not enable a runtime potential by themselves. +force headers in `src/force/nep_spin/`. The GPUMD potential, fixed-spin runtime support and spatial decomposition +also live in `src/force/nep_spin/`. See [runtime validation](spin_runtime_validation.md). ## 3. Main `gpumd` architecture diff --git a/developer/spin_runtime_validation.md b/developer/spin_runtime_validation.md new file mode 100644 index 000000000..8fc8db8bc --- /dev/null +++ b/developer/spin_runtime_validation.md @@ -0,0 +1,57 @@ +# Fixed-spin runtime validation + +Build `gpumd` and `nep` on a CUDA node, choosing the device architecture: + +```sh +make -C src -j8 CUDA_ARCH=-arch=sm_70 gpumd nep +export GPUMD_COMMAND="$PWD/src/gpumd" +export NEP_COMMAND="$PWD/src/nep" +``` + +Use Python with NumPy and pytest. Set `RESULTS` to a new output directory outside +this checkout. The first two GPU scripts need one GPU; run them with one +allocated device visible. The third requires four allocated GPUs in one process. + +```sh +python -m pytest --noconftest -q tests_pytest/test_spin_model_format.py tests_pytest/test_spin_validation.py +python tests_pytest/validate_nep_spin3_runtime.py +python tests_pytest/validate_mfield_runtime.py +python tests_pytest/validate_spin_multigpu.py --binary "$GPUMD_COMMAND" --model tests_pytest/fixtures/nep_spin3/nep_spin_o3c2_uniform.nep --output "$RESULTS/runtime" +``` + +The runtime oracle checks cover energy, atomic force, magnetic force and all +nine virial components, plus energy-reference invariance, excluded-spin +masking, replication and fixed-spin restart. Model-selection checks include +uppercase fields, attached comments and a model record after numeric blocks. +Small-cell checks cover capacity growth and a 10,000-atom thin cell compared +with an equivalent expanded cell. Split XYZ output is checked after restart. The frozen reference provenance +is documented in [fixtures](../tests_pytest/fixtures/nep_spin3/README.md). +The external-field script checks analytic Zeeman energy and magnetic force, +group and region selection, boundary crossing, excluded spins and invalid inputs. +It also checks additive fields and clearing/redeclaring them across runs. +Ordinary-model checks require extra scalar, vector, nonfinite, string and duplicate +`spin` properties to leave predictions and restart output unchanged. +Multi-GPU comparisons use componentwise tolerances and reject nonfinite values. +They cover small cells, nonuniform and triclinic systems, cutoff boundaries, +partition directions, motion across partition boundaries, additive magnetic fields +and energy references. +An additional supported `n_max 4 8`, `l_max 8 2 1` model needs 51,840 bytes +in the direct angular-force path and therefore exercises its 48 KiB fallback +on one, two and four GPUs with nonzero forces. + +Run the [training validation](spin_validation.md) against the same build to +check the shared descriptor implementation and native/JIT agreement. Keep its +model, seed, budget and GPU settings fixed when comparing revisions. + +For ordinary GPUMD and NEP behavior, build an unmodified upstream checkout and +run the full regression suite. `BASELINE_GPUMD` and `BASELINE_NEP` must identify +those independently built upstream executables: + +```sh +python tests_regression/run_regression.py --repo-root "$PWD" \ + --baseline "$BASELINE_GPUMD" --candidate "$GPUMD_COMMAND" \ + --baseline-nep "$BASELINE_NEP" --candidate-nep "$NEP_COMMAND" --suite full +``` + +These checks validate implementation contracts, not fit quality or performance. +Spin integrators and their checkpoint state are outside this suite. diff --git a/developer/spin_validation.md b/developer/spin_validation.md index 86ebf52f9..e815a6c0a 100644 --- a/developer/spin_validation.md +++ b/developer/spin_validation.md @@ -90,3 +90,5 @@ implementation to make a failing check pass. Models use `model nep_spin`, Keep build logs, source/binary hashes, timing and numerical reports outside the source tree. + +For GPUMD fixed-spin inference and restart checks, see [Spin runtime validation](spin_runtime_validation.md). diff --git a/doc/gpumd/input_files/model_xyz.rst b/doc/gpumd/input_files/model_xyz.rst index b918e4b1a..140c0fc62 100644 --- a/doc/gpumd/input_files/model_xyz.rst +++ b/doc/gpumd/input_files/model_xyz.rst @@ -51,6 +51,7 @@ Essentially any keyword is allowed, but we only read the following ones: * ``mass:R:1`` mass (*Optional*: default mass values will be used when this is missing) * ``charge:R:1`` charge (*Optional*: default charge is zero) * ``vel:R:3`` velocity vector (*Optional*) + * ``spin:R:3`` magnetic-moment vector in Bohr magnetons (muB) (*Required for NEP_Spin*) * ``group:I:number_of_grouping_methods`` grouping methods (*Optional*) @@ -67,6 +68,16 @@ The mass should be given in units of the unified atomic mass unit (amu). The charge is in units of e (proton charge). The cell dimensions and atom coordinates should be given in units of Ångstrom. Velocities need to be specified in units of Å/fs. +Spin is specified in Bohr magnetons (muB). +Magnetic force is the negative energy derivative with respect to spin, in eV/muB. +The spin torque, spin crossed with magnetic force, is in eV. +External magnetic induction B is specified in tesla (T). +Spin vectors store Cartesian magnetic moments, not normalized directions. +Active spins are not normalized; training and runtime data must use the same +spin convention. Elements excluded by the model's ``spin_type`` have zero +effective spin and zero magnetic force, while their geometry is retained. +The ``spin`` property may not be repeated. NEP_Spin fails if ``spin:R:3`` is absent; +missing spin values are never filled with zeros. Example ------- diff --git a/doc/gpumd/input_files/run_in.rst b/doc/gpumd/input_files/run_in.rst index e9c414719..6cf082939 100644 --- a/doc/gpumd/input_files/run_in.rst +++ b/doc/gpumd/input_files/run_in.rst @@ -101,6 +101,10 @@ Simulation setup - Apply an electric field to selected atoms - No - No + * - :ref:`add_mfield ` + - Apply a magnetic field to an atom group or fractional region + - No + - Yes * - :ref:`add_spring ` - Add spring interactions to selected atom groups - No diff --git a/doc/gpumd/input_parameters/add_mfield.rst b/doc/gpumd/input_parameters/add_mfield.rst new file mode 100644 index 000000000..a289d76ec --- /dev/null +++ b/doc/gpumd/input_parameters/add_mfield.rst @@ -0,0 +1,55 @@ +.. _kw_add_mfield: + +:attr:`add_mfield` +================== + +Apply a constant external magnetic induction to an atom group or a fractional-coordinate region:: + + add_mfield + add_mfield region + +Grouping method and group ID select an existing model.xyz group, as in +add_efield. The three finite field components are in tesla (T). There is no +separate magnitude or direction normalization. For a 10 T field along z:: + + add_mfield 0 0 0 0 10 + +Group selection follows atom identities even when atoms move. Region selection +instead follows current fractional coordinates along cell vectors a, b, and c, +using the same coordinate conversion and membership test as heat_lan. +Bounds must be finite, in [0, 1], and satisfy min < max. Intervals are half-open: +[min, max). Coordinates are evaluated after periodic boundary wrapping. +No group labels are needed for region mode. + +For a 10 T field along z in the first half of the cell along a:: + + add_mfield region 0 0.5 0 1 0 1 0 0 10 + +Membership is recomputed on every force evaluation: atoms entering the region +receive the field, and atoms leaving it no longer do. The sharp boundary gives +a discontinuous Zeeman energy. No mechanical boundary impulse is implemented; +energy conservation is not guaranteed for atoms crossing this boundary. This +is a local magnetic drive, not a smooth conservative gradient-field model. + +A Spin potential and spin:R:3 are required. The field acts +only on selected atoms whose elements belong to the model's ``spin_type``. +Other elements retain zero effective spin and zero magnetic force. + +Declare the field before each ``run`` that uses it. Multiple ``add_mfield`` +commands in the same run add their contributions, including overlapping groups +or regions. Field actions are cleared after the run. To run without a field, +omit ``add_mfield`` before the next ``run``; a zero-valued command contributes +zero and does not cancel other field actions. + +Spin is in muB, magnetic force is in eV/muB, and energy is in eV. With +c_B = 5.7883818060e-5, the added magnetic force is c_B times the field vector; +the added per-atom Zeeman energy is minus spin dotted with this force. +Both contributions are recomputed at every force evaluation using the actual +stored moments and do not accumulate across evaluations. +Output energies and magnetic forces include these contributions. +No additional g factor is applied. There is no direct mechanical force or +virial contribution. Time-dependent fields and linear gradients are not supported. + +The field is magnetic induction B, not H in A/m. +For example, a 2 muB moment aligned with 1 T has Zeeman energy approximately +-1.1576e-4 eV and magnetic force +5.7884e-5 eV/muB along the field. diff --git a/doc/gpumd/input_parameters/dump_restart.rst b/doc/gpumd/input_parameters/dump_restart.rst index 80242bc36..7896c2dfb 100644 --- a/doc/gpumd/input_parameters/dump_restart.rst +++ b/doc/gpumd/input_parameters/dump_restart.rst @@ -28,3 +28,5 @@ Caveats ------- This keyword is not propagating. That means, its effect will not be passed from one run to the next. + +NEP-spin models also save the per-atom Cartesian ``spin:R:3`` values. diff --git a/doc/gpumd/input_parameters/dump_xyz.rst b/doc/gpumd/input_parameters/dump_xyz.rst index a4075c88b..f30e4284b 100644 --- a/doc/gpumd/input_parameters/dump_xyz.rst +++ b/doc/gpumd/input_parameters/dump_xyz.rst @@ -34,7 +34,7 @@ If it is omitted, data for the whole system will be output. :attr:`single` (the default) writes nine significant digits, which is enough to recover a single-precision value exactly. :attr:`double` writes 17 significant digits, which is enough to recover a double-precision value exactly. -* Then one can write the properties to be output, and the allowed properties include: :attr:`mass`, :attr:`velocity`, :attr:`force`, :attr:`potential`, :attr:`virial`, :attr:`charge`, :attr:`bec`, :attr:`group_labels`, and :attr:`unwrapped_position`. +* Then one can write the properties to be output, and the allowed properties include: :attr:`mass`, :attr:`velocity`, :attr:`force`, :attr:`potential`, :attr:`virial`, :attr:`charge`, :attr:`bec`, :attr:`group_labels`, :attr:`spin`, :attr:`mforce`, and :attr:`unwrapped_position`. The :attr:`group_labels` property writes one integer column per grouping method, holding the label of the group each atom belongs to. @@ -72,3 +72,8 @@ Caveats * For qNEP models, the charge values dumped out are predicted by the qNEP models. For other models, the charge values are those specified in :attr:`model.xyz` via :attr:`charge:R:1`. * The Born effective charge (:term:`BEC`) is only meaningful for a qNEP model trained with target :term:`BEC`. +* ``spin`` and ``mforce`` are three-component vectors available when + ``model.xyz`` contains ``spin:R:3``. +* The per-atom ``virial`` property contains all nine row-major components + ``xx xy xz yx yy yz zx zy zz``. NEP_Spin does not symmetrize this per-atom + tensor. diff --git a/doc/gpumd/input_parameters/index.rst b/doc/gpumd/input_parameters/index.rst index f9b68e302..88fe99e78 100644 --- a/doc/gpumd/input_parameters/index.rst +++ b/doc/gpumd/input_parameters/index.rst @@ -38,6 +38,7 @@ Below you can find a listing of keywords for the ``run.in`` input file. ensemble_nphug add_force add_efield + add_mfield add_spring deposit electron_stop diff --git a/doc/gpumd/input_parameters/potential.rst b/doc/gpumd/input_parameters/potential.rst index 96324fd62..1b221dcf6 100644 --- a/doc/gpumd/input_parameters/potential.rst +++ b/doc/gpumd/input_parameters/potential.rst @@ -30,7 +30,7 @@ relative or absolute path) of the potential file to be used. Some potential types require additional parameters, as described in their corresponding sections. -For standard NEP and temperature-dependent NEP models, a multi-GPU run can +For standard NEP, temperature-dependent NEP, and ``nep_spin`` models, a multi-GPU run can optionally specify the spatial partition direction:: potential [partition_direction] @@ -41,7 +41,7 @@ triclinic box. If the direction is omitted, GPUMD selects it automatically, normally along the thickest direction of the box. The optional partition direction applies only to the multi-GPU implementation -of standard NEP and temperature-dependent NEP models. Charge NEP models +of standard NEP, temperature-dependent NEP, and ``nep_spin`` models. Charge NEP models currently use the single-GPU implementation in ``gpumd``. Example @@ -55,3 +55,17 @@ For a multi-GPU NEP run, one can force the decomposition along a specific direction, for example:: potential Si_NEP.txt x + +NEP-spin +-------- + +A :ref:`NEP-spin model ` uses all GPUs visible to the process to +compute the same atomic system. The optional partition direction also applies +to NEP-spin. Small systems may not benefit from multiple GPUs; global atom +arrays remain on GPU 0. + +Provide ``spin:R:3`` in :ref:`model.xyz `. Fully periodic orthogonal +and triclinic cells are supported. Spins remain fixed under ``nve`` and +``nvt_nhc``; other ensembles and multiple-potential runs are not supported. +See :ref:`the Spin runtime instructions ` for output, +restart and external-field usage. diff --git a/doc/installation.rst b/doc/installation.rst index ead19f708..56cb691e1 100644 --- a/doc/installation.rst +++ b/doc/installation.rst @@ -97,7 +97,7 @@ use ``CUDA_VISIBLE_DEVICES`` to select a subset, for example:: CUDA_VISIBLE_DEVICES=0,1 ./nep The ``gpumd`` executable also uses all visible GPUs when running standard NEP -or temperature-dependent NEP models that support the multi-GPU implementation. +or temperature-dependent NEP models, and ``nep_spin`` spin models. For example:: CUDA_VISIBLE_DEVICES=0,1 ./gpumd diff --git a/doc/potentials/nep_spin.rst b/doc/potentials/nep_spin.rst index 1f05838f3..c9b3b3638 100644 --- a/doc/potentials/nep_spin.rst +++ b/doc/potentials/nep_spin.rst @@ -609,3 +609,45 @@ selected Spin layout. The magnetic kernels share their formulas with the native path. If runtime compilation is unavailable, native kernels remain available. Compilation adds startup time, so very short runs may be slower; ``nep_compile 0`` disables specialization. Prediction uses the native path. + +.. _nep_spin_runtime: + +Fixed-spin molecular dynamics +----------------------------- + +Load the exported model with ``potential nep.txt`` and provide Cartesian +magnetic moments through ``spin:R:3`` in :ref:`model.xyz `. +The ``nve`` and ``nvt_nhc`` ensembles evolve positions and velocities while +keeping the magnetic moments fixed. Only one potential is allowed, and all +three cell directions must be periodic. + +Atoms whose elements are absent from ``spin_type`` have zero effective spin. +Their input spin values are ignored; they still contribute geometric +environments. Their magnetic forces remain zero, including in an applied +external field. The reader does not normalize the active spins. + +For example:: + + potential nep.txt + velocity 300 + ensemble nve + time_step 1 + dump_xyz 100 trajectory.xyz spin mforce + dump_restart 1000 + run 1000 + +Use ``dump_xyz`` with ``spin`` and ``mforce`` to write the moments and magnetic +forces. ``dump_restart`` preserves the fixed spins in ``restart.xyz``; +copy this file to ``model.xyz`` and use the same potential to resume. +``replicate`` copies the spin vectors with their atoms. + +The :ref:`add_mfield keyword ` applies a magnetic induction in +tesla to a group or fractional-coordinate region. Magnetic force is +:math:`-\partial E/\partial \mathbf{s}` in eV/Bohr magneton; it is evaluated +although spins are fixed. Spin dynamics and spin velocities are not part of +this runtime interface. + +Multi-GPU execution uses every visible GPU in a single process. Use +``potential nep.txt x`` (or ``y``/``z``) to select the spatial partition +direction. Small cells may require replicated geometry and offer little +speedup. GPU 0 retains the global arrays. diff --git a/src/force/force.cu b/src/force/force.cu index 7e0887737..12229ef32 100644 --- a/src/force/force.cu +++ b/src/force/force.cu @@ -46,6 +46,8 @@ The driver class calculating force and related quantities. #include "utilities/gpu_macro.cuh" #include "utilities/nep_model_file.cuh" #include "utilities/nep_model_input.cuh" +#include "utilities/nep_spin_model_file.cuh" +#include "nep_spin/runtime.cuh" #include "utilities/read_file.cuh" #include "utilities/run_input.cuh" #include @@ -54,6 +56,15 @@ The driver class calculating force and related quantities. Force::Force(void) {} +Force::~Force() = default; + +void Force::initialize_spin(Atom& atom) +{ + if (spin_runtime_) { + spin_runtime_->initialize(atom); + } +} + void Force::check_types(const std::string& file_potential) { std::ifstream input(file_potential); @@ -114,6 +125,13 @@ std::unique_ptr Force::create_potential( strcmp(potential_name, "qnep2_vdw") == 0) { potential.reset(new NEP_Charge_VDW(tokens[1].c_str(), number_of_atoms, run_input)); is_nep = true; + } else if (strcmp(potential_name, "nep_spin") == 0) { + if (!potentials.empty()) { + PRINT_INPUT_ERROR("NEP_Spin must be the only potential in a run."); + } + spin_runtime_.reset(new Spin_Runtime(run_input)); + potential = spin_runtime_->create_potential(tokens, number_of_atoms); + is_nep = true; } else if ( strcmp(potential_name, "nep4_charge1") == 0 || strcmp(potential_name, "nep4_charge2") == 0 || @@ -234,7 +252,9 @@ void Force::parse_potential( PRINT_INPUT_ERROR("reading error for potential file."); } int num_types = get_number_of_types(fid_potential); - if (NEP_Model_Input::is_new_format(potential_name)) { + if (NEP_Model_Input::is_new_format(potential_name) && NEP_Spin_Model_File::is_spin_model(tokens[1])) { + std::strcpy(potential_name, "nep_spin"); + } else if (NEP_Model_Input::is_new_format(potential_name)) { NEP_Model_File model; std::ifstream input(tokens[1]); std::string error; @@ -276,6 +296,9 @@ void Force::parse_potential( // Move the pointer into the list of potentials potentials.push_back(std::move(potential)); // Check if a non-NEP potential has previously been defined + if (spin_runtime_ && potentials.size() != 1) { + PRINT_INPUT_ERROR("NEP_Spin must be the only potential in a run."); + } has_non_nep = has_non_nep || !is_nep; if (potentials.size() > 1 && has_non_nep) { PRINT_INPUT_ERROR("Multiple potentials may only be used with NEP potentials.\n"); @@ -613,3 +636,41 @@ void Force::compute( virial_per_atom); } + +void Force::compute_spin( + Box& box, + GPU_Vector& position_per_atom, + GPU_Vector& type, + GPU_Vector& potential_per_atom, + GPU_Vector& force_per_atom, + GPU_Vector& virial_per_atom, + GPU_Vector& spin_per_atom, + GPU_Vector& mforce_per_atom, + int* position_image) +{ + if (!spin_runtime_ || potentials.size() != 1) { + PRINT_INPUT_ERROR("Spin-aware force dispatch requires one NEP_Spin potential."); + } + + const int number_of_atoms = type.size(); + if ( + spin_per_atom.size() != static_cast(number_of_atoms) * 3 || + mforce_per_atom.size() != static_cast(number_of_atoms) * 3) { + PRINT_INPUT_ERROR("NEP_Spin requires three spin and mforce components per atom."); + } + prepare_compute( + number_of_atoms, box, position_per_atom, potential_per_atom, + force_per_atom, virial_per_atom, position_image); + mforce_per_atom.fill(0.0); + GPU_CHECK_KERNEL + + potentials[0]->compute( + box, + type, + position_per_atom, + spin_per_atom, + potential_per_atom, + force_per_atom, + virial_per_atom, + mforce_per_atom); +} diff --git a/src/force/force.cuh b/src/force/force.cuh index d53247c95..c017f109f 100644 --- a/src/force/force.cuh +++ b/src/force/force.cuh @@ -28,11 +28,14 @@ #include class RunInput; +class Spin_Runtime; +class Atom; class Force { public: Force(void); + ~Force(); void parse_potential( const std::vector& tokens, @@ -61,6 +64,21 @@ public: GPU_Vector& mass_per_atom, int* position_image = nullptr); + void compute_spin( + Box& box, + GPU_Vector& position_per_atom, + GPU_Vector& type, + GPU_Vector& potential_per_atom, + GPU_Vector& force_per_atom, + GPU_Vector& virial_per_atom, + GPU_Vector& spin_per_atom, + GPU_Vector& mforce_per_atom, + int* position_image = nullptr); + + bool has_spin_potential(void) const { return spin_runtime_ != nullptr; } + + void initialize_spin(Atom& atom); + // Zeroes the given arrays and evaluates one potential alone into them. The positions and the // float box are read as they are. void compute_one_potential( @@ -106,6 +124,7 @@ private: double temperature = 0; double delta_T; std::vector> potentials; + std::unique_ptr spin_runtime_; bool is_fcp = false; bool has_non_nep = false; std::string multiple_potentials_mode_ = "observe"; // "observe" or "average" diff --git a/src/force/nep.cu b/src/force/nep.cu index dd37b0070..86b8b2ae0 100644 --- a/src/force/nep.cu +++ b/src/force/nep.cu @@ -22,6 +22,7 @@ heat transport, Phys. Rev. B. 104, 104309 (2021). #include "neighbor.cuh" #include "nep.cuh" +#include "nep_kernels.cuh" #include "nep_small_box.cuh" #include "utilities/common.cuh" #include "utilities/compact_nep.cuh" @@ -1719,3 +1720,67 @@ void NEP::compute( const GPU_Vector& NEP::get_NN_radial_ptr() { return nep_data.NN_radial; } const GPU_Vector& NEP::get_NL_radial_ptr() { return nep_data.NL_radial; } + +void nep_internal::launch_find_force_radial( + NEP::ParaMB paramb, + NEP::ANN annmb, + int N, + int N1, + int N2, + Box box, + const int* NN, + const int* NL, + const int* type, + const double* x, + const double* y, + const double* z, + const float* Fp, + double* fx, + double* fy, + double* fz, + double* virial) +{ + find_force_radial<<<(N2 - N1 + 63) / 64, 64>>>( + paramb, annmb, N, N1, N2, box, NN, NL, type, x, y, z, Fp, fx, fy, fz, virial); + GPU_CHECK_KERNEL +} + +void nep_internal::launch_find_partial_force_angular( + NEP::ParaMB paramb, + NEP::ANN annmb, + int N, + int N1, + int N2, + Box box, + const int* NN_angular, + const int* NL_angular, + const int* type, + const double* x, + const double* y, + const double* z, + const float* Fp, + const float* sum_fxyz, + float* f12x, + float* f12y, + float* f12z) +{ + find_partial_force_angular<<<(N2 - N1 + 63) / 64, 64>>>( + paramb, + annmb, + N, + N1, + N2, + box, + NN_angular, + NL_angular, + type, + x, + y, + z, + Fp, + sum_fxyz, + f12x, + f12y, + f12z); + GPU_CHECK_KERNEL +} diff --git a/src/force/nep_kernels.cuh b/src/force/nep_kernels.cuh new file mode 100644 index 000000000..48df60c5e --- /dev/null +++ b/src/force/nep_kernels.cuh @@ -0,0 +1,46 @@ +/* Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. SPDX-License-Identifier: LGPL-3.0-or-later */ +#pragma once +#include "nep.cuh" + +// Internal structural-force launchers. CUDA kernels remain translation-unit local. +namespace nep_internal +{ +void launch_find_force_radial( + NEP::ParaMB paramb, + NEP::ANN annmb, + int N, + int N1, + int N2, + Box box, + const int* NN, + const int* NL, + const int* type, + const double* x, + const double* y, + const double* z, + const float* Fp, + double* fx, + double* fy, + double* fz, + double* virial); +void launch_find_partial_force_angular( + NEP::ParaMB paramb, + NEP::ANN annmb, + int N, + int N1, + int N2, + Box box, + const int* NN_angular, + const int* NL_angular, + const int* type, + const double* x, + const double* y, + const double* z, + const float* Fp, + const float* sum_fxyz, + float* f12x, + float* f12y, + float* f12z); + +} // namespace nep_internal diff --git a/src/force/nep_spin/mfield.cu b/src/force/nep_spin/mfield.cu new file mode 100644 index 000000000..8af456682 --- /dev/null +++ b/src/force/nep_spin/mfield.cu @@ -0,0 +1,254 @@ +/* + Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. + GPUMD is free software: you can redistribute it and/or modify + it under the terms of the GNU Lesser General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + GPUMD is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Lesser General Public License for more details. + You should have received a copy of the GNU Lesser General Public License + along with GPUMD. If not, see . +*/ + +#include "force/force.cuh" +#include "mfield.cuh" +#include "model/box.cuh" +#include "nep_spin.cuh" +#include "nep_spin_multigpu.cuh" +#include "utilities/error.cuh" +#include "utilities/gpu_macro.cuh" +#include "utilities/read_file.cuh" +#include +#include + +Spin_Magnetic_Field::Spin_Magnetic_Field( + const std::vector& tokens, const std::vector& group) +{ + action_name = "add_mfield"; + std::vector fields; + for (const auto& token : tokens) { + fields.push_back(token.c_str()); + } + parse(fields.data(), fields.size(), group); +} + +void Spin_Magnetic_Field::pre_run( + int, double, Integrate&, std::vector&, Atom& atom, Box&, Force& force) +{ + const std::vector* types = nullptr; + if (force.get_number_of_potentials() == 1) { + auto& potential = force.get_potential(0); + if (const auto* spin = dynamic_cast(&potential)) { + types = &spin->spin_types(); + } else if (const auto* spin = dynamic_cast(&potential)) { + types = &spin->spin_types(); + } + } + if (!types || !atom.has_spin) { + PRINT_INPUT_ERROR("add_mfield requires a Spin potential and spin:R:3."); + } + active_types_.resize(types->size()); + active_types_.copy_from_host(types->data()); +} + +void Spin_Magnetic_Field::setup_force( + double, Integrate&, std::vector& group, Atom& atom, Box& box, Force&) +{ + compute(box, group, atom); +} + +void Spin_Magnetic_Field::post_force( + int, double, Integrate&, std::vector& group, Atom& atom, Box& box, Force&) +{ + compute(box, group, atom); +} + +static __device__ void get_fractional_position( + const Box& box, + const double x, + const double y, + const double z, + double& sa, + double& sb, + double& sc) +{ + sa = box.cpu_h[9] * x + box.cpu_h[10] * y + box.cpu_h[11] * z; + sb = box.cpu_h[12] * x + box.cpu_h[13] * y + box.cpu_h[14] * z; + sc = box.cpu_h[15] * x + box.cpu_h[16] * y + box.cpu_h[17] * z; +} + +static __device__ bool is_in_region( + const double sa, + const double sb, + const double sc, + const double amin, + const double amax, + const double bmin, + const double bmax, + const double cmin, + const double cmax) +{ + return sa >= amin && sa < amax && sb >= bmin && sb < bmax && sc >= cmin && sc < cmax; +} + +// spin is in muB; this converts a tesla-valued B to magnetic force in eV/muB. +void Spin_Magnetic_Field::parse(const char** param, int num_param, const std::vector& group) +{ + mfield_use_region_ = num_param > 1 && strcmp(param[1], "region") == 0; + if (num_param != (mfield_use_region_ ? 11 : 6)) { + PRINT_INPUT_ERROR("add_mfield requires group_method group_id Bx By Bz or " + "region amin amax bmin bmax cmin cmax Bx By Bz."); + } + if (mfield_use_region_) { + for (int k = 0; k < 6; ++k) { + if ( + !is_valid_real(param[k + 2], &mfield_region_[k]) || !std::isfinite(mfield_region_[k]) || + mfield_region_[k] < 0.0 || mfield_region_[k] > 1.0) { + PRINT_INPUT_ERROR("add_mfield region bounds must be finite numbers in [0, 1]."); + } + } + for (int k = 0; k < 6; k += 2) { + if (mfield_region_[k] >= mfield_region_[k + 1]) { + PRINT_INPUT_ERROR("add_mfield region minimum must be smaller than maximum."); + } + } + } else { + if ( + !is_valid_int(param[1], &mfield_group_method_) || mfield_group_method_ < 0 || + mfield_group_method_ >= group.size()) { + PRINT_INPUT_ERROR("add_mfield grouping method is invalid."); + } + if ( + !is_valid_int(param[2], &mfield_group_id_) || mfield_group_id_ < 0 || + mfield_group_id_ >= group[mfield_group_method_].number) { + PRINT_INPUT_ERROR("add_mfield group id is invalid."); + } + } + const int field_start = mfield_use_region_ ? 8 : 3; + constexpr double muB_eV_per_T = 5.7883818060e-5; + double field[3]; + mfield_enabled_ = false; + for (int k = 0; k < 3; ++k) { + if (!is_valid_real(param[k + field_start], &field[k]) || !std::isfinite(field[k])) { + PRINT_INPUT_ERROR("add_mfield components must be finite numbers in tesla."); + } + mfield_[k] = field[k] * muB_eV_per_T; + mfield_enabled_ = mfield_enabled_ || field[k] != 0.0; + } + printf("External magnetic field: (%g, %g, %g) T.\n", field[0], field[1], field[2]); + if (mfield_use_region_) { + printf( + " fractional region [%g, %g) [%g, %g) [%g, %g).\n", + mfield_region_[0], + mfield_region_[1], + mfield_region_[2], + mfield_region_[3], + mfield_region_[4], + mfield_region_[5]); + } else { + printf(" grouping method %d, group %d.\n", mfield_group_method_, mfield_group_id_); + } +} + +template +static __global__ void add_mfield( + const int n, + const int count, + const int offset, + const int* members, + const double hx, + const double hy, + const double hz, + const Box box, + const double* position, + const double amin, + const double amax, + const double bmin, + const double bmax, + const double cmin, + const double cmax, + const double* spin, + const int* type, + const int* active, + double* energy, + double* mforce) +{ + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= count) + return; + const int i = UseRegion ? tid : members[offset + tid]; + if (!active[type[i]]) { + return; + } + if (UseRegion) { + double sa, sb, sc; + get_fractional_position(box, position[i], position[n + i], position[2 * n + i], sa, sb, sc); + if (!is_in_region(sa, sb, sc, amin, amax, bmin, bmax, cmin, cmax)) + return; + } + energy[i] -= spin[i] * hx + spin[n + i] * hy + spin[2 * n + i] * hz; + mforce[i] += hx; + mforce[n + i] += hy; + mforce[2 * n + i] += hz; +} + +void Spin_Magnetic_Field::compute(Box& box, const std::vector& group, Atom& atom) const +{ + const int number_of_atoms = atom.number_of_atoms; + if (mfield_enabled_) { + if (mfield_use_region_) { + add_mfield<<<(number_of_atoms + 127) / 128, 128>>>( + number_of_atoms, + number_of_atoms, + 0, + nullptr, + mfield_[0], + mfield_[1], + mfield_[2], + box, + atom.position_per_atom.data(), + mfield_region_[0], + mfield_region_[1], + mfield_region_[2], + mfield_region_[3], + mfield_region_[4], + mfield_region_[5], + atom.spin_per_atom.data(), + atom.type.data(), + active_types_.data(), + atom.potential_per_atom.data(), + atom.mforce_per_atom.data()); + GPU_CHECK_KERNEL + } else { + const Group& selected = group[mfield_group_method_]; + const int count = selected.cpu_size[mfield_group_id_]; + if (count > 0) { + add_mfield<<<(count + 127) / 128, 128>>>( + number_of_atoms, + count, + selected.cpu_size_sum[mfield_group_id_], + selected.contents.data(), + mfield_[0], + mfield_[1], + mfield_[2], + box, + nullptr, + 0, + 0, + 0, + 0, + 0, + 0, + atom.spin_per_atom.data(), + atom.type.data(), + active_types_.data(), + atom.potential_per_atom.data(), + atom.mforce_per_atom.data()); + GPU_CHECK_KERNEL + } + } + } +} diff --git a/src/force/nep_spin/mfield.cuh b/src/force/nep_spin/mfield.cuh new file mode 100644 index 000000000..88087034a --- /dev/null +++ b/src/force/nep_spin/mfield.cuh @@ -0,0 +1,60 @@ +/* + Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. + GPUMD is free software: you can redistribute it and/or modify + it under the terms of the GNU Lesser General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + GPUMD is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Lesser General Public License for more details. + You should have received a copy of the GNU Lesser General Public License + along with GPUMD. If not, see . +*/ + +#pragma once + +#include "measure/action.cuh" +#include + +class Spin_Magnetic_Field : public Action +{ +public: + Spin_Magnetic_Field(const std::vector& tokens, const std::vector& group); + + void pre_run( + int number_of_steps, + double time_step, + Integrate& integrate, + std::vector& group, + Atom& atom, + Box& box, + Force& force) override; + void setup_force( + double time_step, + Integrate& integrate, + std::vector& group, + Atom& atom, + Box& box, + Force& force) override; + void post_force( + int step, + double time_step, + Integrate& integrate, + std::vector& group, + Atom& atom, + Box& box, + Force& force) override; + +private: + void parse(const char** param, int num_param, const std::vector& group); + void compute(Box& box, const std::vector& group, Atom& atom) const; + GPU_Vector active_types_; + int mfield_group_method_ = 0; + int mfield_group_id_ = 0; + bool mfield_use_region_ = false; + double mfield_region_[6] = {}; + bool mfield_enabled_ = false; + double mfield_[3] = {}; // eV/muB +}; diff --git a/src/force/nep_spin/neighbors.cu b/src/force/nep_spin/neighbors.cu new file mode 100644 index 000000000..2fb68347e --- /dev/null +++ b/src/force/nep_spin/neighbors.cu @@ -0,0 +1,185 @@ +/* Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. SPDX-License-Identifier: LGPL-3.0-or-later */ +#include "nep_spin.cuh" +#include "utilities/error.cuh" +#include + +namespace +{ +// The same cell traversal as Neighbor, with checked writes and per-instance state. +static __global__ void find_spin_neighbors( + const Box box, + const int N, + const int N1, + const int N2, + const int* __restrict__ type, + const int* __restrict__ cell_counts, + const int* __restrict__ cell_count_sum, + const int* __restrict__ cell_contents, + int* NN, + int* NL, + const double* __restrict__ x, + const double* __restrict__ y, + const double* __restrict__ z, + const int nx, + const int ny, + const int nz, + const double rc_inv, + const float cutoff_square, + const int neighbor_capacity, + int* required_capacity) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x + N1; + int count = 0; + if (n1 < N2) { + const double x1 = x[n1]; + const double y1 = y[n1]; + const double z1 = z[n1]; + int cell_id; + int cell_id_x; + int cell_id_y; + int cell_id_z; + find_cell_id(box, x1, y1, z1, rc_inv, nx, ny, nz, cell_id_x, cell_id_y, cell_id_z, cell_id); + + int z_begin, z_end; + int y_begin, y_end; + int x_begin, x_end; + get_neighbor_cell_offset_range(box.pbc_z, nz, z_begin, z_end); + get_neighbor_cell_offset_range(box.pbc_y, ny, y_begin, y_end); + get_neighbor_cell_offset_range(box.pbc_x, nx, x_begin, x_end); + + // get radial descriptors + for (int k = z_begin; k < z_end; ++k) { + for (int j = y_begin; j < y_end; ++j) { + for (int i = x_begin; i < x_end; ++i) { + int neighbor_cell = cell_id + k * nx * ny + j * nx + i; + if (cell_id_x + i < 0) + neighbor_cell += nx; + else if (cell_id_x + i >= nx) + neighbor_cell -= nx; + if (cell_id_y + j < 0) + neighbor_cell += ny * nx; + else if (cell_id_y + j >= ny) + neighbor_cell -= ny * nx; + if (cell_id_z + k < 0) + neighbor_cell += nz * ny * nx; + else if (cell_id_z + k >= nz) + neighbor_cell -= nz * ny * nx; + + const int num_atoms_neighbor_cell = cell_counts[neighbor_cell]; + const int num_atoms_previous_cells = cell_count_sum[neighbor_cell]; + + for (int m = 0; m < num_atoms_neighbor_cell; ++m) { + const int n2 = cell_contents[num_atoms_previous_cells + m]; + if (n2 >= N1 && n2 < N2 && n1 != n2) { + + float x12 = x[n2] - x1; + float y12 = y[n2] - y1; + float z12 = z[n2] - z1; + apply_mic(box, x12, y12, z12); + const float d2 = x12 * x12 + y12 * y12 + z12 * z12; + + if (d2 < cutoff_square) { + if (count < neighbor_capacity) { + NL[static_cast(N) * count + n1] = n2; + } + ++count; + } + } + } + } + } + } + NN[n1] = min(count, neighbor_capacity); + atomicMax(required_capacity, count); + } +} + +static __global__ void +check_displacement(int n, Box box, const double* position, const double* reference, int* changed) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + double x = position[i] - reference[i]; + double y = position[n + i] - reference[n + i]; + double z = position[2 * n + i] - reference[2 * n + i]; + apply_mic(box, x, y, z); + if (x * x + y * y + z * z > 0.25) + atomicExch(changed, 1); +} +} // namespace + +void NEP_Spin::Neighbors::initialize(double rc, int atoms, int capacity) +{ + valid_ = false; + const double buffered = rc + skin_; + capacity_ = capacity * buffered * buffered * buffered / (rc * rc * rc); + if (NN.size() != atoms) { + NN.resize(atoms); + cell_count_.resize(atoms); + cell_sum_.resize(atoms); + cell_contents_.resize(atoms); + reference_.resize(size_t(atoms) * 3); + } + if (NL.size() != size_t(atoms) * capacity_) + NL.resize(size_t(atoms) * capacity_); + if (flag_.size() != 1) + flag_.resize(1); +} + +void NEP_Spin::Neighbors::find_neighbor_global( + double rc, Box& box, const GPU_Vector& type, const GPU_Vector& position) +{ + const int n = type.size(); + const int grid = (n + 255) / 256; + int changed = !valid_ || !std::equal(box_, box_ + 9, box.cpu_h); + if (!changed) { + flag_.fill(0); + check_displacement<<>>(n, box, position.data(), reference_.data(), flag_.data()); + GPU_CHECK_KERNEL + flag_.copy_to_host(&changed); + } + if (!changed) + return; + const double cutoff = rc + skin_; + int bins[3]; + box.get_num_bins(0.5 * cutoff, bins); + find_cell_list(0.5 * cutoff, bins, box, position, cell_count_, cell_sum_, cell_contents_); + flag_.fill(0); + find_spin_neighbors<<>>( + box, + n, + 0, + n, + type.data(), + cell_count_.data(), + cell_sum_.data(), + cell_contents_.data(), + NN.data(), + NL.data(), + position.data(), + position.data() + n, + position.data() + 2 * n, + bins[0], + bins[1], + bins[2], + 2.0 / cutoff, + cutoff * cutoff, + capacity_, + flag_.data()); + GPU_CHECK_KERNEL + int required; + flag_.copy_to_host(&required); + if (required > capacity_) { + PRINT_INPUT_ERROR( + "NEP-spin neighbor capacity is insufficient; increase the model neighbor capacity."); + } + gpu_sort_neighbor_list<<>>( + n, NN.data(), NL.data()); + GPU_CHECK_KERNEL + CHECK(gpuMemcpy( + reference_.data(), position.data(), size_t(n) * 3 * sizeof(double), gpuMemcpyDeviceToDevice)); + std::copy(box.cpu_h, box.cpu_h + 9, box_); + valid_ = true; +} diff --git a/src/force/nep_spin/nep_spin.cu b/src/force/nep_spin/nep_spin.cu new file mode 100644 index 000000000..c653e2ce0 --- /dev/null +++ b/src/force/nep_spin/nep_spin.cu @@ -0,0 +1,1974 @@ +/* + Copyright 2017 Zheyong Fan and GPUMD development team + Copyright 2026 NEPAdapters contributors + This file is part of GPUMD and is distributed under LGPLv3 or later. + +*/ + +#include "force/nep_kernels.cuh" +#include "nep_spin.cuh" +#include "nep_spin_small_box.cuh" +#include "utilities/error.cuh" +#include "utilities/gpu_macro.cuh" +#include "utilities/nep_spin_model_file.cuh" +#include "utilities/nep_utilities.cuh" +#include +#include +#include +#include +#include + +namespace +{ + +using SimulationBox = Box; + +struct Spin_Polynomial_Layout { + int channels = 0; + int pair_count = 0; + int density_stride = 38; + int moment_count = 0; + int local_s2 = -1; + int edge_l0_dot = -1; + int edge_l0_neighbor_s2 = -1; + int edge_l2_pair = -1; + int center_l2_environment = -1; + int edge_l0_dot2 = -1; + int density_l0_self = -1; + int density_l1_longitudinal_self = -1; + int density_l1_axial_self = -1; + int density_l1_stf_self = -1; + int density_l1_product_self = -1; + int density_l2_product_self = -1; + int density_l0_dot_response = -1; + int correlation_same_edge = -1; + int correlation_distinct_neighbor = -1; + int correlation_distinct_l1 = -1; + int correlation_distinct_l2 = -1; + int angular_l1_moment_offset = -1; + int angular_l2_moment_offset = -1; + int coupling_l11_axial = -1; + int edge_l11_axial = -1; + int coupling_l22_axial = -1; + int edge_l22_axial = -1; + int edge_l0_moment_gate = -1; + int coupling_l11_dot_response = -1; + int coupling_l111_p_m_x = -1; + int coupling_l22_dot_response = -1; + int coupling_l111_p_qs_x = -1; + int coupling_l112_edge_response = -1; + int coupling_l111_bulk = -1; + int descriptor_dim = 0; +}; + +#include "nep_spin3_common.cuh" + +enum class SpinVirialMode : int { + disabled, + center_owned, + neighbor_owned, + center_and_neighbor_float_sink, +}; + +constexpr const char* kElementSymbols[NUM_ELEMENTS] = { + "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P", "S", + "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn", "Ga", "Ge", + "As", "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", + "In", "Sn", "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", + "Tb", "Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg", + "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa", "U", "Np", "Pu"}; + +__device__ __forceinline__ void atomic_add_per_atom_virial_double( + const int atom_stride, + const int atom, + const float x12, + const float y12, + const float z12, + const float fx, + const float fy, + const float fz, + double* virial) +{ + atomicAdd(&virial[atom], -static_cast(x12 * fx)); + atomicAdd(&virial[atom_stride + atom], -static_cast(y12 * fy)); + atomicAdd(&virial[2 * atom_stride + atom], -static_cast(z12 * fz)); + atomicAdd(&virial[3 * atom_stride + atom], -static_cast(x12 * fy)); + atomicAdd(&virial[4 * atom_stride + atom], -static_cast(x12 * fz)); + atomicAdd(&virial[5 * atom_stride + atom], -static_cast(y12 * fz)); + atomicAdd(&virial[6 * atom_stride + atom], -static_cast(y12 * fx)); + atomicAdd(&virial[7 * atom_stride + atom], -static_cast(z12 * fx)); + atomicAdd(&virial[8 * atom_stride + atom], -static_cast(z12 * fy)); +} + +__device__ __forceinline__ void atomic_add_per_atom_virial_float( + const int atom_stride, + const int atom, + const float x12, + const float y12, + const float z12, + const float fx, + const float fy, + const float fz, + float* virial) +{ + atomicAdd(&virial[atom], -x12 * fx); + atomicAdd(&virial[atom_stride + atom], -y12 * fy); + atomicAdd(&virial[2 * atom_stride + atom], -z12 * fz); + atomicAdd(&virial[3 * atom_stride + atom], -x12 * fy); + atomicAdd(&virial[4 * atom_stride + atom], -x12 * fz); + atomicAdd(&virial[5 * atom_stride + atom], -y12 * fz); + atomicAdd(&virial[6 * atom_stride + atom], -y12 * fx); + atomicAdd(&virial[7 * atom_stride + atom], -z12 * fx); + atomicAdd(&virial[8 * atom_stride + atom], -z12 * fy); +} + +__device__ __forceinline__ void +atomic_add_spin_transfer_float(int, int, float, float, float, const float*, float*) +{ +} + +__global__ void mask_inactive_spin_mforce( + const int atom_count, + const int atom_stride, + const int* __restrict__ types, + const int* __restrict__ spin_type_active, + double* __restrict__ mforce) +{ + const int atom = blockIdx.x * blockDim.x + threadIdx.x; + if (atom < atom_count && spin_type_active[types[atom]] == 0) { + mforce[atom] = 0.0; + mforce[atom_stride + atom] = 0.0; + mforce[2 * atom_stride + atom] = 0.0; + } +} + +template +__device__ __forceinline__ void evaluate_spin3_edge_weights_b9_f32( + const float spin_cutoff, + const float dist, + const int num_types, + const int type_pair, + const float* __restrict__ descriptor_coefficients, + const int spin_coefficient_offset, + float* weights, + float* weight_derivatives) +{ + constexpr int BasisCount = 9; + constexpr float Pi = 3.14159265358979323846f; + const int pair_count = num_types * num_types; + const float rcinv = 1.0f / spin_cutoff; + const float r_scaled = dist * rcinv; + const float phase = Pi * r_scaled; + const float fc = 0.5f * cosf(phase) + 0.5f; + const float shifted = r_scaled - 1.0f; + const float x = 2.0f * shifted * shifted - 1.0f; + const float dxdr = 2.0f * shifted * rcinv; + const float fcp = NeedDerivatives ? -0.5f * Pi * sinf(phase) * rcinv : 0.0f; + float fn[BasisCount] = {}; + float fnp[BasisCount] = {}; + fn[0] = fc; + fnp[0] = fcp; + const float raw1 = 0.5f * (x + 1.0f); + fn[1] = raw1 * fc; + fnp[1] = dxdr * fc + raw1 * fcp; + float t0 = 1.0f; + float t1 = x; + float u0 = 1.0f; + float u1 = 2.0f * x; + for (int n = 2; n < BasisCount; ++n) { + const float t2 = 2.0f * x * t1 - t0; + const float raw = 0.5f * (t2 + 1.0f); + fn[n] = raw * fc; + fnp[n] = static_cast(n) * u1 * dxdr * fc + raw * fcp; + const float u2 = 2.0f * x * u1 - u0; + t0 = t1; + t1 = t2; + u0 = u1; + u1 = u2; + } + for (int c = 0; c < C; ++c) { + float weight = 0.0f; + float derivative = 0.0f; + for (int k = 0; k < BasisCount; ++k) { + const int index = spin_coefficient_offset + (c * BasisCount + k) * pair_count + type_pair; + weight += fn[k] * descriptor_coefficients[index]; + if (NeedDerivatives) { + derivative += fnp[k] * descriptor_coefficients[index]; + } + } + weights[c] = weight; + if (NeedDerivatives) { + weight_derivatives[c] = derivative; + } + } +} + +template +__device__ __forceinline__ bool load_spin3_edge_f32( + const int atom, + const int neighbor, + const int atom_stride, + const int num_types, + const int spin_basis_size, + const float spin_cutoff, + const float* __restrict__ spin_cutoff_pair, + const SimulationBox box, + const int* __restrict__ types, + const double* __restrict__ positions_soa3, + const double* __restrict__ slot_r12, + const int r12_plane_size, + const int slot_index, + const double* __restrict__ spins_soa3, + const float* __restrict__ descriptor_coefficients, + const int spin_coefficient_offset, + float* rhat, + float& dist, + float* si, + float* sj, + float* weights, + float* weight_derivatives) +{ + if (slot_r12) { + const float dx = static_cast(slot_r12[slot_index]); + const float dy = static_cast(slot_r12[r12_plane_size + slot_index]); + const float dz = static_cast(slot_r12[2 * r12_plane_size + slot_index]); + dist = sqrtf(dx * dx + dy * dy + dz * dz); + if (dist > 1.0e-12f) { + const float inverse_distance = 1.0f / dist; + rhat[0] = dx * inverse_distance; + rhat[1] = dy * inverse_distance; + rhat[2] = dz * inverse_distance; + } + } else { + compute_spin_edge_geometry_f32(atom, neighbor, atom_stride, box, positions_soa3, rhat, dist); + } + for (int d = 0; d < 3; ++d) { + si[d] = static_cast(spins_soa3[d * atom_stride + atom]); + sj[d] = static_cast(spins_soa3[d * atom_stride + neighbor]); + } + const int type_pair = types[atom] * num_types + types[neighbor]; + const float resolved_spin_cutoff = + spin_cutoff_pair == nullptr ? spin_cutoff : spin_cutoff_pair[type_pair]; + if (!(dist > 1.0e-12f && dist < resolved_spin_cutoff) || spin_basis_size != 8) { + return false; + } + evaluate_spin3_edge_weights_b9_f32( + resolved_spin_cutoff, + dist, + num_types, + type_pair, + descriptor_coefficients, + spin_coefficient_offset, + weights, + weight_derivatives); + return true; +} + +#include "nep_spin3_descriptor.cuh" +#include "nep_spin3_force.cuh" +#include "nep_spin3_layout.cuh" + +__global__ void find_local_neighbors_spin( + const int N, + const Box box, + const float cutoff_radial_square, + const float cutoff_angular_square, + const float cutoff_spin_square, + const double* __restrict__ x, + const double* __restrict__ y, + const double* __restrict__ z, + const int* __restrict__ NN_global, + const int* __restrict__ NL_global, + int* NN_radial, + int* NL_radial, + int* NN_angular, + int* NL_angular, + int* NN_spin, + int* NL_spin) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x; + if (n1 >= N) { + return; + } + + const double x1 = x[n1]; + const double y1 = y[n1]; + const double z1 = z[n1]; + int count_radial = 0; + int count_angular = 0; + int count_spin = 0; + + for (int slot = 0; slot < NN_global[n1]; ++slot) { + const int n2 = NL_global[static_cast(N) * slot + n1]; + float x12 = x[n2] - x1; + float y12 = y[n2] - y1; + float z12 = z[n2] - z1; + apply_mic(box, x12, y12, z12); + const float distance_square = x12 * x12 + y12 * y12 + z12 * z12; + + if (distance_square < cutoff_radial_square) { + NL_radial[static_cast(N) * count_radial++ + n1] = n2; + } + if (distance_square < cutoff_angular_square) { + NL_angular[static_cast(N) * count_angular++ + n1] = n2; + } + if (distance_square < cutoff_spin_square) { + NL_spin[static_cast(N) * count_spin++ + n1] = n2; + } + } + + NN_radial[n1] = count_radial; + NN_angular[n1] = count_angular; + NN_spin[n1] = count_spin; +} + +__device__ __forceinline__ void apply_mic_spin_small_box( + const Box& box, const NEP::ExpandedBox& ebox, double& x12, double& y12, double& z12) +{ + double sx12 = + (box.cpu_h[9] * x12 + box.cpu_h[10] * y12 + box.cpu_h[11] * z12) / ebox.num_cells[0]; + double sy12 = + (box.cpu_h[12] * x12 + box.cpu_h[13] * y12 + box.cpu_h[14] * z12) / ebox.num_cells[1]; + double sz12 = + (box.cpu_h[15] * x12 + box.cpu_h[16] * y12 + box.cpu_h[17] * z12) / ebox.num_cells[2]; + if (box.pbc_x == 1) + sx12 -= nearbyint(sx12); + if (box.pbc_y == 1) + sy12 -= nearbyint(sy12); + if (box.pbc_z == 1) + sz12 -= nearbyint(sz12); + x12 = box.cpu_h[0] * ebox.num_cells[0] * sx12 + box.cpu_h[1] * ebox.num_cells[1] * sy12 + + box.cpu_h[2] * ebox.num_cells[2] * sz12; + y12 = box.cpu_h[3] * ebox.num_cells[0] * sx12 + box.cpu_h[4] * ebox.num_cells[1] * sy12 + + box.cpu_h[5] * ebox.num_cells[2] * sz12; + z12 = box.cpu_h[6] * ebox.num_cells[0] * sx12 + box.cpu_h[7] * ebox.num_cells[1] * sy12 + + box.cpu_h[8] * ebox.num_cells[2] * sz12; +} + +__global__ void find_neighbor_list_spin_small_box( + const NEP::ParaMB paramb, + const float spin_cutoff, + const int N, + const Box box, + const NEP::ExpandedBox ebox, + const int capacity, + const int* type, + const double* x, + const double* y, + const double* z, + int* NN_radial, + int* NL_radial, + int* NN_angular, + int* NL_angular, + int* NN_spin, + int* NL_spin, + double* r12_radial, + double* r12_angular, + double* r12_spin, + int* required_capacity) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x; + if (n1 >= N) { + return; + } + const double x1 = x[n1]; + const double y1 = y[n1]; + const double z1 = z[n1]; + const int t1 = type[n1]; + const int plane_size = N * capacity; + int count_radial = 0; + int count_angular = 0; + int count_spin = 0; + for (int n2 = 0; n2 < N; ++n2) { + for (int ia = 0; ia < ebox.num_cells[0]; ++ia) { + for (int ib = 0; ib < ebox.num_cells[1]; ++ib) { + for (int ic = 0; ic < ebox.num_cells[2]; ++ic) { + if (n1 == n2 && ia == 0 && ib == 0 && ic == 0) { + continue; + } + double x12 = x[n2] + box.cpu_h[0] * ia + box.cpu_h[1] * ib + box.cpu_h[2] * ic - x1; + double y12 = y[n2] + box.cpu_h[3] * ia + box.cpu_h[4] * ib + box.cpu_h[5] * ic - y1; + double z12 = z[n2] + box.cpu_h[6] * ia + box.cpu_h[7] * ib + box.cpu_h[8] * ic - z1; + apply_mic_spin_small_box(box, ebox, x12, y12, z12); + const double d2 = x12 * x12 + y12 * y12 + z12 * z12; + const int t2 = type[n2]; + const float rc_radial = 0.5f * (paramb.rc_radial[t1] + paramb.rc_radial[t2]); + const float rc_angular = 0.5f * (paramb.rc_angular[t1] + paramb.rc_angular[t2]); + if (d2 < rc_radial * rc_radial) { + const int slot = count_radial++; + if (slot < capacity) { + const int index = n1 + N * slot; + NL_radial[index] = n2; + r12_radial[index] = x12; + r12_radial[plane_size + index] = y12; + r12_radial[2 * plane_size + index] = z12; + } + } + if (d2 < rc_angular * rc_angular) { + const int slot = count_angular++; + if (slot < capacity) { + const int index = n1 + N * slot; + NL_angular[index] = n2; + r12_angular[index] = x12; + r12_angular[plane_size + index] = y12; + r12_angular[2 * plane_size + index] = z12; + } + } + if (d2 < spin_cutoff * spin_cutoff) { + const int slot = count_spin++; + if (slot < capacity) { + const int index = n1 + N * slot; + NL_spin[index] = n2; + r12_spin[index] = x12; + r12_spin[plane_size + index] = y12; + r12_spin[2 * plane_size + index] = z12; + } + } + } + } + } + } + NN_radial[n1] = min(count_radial, capacity); + NN_angular[n1] = min(count_angular, capacity); + NN_spin[n1] = min(count_spin, capacity); + atomicMax(required_capacity, max(count_radial, max(count_angular, count_spin))); +} + +__global__ void find_structural_descriptor( + const NEP::ParaMB paramb, + const NEP::ANN ann, + const int N, + const Box box, + const int* NN_radial, + const int* NL_radial, + const int* NN_angular, + const int* NL_angular, + const double* r12_radial, + const double* r12_angular, + const int r12_plane_size, + const int* type, + const double* x, + const double* y, + const double* z, + float* descriptor, + float* sum_fxyz) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x; + if (n1 >= N) { + return; + } + const int t1 = type[n1]; + const double x1 = x[n1]; + const double y1 = y[n1]; + const double z1 = z[n1]; + float q[MAX_DIM] = {0.0f}; + + for (int i1 = 0; i1 < NN_radial[n1]; ++i1) { + const int n2 = NL_radial[static_cast(N) * i1 + n1]; + const int slot = N * i1 + n1; + double x12 = r12_radial ? r12_radial[slot] : x[n2] - x1; + double y12 = r12_radial ? r12_radial[r12_plane_size + slot] : y[n2] - y1; + double z12 = r12_radial ? r12_radial[2 * r12_plane_size + slot] : z[n2] - z1; + if (!r12_radial) { + apply_mic(box, x12, y12, z12); + } + const float d12 = sqrtf(x12 * x12 + y12 * y12 + z12 * z12); + const int t2 = type[n2]; + const float rc = (paramb.rc_radial[t1] + paramb.rc_radial[t2]) * 0.5f; + const float rcinv = 1.0f / rc; + float fc12; + find_fc(rc, rcinv, d12, fc12); + float fn12[MAX_NUM_N]; + find_fn(paramb.basis_size_radial, rcinv, d12, fc12, fn12); + const int basis_count = (paramb.n_max_radial + 1) * (paramb.basis_size_radial + 1); + for (int n = 0; n <= paramb.n_max_radial; ++n) { + float gn12 = 0.0f; + for (int k = 0; k <= paramb.basis_size_radial; ++k) { + const int basis = n * (paramb.basis_size_radial + 1) + k; + gn12 += fn12[k] * ann.c_type_pair[(t1 * paramb.num_types + t2) * basis_count + basis]; + } + q[n] += gn12; + } + } + + for (int n = 0; n <= paramb.n_max_angular; ++n) { + float s[NUM_OF_ABC] = {0.0f}; + const int abc_count = (paramb.L_max + 1) * (paramb.L_max + 1) - 1; + for (int i1 = 0; i1 < NN_angular[n1]; ++i1) { + const int n2 = NL_angular[n1 + N * i1]; + const int slot = N * i1 + n1; + double x12 = r12_angular ? r12_angular[slot] : x[n2] - x1; + double y12 = r12_angular ? r12_angular[r12_plane_size + slot] : y[n2] - y1; + double z12 = r12_angular ? r12_angular[2 * r12_plane_size + slot] : z[n2] - z1; + if (!r12_angular) { + apply_mic(box, x12, y12, z12); + } + const float d12 = sqrtf(x12 * x12 + y12 * y12 + z12 * z12); + const int t2 = type[n2]; + const float rc = (paramb.rc_angular[t1] + paramb.rc_angular[t2]) * 0.5f; + const float rcinv = 1.0f / rc; + float fc12; + find_fc(rc, rcinv, d12, fc12); + float fn12[MAX_NUM_N]; + find_fn(paramb.basis_size_angular, rcinv, d12, fc12, fn12); + float gn12 = 0.0f; + for (int k = 0; k <= paramb.basis_size_angular; ++k) { + int c_index = paramb.num_c_radial; + c_index += (t1 * paramb.num_types + t2) * + ((paramb.n_max_angular + 1) * (paramb.basis_size_angular + 1)); + c_index += n * (paramb.basis_size_angular + 1) + k; + gn12 += fn12[k] * ann.c_type_pair[c_index]; + } + accumulate_s(paramb.L_max, d12, x12, y12, z12, gn12, s); + } + find_q( + paramb.L_max, + paramb.has_q_222, + paramb.has_q_1111, + paramb.has_q_112, + paramb.has_q_123, + paramb.has_q_233, + paramb.has_q_134, + paramb.n_max_angular + 1, + n, + s, + q + paramb.n_max_radial + 1); + for (int abc = 0; abc < abc_count; ++abc) { + sum_fxyz[static_cast(N) * (n * abc_count + abc) + n1] = s[abc]; + } + } + const int struct_dim = paramb.n_max_radial + 1 + paramb.dim_angular; + for (int d = 0; d < struct_dim; ++d) { + descriptor[static_cast(N) * d + n1] = q[d]; + } +} + +__global__ void find_potential( + const int N, + const int descriptor_dim, + const int hidden_neurons, + const int num_types, + const int* type, + const float* descriptor, + const float* parameters, + const float* q_scaler, + const double* energy_reference, + double* potential, + float* Fp) +{ + const int atom = blockIdx.x * blockDim.x + threadIdx.x; + if (atom >= N) { + return; + } + const int atom_type = type[atom]; + const std::size_t type_block = static_cast(hidden_neurons) * (descriptor_dim + 2); + const float* w0 = parameters + static_cast(atom_type) * type_block; + const float* b0 = w0 + static_cast(hidden_neurons) * descriptor_dim; + const float* w1 = b0 + hidden_neurons; + const float* b1 = parameters + static_cast(num_types) * type_block; + float q[NEP_Spin_Model_File::MAX_DIM_SPIN] = {0.0f}; + for (int d = 0; d < descriptor_dim; ++d) { + q[d] = descriptor[static_cast(N) * d + atom] * q_scaler[d]; + } + float energy = 0.0f; + float fp[NEP_Spin_Model_File::MAX_DIM_SPIN] = {0.0f}; + for (int neuron = 0; neuron < hidden_neurons; ++neuron) { + float projection = 0.0f; + for (int d = 0; d < descriptor_dim; ++d) { + projection += w0[neuron * descriptor_dim + d] * q[d]; + } + const float activation = tanhf(projection - b0[neuron]); + const float activation_derivative = 1.0f - activation * activation; + energy += w1[neuron] * activation; + for (int d = 0; d < descriptor_dim; ++d) { + fp[d] += w1[neuron] * activation_derivative * w0[neuron * descriptor_dim + d]; + } + } + energy -= b1[0]; + potential[atom] += static_cast(energy) + static_cast(energy_reference[atom_type]); + for (int d = 0; d < descriptor_dim; ++d) { + Fp[static_cast(N) * d + atom] = fp[d] * q_scaler[d]; + } +} + +template +__global__ void find_potential_warp( + const int N, + const int descriptor_dim, + const int hidden_neurons, + const int num_types, + const int* __restrict__ type, + const float* __restrict__ descriptor, + const float* __restrict__ parameters, + const float* __restrict__ q_scaler, + const double* __restrict__ energy_reference, + double* potential, + float* __restrict__ Fp) +{ + static_assert(32 % lanes_per_atom == 0, "one warp must contain complete atom tiles"); + constexpr int warps_per_block = 4; + constexpr int atoms_per_warp = 32 / lanes_per_atom; + constexpr int atoms_per_block = warps_per_block * atoms_per_warp; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int atom_lane = lane / lanes_per_atom; + const int work_lane = lane - atom_lane * lanes_per_atom; + const int atom_local = warp * atoms_per_warp + atom_lane; + const int atom = blockIdx.x * atoms_per_block + atom_local; + const unsigned int atom_mask = ((1u << lanes_per_atom) - 1u) << (atom_lane * lanes_per_atom); + + extern __shared__ float state[]; + float* q = state + static_cast(atom_local) * descriptor_dim; + float* hidden_delta = state + static_cast(atoms_per_block) * descriptor_dim + + static_cast(atom_local) * hidden_neurons; + + if (atom < N) { + for (int d = work_lane; d < descriptor_dim; d += lanes_per_atom) { + q[d] = descriptor[static_cast(N) * d + atom] * q_scaler[d]; + } + } + __syncwarp(atom_mask); + if (atom >= N) { + return; + } + + const int atom_type = type[atom]; + const std::size_t type_block = static_cast(hidden_neurons) * (descriptor_dim + 2); + const float* w0 = parameters + static_cast(atom_type) * type_block; + const float* b0 = w0 + static_cast(hidden_neurons) * descriptor_dim; + const float* w1 = b0 + hidden_neurons; + const float* b1 = parameters + static_cast(num_types) * type_block; + + float energy = 0.0f; + for (int neuron = work_lane; neuron < hidden_neurons; neuron += lanes_per_atom) { + float projection = 0.0f; + for (int d = 0; d < descriptor_dim; ++d) { + projection += w0[neuron * descriptor_dim + d] * q[d]; + } + const float activation = tanhf(projection - b0[neuron]); + const float delta = w1[neuron] * (1.0f - activation * activation); + hidden_delta[neuron] = delta; + energy += w1[neuron] * activation; + } + for (int offset = lanes_per_atom / 2; offset > 0; offset >>= 1) { + energy += __shfl_down_sync(atom_mask, energy, offset, lanes_per_atom); + } + __syncwarp(atom_mask); + + if (work_lane == 0) { + energy -= b1[0]; + potential[atom] += + static_cast(energy) + static_cast(energy_reference[atom_type]); + } + for (int d = work_lane; d < descriptor_dim; d += lanes_per_atom) { + float fp = 0.0f; + for (int neuron = 0; neuron < hidden_neurons; ++neuron) { + fp += hidden_delta[neuron] * w0[neuron * descriptor_dim + d]; + } + Fp[static_cast(N) * d + atom] = fp * q_scaler[d]; + } +} + +template +__global__ void find_force_angular_warp( + const NEP::ParaMB paramb, + const NEP::ANN ann, + const int N, + const Box box, + const int* NN_angular, + const int* NL_angular, + const int* __restrict__ type, + const double* __restrict__ x, + const double* __restrict__ y, + const double* __restrict__ z, + const float* __restrict__ Fp, + const float* __restrict__ sum_fxyz, + double* force, + double* virial) +{ + static_assert(atoms_per_warp * edge_lanes == 32, "one tile must fill one warp"); + constexpr int warps_per_block = 4; + constexpr int centers_per_block = warps_per_block * atoms_per_warp; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int atom_lane = lane / edge_lanes; + const int edge_lane = lane - atom_lane * edge_lanes; + const unsigned int center_mask = (0xffffffffu >> (32 - edge_lanes)) << (atom_lane * edge_lanes); + const int center_local = warp * atoms_per_warp + atom_lane; + const int n1 = blockIdx.x * centers_per_block + center_local; + + const int abc_count = (paramb.L_max + 1) * (paramb.L_max + 1) - 1; + const int sum_count = (paramb.n_max_angular + 1) * NUM_OF_ABC; + const int state_size = paramb.dim_angular + sum_count; + extern __shared__ float center_state[]; + float* center_Fp = center_state + center_local * state_size; + float* center_sum_fxyz = center_Fp + paramb.dim_angular; + + if (n1 < N) { + for (int d = edge_lane; d < paramb.dim_angular; d += edge_lanes) { + center_Fp[d] = Fp[static_cast(N) * (paramb.n_max_radial + 1 + d) + n1]; + } + const int compact_sum_count = (paramb.n_max_angular + 1) * abc_count; + for (int d = edge_lane; d < compact_sum_count; d += edge_lanes) { + const int n = d / abc_count; + const int abc = d - n * abc_count; + center_sum_fxyz[n * NUM_OF_ABC + abc] = sum_fxyz[static_cast(N) * d + n1]; + } + } + __syncwarp(); + + if (n1 >= N) { + return; + } + const int t1 = type[n1]; + const double x1 = x[n1]; + const double y1 = y[n1]; + const double z1 = z[n1]; + float center_force[3] = {}; + for (int i1 = edge_lane; i1 < NN_angular[n1]; i1 += edge_lanes) { + const int index = i1 * N + n1; + const int n2 = NL_angular[index]; + float x12 = x[n2] - x1; + float y12 = y[n2] - y1; + float z12 = z[n2] - z1; + apply_mic(box, x12, y12, z12); + float r12[3] = {x12, y12, z12}; + const float d12 = sqrt(r12[0] * r12[0] + r12[1] * r12[1] + r12[2] * r12[2]); + float f12[3] = {0.0f}; + float fc12; + float fcp12; + const int t2 = type[n2]; + const float rc = (paramb.rc_angular[t1] + paramb.rc_angular[t2]) * 0.5f; + const float rcinv = 1.0f / rc; + find_fc_and_fcp(rc, rcinv, d12, fc12, fcp12); + + float fn12[MAX_NUM_N]; + float fnp12[MAX_NUM_N]; + find_fn_and_fnp(paramb.basis_size_angular, rcinv, d12, fc12, fcp12, fn12, fnp12); + for (int n = 0; n <= paramb.n_max_angular; ++n) { + float gn12 = 0.0f; + float gnp12 = 0.0f; + for (int k = 0; k <= paramb.basis_size_angular; ++k) { + int c_index = paramb.num_c_radial; + c_index += (t1 * paramb.num_types + t2) * + ((paramb.n_max_angular + 1) * (paramb.basis_size_angular + 1)); + c_index += n * (paramb.basis_size_angular + 1) + k; + gn12 += fn12[k] * ann.c_type_pair[c_index]; + gnp12 += fnp12[k] * ann.c_type_pair[c_index]; + } + if (l4_q222_only) { + accumulate_f12( + 4, + true, + false, + false, + false, + false, + false, + 5, + n, + paramb.n_max_angular + 1, + d12, + r12, + gn12, + gnp12, + center_Fp, + center_sum_fxyz, + f12); + } else { + accumulate_f12( + paramb.L_max, + paramb.has_q_222, + paramb.has_q_1111, + paramb.has_q_112, + paramb.has_q_123, + paramb.has_q_233, + paramb.has_q_134, + paramb.num_L, + n, + paramb.n_max_angular + 1, + d12, + r12, + gn12, + gnp12, + center_Fp, + center_sum_fxyz, + f12); + } + } + // One directed response contributes +f12 to its center and -f12 to its + // neighbor. The neighbor-owned virial is (-r12) tensor f12. + for (int d = 0; d < 3; ++d) { + center_force[d] += f12[d]; + atomicAdd(force + d * N + n2, -static_cast(f12[d])); + } + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + const int component = virial_internal_component(3 * a + b); + atomicAdd(virial + component * N + n2, -static_cast(r12[a] * f12[b])); + } + } + } + for (int offset = edge_lanes / 2; offset > 0; offset >>= 1) { + for (int d = 0; d < 3; ++d) { + center_force[d] += __shfl_down_sync(center_mask, center_force[d], offset, edge_lanes); + } + } + if (edge_lane == 0) { + for (int d = 0; d < 3; ++d) { + atomicAdd(force + d * N + n1, static_cast(center_force[d])); + } + } +} + +template +__global__ void find_spin_zbl_force( + const NEP::ParaMB paramb, + const int atom_count, + const NEP::ZBL zbl, + const Box box, + const int* neighbor_count, + const int* neighbor_list, + const int* type, + const double* x, + const double* y, + const double* z, + const double* r12x, + const double* r12y, + const double* r12z, + double* force_x, + double* force_y, + double* force_z, + double* virial, + double* potential) +{ + const int atom1 = blockIdx.x * blockDim.x + threadIdx.x; + if (atom1 >= atom_count) { + return; + } + float energy = 0.0f; + float force_sum[3] = {}; + float virial_sum[9] = {}; + const int type1 = type[atom1]; + const int atomic_number1 = zbl.atomic_numbers[type1]; + const float atomic_number_power1 = powf(static_cast(atomic_number1), 0.23f); + for (int neighbor = 0; neighbor < neighbor_count[atom1]; ++neighbor) { + const int index = atom1 + atom_count * neighbor; + const int atom2 = neighbor_list[index]; + float displacement[3]; + if (SmallBox) { + displacement[0] = static_cast(r12x[index]); + displacement[1] = static_cast(r12y[index]); + displacement[2] = static_cast(r12z[index]); + } else { + displacement[0] = static_cast(x[atom2] - x[atom1]); + displacement[1] = static_cast(y[atom2] - y[atom1]); + displacement[2] = static_cast(z[atom2] - z[atom1]); + apply_mic(box, displacement[0], displacement[1], displacement[2]); + } + const float distance = sqrtf( + displacement[0] * displacement[0] + displacement[1] * displacement[1] + + displacement[2] * displacement[2]); + const float inverse_distance = 1.0f / distance; + const int type2 = type[atom2]; + const int atomic_number2 = zbl.atomic_numbers[type2]; + const float inverse_screening_length = + (atomic_number_power1 + powf(static_cast(atomic_number2), 0.23f)) * 2.134563f; + const float charge_product = K_C_SP * atomic_number1 * atomic_number2; + float pair_energy = 0.0f; + float pair_derivative = 0.0f; + float inner_cutoff = zbl.rc_inner; + float outer_cutoff = zbl.rc_outer; + if (paramb.use_typewise_cutoff_zbl) { + outer_cutoff = min( + (COVALENT_RADIUS[atomic_number1 - 1] + COVALENT_RADIUS[atomic_number2 - 1]) * + paramb.typewise_cutoff_zbl_factor, + outer_cutoff); + inner_cutoff = 0.0f; + } + find_f_and_fp_zbl( + charge_product, + inverse_screening_length, + inner_cutoff, + outer_cutoff, + distance, + inverse_distance, + pair_energy, + pair_derivative); + const float force_scale = 0.5f * pair_derivative * inverse_distance; + const float pair_force[3] = { + displacement[0] * force_scale, displacement[1] * force_scale, displacement[2] * force_scale}; + energy += 0.5f * pair_energy; + if (SmallBox) { + atomicAdd(&force_x[atom1], static_cast(pair_force[0])); + atomicAdd(&force_y[atom1], static_cast(pair_force[1])); + atomicAdd(&force_z[atom1], static_cast(pair_force[2])); + atomicAdd(&force_x[atom2], static_cast(-pair_force[0])); + atomicAdd(&force_y[atom2], static_cast(-pair_force[1])); + atomicAdd(&force_z[atom2], static_cast(-pair_force[2])); + atomicAdd(&virial[atom2], static_cast(-displacement[0] * pair_force[0])); + atomicAdd(&virial[atom_count + atom2], static_cast(-displacement[1] * pair_force[1])); + atomicAdd( + &virial[2 * atom_count + atom2], static_cast(-displacement[2] * pair_force[2])); + atomicAdd( + &virial[3 * atom_count + atom2], static_cast(-displacement[0] * pair_force[1])); + atomicAdd( + &virial[4 * atom_count + atom2], static_cast(-displacement[0] * pair_force[2])); + atomicAdd( + &virial[5 * atom_count + atom2], static_cast(-displacement[1] * pair_force[2])); + atomicAdd( + &virial[6 * atom_count + atom2], static_cast(-displacement[1] * pair_force[0])); + atomicAdd( + &virial[7 * atom_count + atom2], static_cast(-displacement[2] * pair_force[0])); + atomicAdd( + &virial[8 * atom_count + atom2], static_cast(-displacement[2] * pair_force[1])); + } else { + force_sum[0] += 2.0f * pair_force[0]; + force_sum[1] += 2.0f * pair_force[1]; + force_sum[2] += 2.0f * pair_force[2]; + virial_sum[0] -= displacement[0] * pair_force[0]; + virial_sum[1] -= displacement[1] * pair_force[1]; + virial_sum[2] -= displacement[2] * pair_force[2]; + virial_sum[3] -= displacement[0] * pair_force[1]; + virial_sum[4] -= displacement[0] * pair_force[2]; + virial_sum[5] -= displacement[1] * pair_force[2]; + virial_sum[6] -= displacement[1] * pair_force[0]; + virial_sum[7] -= displacement[2] * pair_force[0]; + virial_sum[8] -= displacement[2] * pair_force[1]; + } + } + potential[atom1] += energy; + if (!SmallBox) { + force_x[atom1] += force_sum[0]; + force_y[atom1] += force_sum[1]; + force_z[atom1] += force_sum[2]; +#pragma unroll + for (int component = 0; component < 9; ++component) { + virial[component * atom_count + atom1] += virial_sum[component]; + } + } +} + +} // namespace + +template +void NEP_Spin::launch_force_angular_warp( + const NEP::ParaMB& paramb, + const NEP::ANN& ann, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const Data& data, + GPU_Vector& force, + GPU_Vector& virial, + const int grid_size, + const int block_size, + const std::size_t shared_size) +{ + const int N = type.size(); + find_force_angular_warp + <<>>( + paramb, + ann, + N, + box, + data.NN_angular.data(), + data.NL_angular.data(), + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data.Fp.data(), + data.sum_fxyz.data(), + force.data(), + virial.data()); +} + +template +void NEP_Spin::launch_spin3_descriptors_typed( + const NEP_Spin::Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + const double* slot_r12, + int r12_plane_size, + gpuStream_t stream) +{ + constexpr int block_size = 128; + const int N = type.size(); + const auto layout = make_spin_polynomial_layout( + C, model.spin_l_max[0], model.spin_order, model.spin_soc); + const int density_work = C * N; + build_spin3_oc_density_bank + <<<(density_work + block_size - 1) / block_size, block_size, 0, stream>>>( + layout, + N, + N, + model.struct_descriptor_dim, + model.num_types, + model.spin_basis_size[0], + static_cast(model.spin_cutoff), + data.spin_cutoff_pair.data(), + static_cast(model.radial_parameter_count + model.angular_parameter_count), + box, + type.data(), + data.spin_type_active.data(), + position.data(), + slot_r12, + r12_plane_size, + spin.data(), + data.NN_spin.data(), + data.NL_spin.data(), + data.descriptor_parameters_type_pair.data(), + data.descriptor.data(), + data.spin3_moments.data()); + contract_spin3_oc_descriptors<<<(N + block_size - 1) / block_size, block_size, 0, stream>>>( + layout, + N, + N, + model.struct_descriptor_dim, + type.data(), + data.spin_type_active.data(), + spin.data(), + data.spin_projection_parameters.data(), + data.spin3_moments.data(), + data.descriptor.data()); + GPU_CHECK_KERNEL +} + +void NEP_Spin::launch_spin3_descriptors( + const NEP_Spin::Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + const double* slot_r12, + int r12_plane_size, + gpuStream_t stream) +{ + switch (model.spin_compress) { + case 1: + launch_spin3_descriptors_typed<1>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 2: + launch_spin3_descriptors_typed<2>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 3: + launch_spin3_descriptors_typed<3>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 4: + launch_spin3_descriptors_typed<4>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 5: + launch_spin3_descriptors_typed<5>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 6: + launch_spin3_descriptors_typed<6>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 7: + launch_spin3_descriptors_typed<7>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 8: + launch_spin3_descriptors_typed<8>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + case 9: + launch_spin3_descriptors_typed<9>( + model, box, type, position, spin, data, slot_r12, r12_plane_size, stream); + break; + default: + PRINT_INPUT_ERROR("Unsupported spin_compress in descriptor dispatch."); + } +} + +template +void NEP_Spin::launch_spin3_forces_typed( + const NEP_Spin::Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + GPU_Vector& force, + GPU_Vector& mforce, + GPU_Vector& virial, + const double* slot_r12, + int r12_plane_size) +{ + const auto layout = make_spin_polynomial_layout( + C, model.spin_l_max[0], model.spin_order, model.spin_soc); + using Tile = Spin3ForceTile; + constexpr int block_size = Tile::block_size; + const int N = type.size(); + constexpr int atoms_per_warp = Tile::atoms_per_warp; + constexpr int centers_per_block = Tile::centers_per_block; + const int grid_size = (N - 1) / centers_per_block + 1; + const std::size_t shared_size = + static_cast(centers_per_block) * layout.moment_count * sizeof(float); + accumulate_spin3_oc_native_forces + <<>>( + layout, + N, + N, + model.struct_descriptor_dim, + model.num_types, + model.spin_basis_size[0], + static_cast(model.spin_cutoff), + data.spin_cutoff_pair.data(), + box, + type.data(), + data.spin_type_active.data(), + position.data(), + slot_r12, + r12_plane_size, + spin.data(), + data.NN_spin.data(), + data.NL_spin.data(), + data.Fp.data(), + data.descriptor_parameters_type_pair.data(), + data.spin_projection_parameters.data(), + data.spin3_moments.data(), + static_cast(model.radial_parameter_count + model.angular_parameter_count), + force.data(), + mforce.data(), + virial.data(), + nullptr, + nullptr); + GPU_CHECK_KERNEL +} + +void NEP_Spin::launch_spin3_forces( + const NEP_Spin::Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + GPU_Vector& force, + GPU_Vector& mforce, + GPU_Vector& virial, + const double* slot_r12, + int r12_plane_size) +{ + switch (model.spin_compress) { + case 1: + launch_spin3_forces_typed<1>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 2: + launch_spin3_forces_typed<2>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 3: + launch_spin3_forces_typed<3>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 4: + launch_spin3_forces_typed<4>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 5: + launch_spin3_forces_typed<5>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 6: + launch_spin3_forces_typed<6>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 7: + launch_spin3_forces_typed<7>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 8: + launch_spin3_forces_typed<8>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + case 9: + launch_spin3_forces_typed<9>( + model, box, type, position, spin, data, force, mforce, virial, slot_r12, r12_plane_size); + break; + default: + PRINT_INPUT_ERROR("Unsupported spin_compress in force dispatch."); + } + const bool has_inactive_dof = + std::any_of(model.spin_type_active.begin(), model.spin_type_active.end(), [](const int active) { + return active == 0; + }); + if (has_inactive_dof) { + constexpr int block_size = 128; + const int grid_size = (type.size() - 1) / block_size + 1; + mask_inactive_spin_mforce<<>>( + type.size(), type.size(), type.data(), data.spin_type_active.data(), mforce.data()); + } + GPU_CHECK_KERNEL +} + +int NEP_Spin::Body_Channels::count(void) const +{ + return l_max_3body + static_cast(has_q_222) + static_cast(has_q_1111) + + static_cast(has_q_112) + static_cast(has_q_123) + static_cast(has_q_233) + + static_cast(has_q_134); +} + +NEP_Spin::NEP_Spin(const char* file_potential, const int num_atoms) : num_atoms_(num_atoms) +{ + try { + read_model(file_potential); + } catch (const std::exception& error) { + PRINT_INPUT_ERROR(error.what()); + } + CHECK(gpuStreamCreate(&structural_descriptor_stream_)); + CHECK(gpuStreamCreate(&spin_descriptor_stream_)); + + rc = std::max({model_.cutoff_radial, model_.cutoff_angular, model_.spin_cutoff}); + input_neighbor_capacity_ = model_.neighbor_capacity; + resize_domain(num_atoms_); + std::vector spin_cutoff_pair( + static_cast(model_.num_types) * model_.num_types); + for (int type_i = 0; type_i < model_.num_types; ++type_i) { + for (int type_j = 0; type_j < model_.num_types; ++type_j) { + spin_cutoff_pair[type_i * model_.num_types + type_j] = static_cast( + 0.5 * (model_.spin_cutoff_by_type[type_i] + model_.spin_cutoff_by_type[type_j])); + } + } + data_.spin_cutoff_pair.resize(spin_cutoff_pair.size()); + data_.spin_cutoff_pair.copy_from_host(spin_cutoff_pair.data(), spin_cutoff_pair.size()); + + printf( + "Use NEP-spin potential with %d atom type%s.\n", + model_.num_types, + model_.num_types == 1 ? "" : "s"); + printf( + " structural/spin descriptor dimensions = %d/%d.\n", + model_.struct_descriptor_dim, + model_.spin_descriptor_dim); + printf( + " radial/angular/spin cutoffs = %g/%g/%g A.\n", + model_.cutoff_radial, + model_.cutoff_angular, + model_.spin_cutoff); +} + +void NEP_Spin::resize_domain(int num_atoms) +{ + num_atoms_ = num_atoms; + N1 = 0; + N2 = num_atoms; + small_box_initialized_ = false; + neighbor_.initialize(rc, num_atoms_, input_neighbor_capacity_); + model_.neighbor_capacity = neighbor_.get_capacity(); + const std::size_t slots = + static_cast(num_atoms_) * static_cast(model_.neighbor_capacity); + if (data_.NN_radial.size() != (num_atoms_)) { + data_.NN_radial.resize(num_atoms_); + } + if (data_.NN_angular.size() != (num_atoms_)) { + data_.NN_angular.resize(num_atoms_); + } + if (data_.NN_spin.size() != (num_atoms_)) { + data_.NN_spin.resize(num_atoms_); + } + if (data_.NL_radial.size() != (slots)) { + data_.NL_radial.resize(slots); + } + if (data_.NL_angular.size() != (slots)) { + data_.NL_angular.resize(slots); + } + if (data_.NL_spin.size() != (slots)) { + data_.NL_spin.resize(slots); + } + if (data_.descriptor.size() != (static_cast(num_atoms_) * model_.descriptor_dim)) { + data_.descriptor.resize(static_cast(num_atoms_) * model_.descriptor_dim); + } + if (data_.Fp.size() != (static_cast(num_atoms_) * model_.descriptor_dim)) { + data_.Fp.resize(static_cast(num_atoms_) * model_.descriptor_dim); + } + if (data_.f12x.size() != (slots)) { + data_.f12x.resize(slots); + } + if (data_.f12y.size() != (slots)) { + data_.f12y.resize(slots); + } + if (data_.f12z.size() != (slots)) { + data_.f12z.resize(slots); + } + const std::size_t N = static_cast(num_atoms_); + const std::size_t spin3_state_size = N * static_cast(model_.spin_moment_count); + if (data_.spin3_moments.size() != (spin3_state_size)) { + data_.spin3_moments.resize(spin3_state_size); + } + const std::size_t sum_fxyz_size = + N * static_cast(model_.n_max_angular + 1) * + static_cast((model_.body.l_max_3body + 1) * (model_.body.l_max_3body + 1) - 1); + if (data_.sum_fxyz.size() != (std::max(sum_fxyz_size, 1))) { + data_.sum_fxyz.resize(std::max(sum_fxyz_size, 1)); + } +} + +NEP_Spin::~NEP_Spin(void) +{ + if (structural_descriptor_stream_ != nullptr) { + CHECK(gpuStreamDestroy(structural_descriptor_stream_)); + } + if (spin_descriptor_stream_ != nullptr) { + CHECK(gpuStreamDestroy(spin_descriptor_stream_)); + } +} + +void NEP_Spin::read_model(const char* file_potential) +{ + std::ifstream input(file_potential); + NEP_Spin_Model_File file; + std::string error; + if (!file.read(input, error)) { + throw std::runtime_error(error); + } + model_.elements = file.elements; + model_.num_types = file.elements.size(); + model_.cutoff_radial = file.rc_radial; + model_.cutoff_angular = file.rc_angular; + model_.max_neighbors_global = file.max_radial; + model_.max_neighbors_angular = file.max_angular; + if (file.max_radial <= 0 || file.max_angular <= 0) { + throw std::runtime_error("NEP_Spin neighbor capacities must be positive"); + } + model_.n_max_radial = file.n_max[0]; + model_.n_max_angular = file.n_max[1]; + model_.basis_size_radial = file.basis_size[0]; + model_.basis_size_angular = file.basis_size[1]; + model_.body.l_max_3body = file.l_max[0]; + model_.body.has_q_222 = file.l_max[1] != 0; + model_.body.has_q_1111 = file.l_max[2] != 0; + model_.body.has_q_112 = file.l_max[3] != 0; + model_.body.has_q_123 = file.l_max[4] != 0; + model_.body.has_q_233 = file.l_max[5] != 0; + model_.body.has_q_134 = file.l_max[6] != 0; + model_.hidden_neurons = file.num_neurons; + model_.spin_basis_size[0] = file.spin_basis_size; + model_.spin_l_max[0] = file.spin_l_max; + model_.spin_compress = file.spin_compress; + model_.spin_cutoff_by_type.assign(file.spin_cutoff.begin(), file.spin_cutoff.end()); + model_.spin_cutoff = *std::max_element(file.spin_cutoff.begin(), file.spin_cutoff.end()); + model_.spin_order = file.spin_order; + model_.spin_soc = file.spin_soc; + model_.spin_projection_size = file.spin_projection_size; + model_.spin_type_active = file.spin_type; + model_.energy_reference = file.energy_reference; + model_.struct_descriptor_dim = file.dim_struct; + model_.spin_descriptor_dim = file.dim_spin; + model_.descriptor_dim = file.dim; + model_.spin_moment_count = + make_spin_polynomial_layout( + model_.spin_compress, model_.spin_l_max[0], model_.spin_order, model_.spin_soc) + .moment_count; + + zbl_.enabled = file.enable_zbl; + zbl_.num_types = model_.num_types; + zbl_.rc_inner = file.zbl_rc_inner; + zbl_.rc_outer = file.zbl_rc_outer; + paramb_.use_typewise_cutoff_zbl = file.use_typewise_cutoff_zbl; + paramb_.typewise_cutoff_zbl_factor = file.typewise_cutoff_zbl_factor; + for (int type = 0; type < model_.num_types; ++type) { + for (int element = 0; element < NUM_ELEMENTS; ++element) { + if (model_.elements[type] == kElementSymbols[element]) { + zbl_.atomic_numbers[type] = element + 1; + break; + } + } + } + + const std::size_t types = static_cast(model_.num_types); + const std::size_t type_pairs = types * types; + model_.ann_parameter_count = + static_cast(model_.descriptor_dim + 2) * model_.hidden_neurons * types + 1; + model_.radial_parameter_count = + static_cast(model_.n_max_radial + 1) * (model_.basis_size_radial + 1) * type_pairs; + model_.angular_parameter_count = static_cast(model_.n_max_angular + 1) * + (model_.basis_size_angular + 1) * type_pairs; + model_.spin_parameter_count = + static_cast(model_.spin_compress) * (model_.spin_basis_size[0] + 1) * type_pairs; + model_.spin_projection_parameter_count = static_cast(model_.spin_projection_size); + model_.model_parameter_count = model_.ann_parameter_count + model_.radial_parameter_count + + model_.angular_parameter_count + model_.spin_parameter_count + + model_.spin_projection_parameter_count; + + std::vector parameters = file.parameters; + parameters.insert(parameters.end(), file.q_scaler.begin(), file.q_scaler.end()); + data_.parameters.resize(parameters.size()); + data_.parameters.copy_from_host(parameters.data()); + + const std::size_t descriptor_offset = model_.ann_parameter_count; + const std::size_t ordinary_count = model_.radial_parameter_count + model_.angular_parameter_count; + // Current model blocks already use the runtime structural type-pair order. + std::vector descriptor_parameters( + parameters.begin() + descriptor_offset, + parameters.begin() + descriptor_offset + ordinary_count + model_.spin_parameter_count); + data_.descriptor_parameters_type_pair.resize(descriptor_parameters.size()); + data_.descriptor_parameters_type_pair.copy_from_host(descriptor_parameters.data()); + if (model_.spin_projection_parameter_count > 0) { + const std::size_t projection_offset = + descriptor_offset + ordinary_count + model_.spin_parameter_count; + std::vector projection_parameters( + parameters.begin() + projection_offset, + parameters.begin() + projection_offset + model_.spin_projection_parameter_count); + data_.spin_projection_parameters.resize(projection_parameters.size()); + data_.spin_projection_parameters.copy_from_host(projection_parameters.data()); + } + + const auto& energy_reference = model_.energy_reference; + data_.energy_reference.resize(energy_reference.size()); + data_.energy_reference.copy_from_host(energy_reference.data()); + data_.spin_type_active.resize(model_.spin_type_active.size()); + data_.spin_type_active.copy_from_host(model_.spin_type_active.data()); + + paramb_.num_types = model_.num_types; + paramb_.num_types_sq = model_.num_types * model_.num_types; + paramb_.n_max_radial = model_.n_max_radial; + paramb_.n_max_angular = model_.n_max_angular; + paramb_.basis_size_radial = model_.basis_size_radial; + paramb_.basis_size_angular = model_.basis_size_angular; + paramb_.L_max = model_.body.l_max_3body; + paramb_.num_L = model_.body.count(); + paramb_.has_q_222 = model_.body.has_q_222; + paramb_.has_q_1111 = model_.body.has_q_1111; + paramb_.has_q_112 = model_.body.has_q_112; + paramb_.has_q_123 = model_.body.has_q_123; + paramb_.has_q_233 = model_.body.has_q_233; + paramb_.has_q_134 = model_.body.has_q_134; + paramb_.dim_angular = (model_.n_max_angular + 1) * model_.body.count(); + paramb_.num_c_radial = static_cast(model_.radial_parameter_count); + paramb_.MN_radial = model_.neighbor_capacity; + paramb_.MN_angular = model_.neighbor_capacity; + paramb_.rc_radial_max = model_.cutoff_radial; + paramb_.rc_radial_max_inv = 1.0f / model_.cutoff_radial; + for (int type = 0; type < model_.num_types; ++type) { + paramb_.rc_radial[type] = model_.cutoff_radial; + paramb_.rc_angular[type] = model_.cutoff_angular; + } + + ann_.dim = model_.descriptor_dim; + ann_.num_neurons1 = model_.hidden_neurons; + ann_.num_para_ann = static_cast(model_.ann_parameter_count); + ann_.num_para = static_cast(model_.model_parameter_count); + float* pointer = data_.parameters.data(); + for (int type = 0; type < model_.num_types; ++type) { + ann_.w0[type] = pointer; + pointer += model_.hidden_neurons * model_.descriptor_dim; + ann_.b0[type] = pointer; + pointer += model_.hidden_neurons; + ann_.w1[type] = pointer; + pointer += model_.hidden_neurons; + } + ann_.b1 = pointer; + ann_.c_type_pair = data_.descriptor_parameters_type_pair.data(); + ann_.q_scaler = data_.parameters.data() + model_.model_parameter_count; + + const double enlarged = std::ceil(static_cast(model_.max_neighbors_global) * 1.25); + if (enlarged > std::numeric_limits::max()) { + throw std::runtime_error("neighbor capacity is too large"); + } + model_.neighbor_capacity = static_cast(enlarged); +} + +void NEP_Spin::compute( + Box&, + const GPU_Vector&, + const GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&) +{ + PRINT_INPUT_ERROR("NEP_Spin requires spin inputs."); +} + +void NEP_Spin::initialize_small_box(const Box& box) +{ + const double volume = box.get_volume(); + const double thickness[3] = { + volume / box.get_area(0), volume / box.get_area(1), volume / box.get_area(2)}; + expanded_box_.num_cells[0] = box.pbc_x ? static_cast(std::ceil(2.0 * rc / thickness[0])) : 1; + expanded_box_.num_cells[1] = box.pbc_y ? static_cast(std::ceil(2.0 * rc / thickness[1])) : 1; + expanded_box_.num_cells[2] = box.pbc_z ? static_cast(std::ceil(2.0 * rc / thickness[2])) : 1; + + const std::size_t cell_count = static_cast(expanded_box_.num_cells[0]) * + static_cast(expanded_box_.num_cells[1]) * + static_cast(expanded_box_.num_cells[2]); + if ( + cell_count == 0 || cell_count > (std::numeric_limits::max() - 1) / + static_cast(num_atoms_)) { + PRINT_INPUT_ERROR("NEP_Spin small-box image count overflows size_t."); + } + const std::size_t capacity = static_cast(num_atoms_) * cell_count - 1; + if (capacity > static_cast(std::numeric_limits::max())) { + PRINT_INPUT_ERROR("NEP_Spin small-box neighbor capacity exceeds int."); + } + if (!small_box_initialized_) { + const int estimate = static_cast( + std::min(capacity, static_cast(std::max(1, input_neighbor_capacity_)))); + resize_small_box_storage(std::max(1, estimate)); + } + + expanded_box_.h[0] = box.cpu_h[0] * expanded_box_.num_cells[0]; + expanded_box_.h[3] = box.cpu_h[3] * expanded_box_.num_cells[0]; + expanded_box_.h[6] = box.cpu_h[6] * expanded_box_.num_cells[0]; + expanded_box_.h[1] = box.cpu_h[1] * expanded_box_.num_cells[1]; + expanded_box_.h[4] = box.cpu_h[4] * expanded_box_.num_cells[1]; + expanded_box_.h[7] = box.cpu_h[7] * expanded_box_.num_cells[1]; + expanded_box_.h[2] = box.cpu_h[2] * expanded_box_.num_cells[2]; + expanded_box_.h[5] = box.cpu_h[5] * expanded_box_.num_cells[2]; + expanded_box_.h[8] = box.cpu_h[8] * expanded_box_.num_cells[2]; + expanded_box_.h[9] = + expanded_box_.h[4] * expanded_box_.h[8] - expanded_box_.h[5] * expanded_box_.h[7]; + expanded_box_.h[10] = + expanded_box_.h[2] * expanded_box_.h[7] - expanded_box_.h[1] * expanded_box_.h[8]; + expanded_box_.h[11] = + expanded_box_.h[1] * expanded_box_.h[5] - expanded_box_.h[2] * expanded_box_.h[4]; + expanded_box_.h[12] = + expanded_box_.h[5] * expanded_box_.h[6] - expanded_box_.h[3] * expanded_box_.h[8]; + expanded_box_.h[13] = + expanded_box_.h[0] * expanded_box_.h[8] - expanded_box_.h[2] * expanded_box_.h[6]; + expanded_box_.h[14] = + expanded_box_.h[2] * expanded_box_.h[3] - expanded_box_.h[0] * expanded_box_.h[5]; + expanded_box_.h[15] = + expanded_box_.h[3] * expanded_box_.h[7] - expanded_box_.h[4] * expanded_box_.h[6]; + expanded_box_.h[16] = + expanded_box_.h[1] * expanded_box_.h[6] - expanded_box_.h[0] * expanded_box_.h[7]; + expanded_box_.h[17] = + expanded_box_.h[0] * expanded_box_.h[4] - expanded_box_.h[1] * expanded_box_.h[3]; + const double determinant = + expanded_box_.h[0] * + (expanded_box_.h[4] * expanded_box_.h[8] - expanded_box_.h[5] * expanded_box_.h[7]) + + expanded_box_.h[1] * + (expanded_box_.h[5] * expanded_box_.h[6] - expanded_box_.h[3] * expanded_box_.h[8]) + + expanded_box_.h[2] * + (expanded_box_.h[3] * expanded_box_.h[7] - expanded_box_.h[4] * expanded_box_.h[6]); + for (int index = 9; index < 18; ++index) { + expanded_box_.h[index] /= determinant; + } + + small_box_initialized_ = true; +} + +void NEP_Spin::resize_small_box_storage(int capacity) +{ + const std::size_t atoms = static_cast(num_atoms_); + if ( + capacity <= 0 || atoms > static_cast(std::numeric_limits::max()) / + (3 * static_cast(capacity))) { + PRINT_INPUT_ERROR("NEP_Spin small-box displacement storage exceeds int indexing."); + } + const std::size_t plane_size = atoms * static_cast(capacity); + if (plane_size > std::numeric_limits::max() / (3 * sizeof(double))) { + PRINT_INPUT_ERROR("NEP_Spin small-box displacement storage overflows size_t."); + } + small_box_capacity_ = capacity; + data_.NL_radial.resize(plane_size); + data_.NL_angular.resize(plane_size); + data_.NL_spin.resize(plane_size); + data_.r12_radial.resize(3 * plane_size); + data_.r12_angular.resize(3 * plane_size); + data_.r12_spin.resize(3 * plane_size); + data_.f12x.resize(plane_size); + data_.f12y.resize(plane_size); + data_.f12z.resize(plane_size); + if (data_.required_small_box_capacity.size() != 1) { + data_.required_small_box_capacity.resize(1); + } +} + +void NEP_Spin::compute( + Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + GPU_Vector& potential, + GPU_Vector& force, + GPU_Vector& virial, + GPU_Vector& mforce) +{ + const int N = type.size(); + if ( + N != num_atoms_ || position.size() != static_cast(3 * N) || + spin.size() != static_cast(3 * N)) { + PRINT_INPUT_ERROR("NEP_Spin input sizes do not match the initialized atom count."); + } + if (!(box.pbc_x && box.pbc_y && box.pbc_z)) { + PRINT_INPUT_ERROR("NEP_Spin requires periodic boundary conditions in all three directions."); + } + const double volume = box.get_volume(); + const double thickness_x = volume / box.get_area(0); + const double thickness_y = volume / box.get_area(1); + const double thickness_z = volume / box.get_area(2); + const double minimum_safe_thickness = neighbor_.minimum_safe_box_thickness(rc); + const bool is_small_box = (box.pbc_x && thickness_x < minimum_safe_thickness) || + (box.pbc_y && thickness_y < minimum_safe_thickness) || + (box.pbc_z && thickness_z < minimum_safe_thickness); + + constexpr int block_size = 64; + const int grid_size = (N - 1) / block_size + 1; + if (is_small_box) { + initialize_small_box(box); + while (true) { + data_.required_small_box_capacity.fill(0); + find_neighbor_list_spin_small_box<<>>( + paramb_, + static_cast(model_.spin_cutoff), + N, + box, + expanded_box_, + small_box_capacity_, + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data_.NN_radial.data(), + data_.NL_radial.data(), + data_.NN_angular.data(), + data_.NL_angular.data(), + data_.NN_spin.data(), + data_.NL_spin.data(), + data_.r12_radial.data(), + data_.r12_angular.data(), + data_.r12_spin.data(), + data_.required_small_box_capacity.data()); + GPU_CHECK_KERNEL + int required_capacity = 0; + data_.required_small_box_capacity.copy_to_host(&required_capacity); + if (required_capacity <= small_box_capacity_) { + break; + } + // The bounded fill counted all neighbors; no consumer runs until it fits. + resize_small_box_storage(required_capacity); + } + const int plane_size = N * small_box_capacity_; + find_structural_descriptor<<>>( + paramb_, + ann_, + N, + box, + data_.NN_radial.data(), + data_.NL_radial.data(), + data_.NN_angular.data(), + data_.NL_angular.data(), + data_.r12_radial.data(), + data_.r12_angular.data(), + plane_size, + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data_.descriptor.data(), + data_.sum_fxyz.data()); + GPU_CHECK_KERNEL + launch_spin3_descriptors( + model_, box, type, position, spin, data_, data_.r12_spin.data(), plane_size); + find_potential<<>>( + N, + model_.descriptor_dim, + model_.hidden_neurons, + model_.num_types, + type.data(), + data_.descriptor.data(), + data_.parameters.data(), + data_.parameters.data() + model_.model_parameter_count, + data_.energy_reference.data(), + potential.data(), + data_.Fp.data()); + GPU_CHECK_KERNEL + spin_find_force_radial_small_box<<>>( + paramb_, + ann_, + N, + 0, + N, + data_.NN_radial.data(), + data_.NL_radial.data(), + type.data(), + data_.r12_radial.data(), + data_.r12_radial.data() + plane_size, + data_.r12_radial.data() + 2 * plane_size, + data_.Fp.data(), + force.data(), + force.data() + N, + force.data() + 2 * N, + virial.data()); + GPU_CHECK_KERNEL + spin_find_force_angular_small_box<<>>( + paramb_, + ann_, + N, + 0, + N, + data_.NN_angular.data(), + data_.NL_angular.data(), + type.data(), + data_.r12_angular.data(), + data_.r12_angular.data() + plane_size, + data_.r12_angular.data() + 2 * plane_size, + data_.Fp.data(), + data_.sum_fxyz.data(), + force.data(), + force.data() + N, + force.data() + 2 * N, + virial.data()); + GPU_CHECK_KERNEL + launch_spin3_forces( + model_, + box, + type, + position, + spin, + data_, + force, + mforce, + virial, + data_.r12_spin.data(), + plane_size); + if (zbl_.enabled) { + find_spin_zbl_force<<>>( + paramb_, + N, + zbl_, + box, + data_.NN_radial.data(), + data_.NL_radial.data(), + type.data(), + nullptr, + nullptr, + nullptr, + data_.r12_radial.data(), + data_.r12_radial.data() + plane_size, + data_.r12_radial.data() + 2 * plane_size, + force.data(), + force.data() + N, + force.data() + 2 * N, + virial.data(), + potential.data()); + GPU_CHECK_KERNEL + } + return; + } + + if (small_box_initialized_) { + resize_domain(N); + } + neighbor_.find_neighbor_global(rc, box, type, position); + constexpr int neighbor_block_size = 128; + const int neighbor_grid_size = (N - 1) / neighbor_block_size + 1; + find_local_neighbors_spin<<>>( + N, + box, + static_cast(model_.cutoff_radial * model_.cutoff_radial), + static_cast(model_.cutoff_angular * model_.cutoff_angular), + static_cast(model_.spin_cutoff * model_.spin_cutoff), + position.data(), + position.data() + N, + position.data() + 2 * N, + neighbor_.NN.data(), + neighbor_.NL.data(), + data_.NN_radial.data(), + data_.NL_radial.data(), + data_.NN_angular.data(), + data_.NL_angular.data(), + data_.NN_spin.data(), + data_.NL_spin.data()); + GPU_CHECK_KERNEL + + find_structural_descriptor<<>>( + paramb_, + ann_, + N, + box, + data_.NN_radial.data(), + data_.NL_radial.data(), + data_.NN_angular.data(), + data_.NL_angular.data(), + nullptr, + nullptr, + 0, + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data_.descriptor.data(), + data_.sum_fxyz.data()); + GPU_CHECK_KERNEL + + launch_spin3_descriptors( + model_, box, type, position, spin, data_, nullptr, 0, spin_descriptor_stream_); +#ifdef USE_HIP + CHECK(hipStreamSynchronize(structural_descriptor_stream_)); + CHECK(hipStreamSynchronize(spin_descriptor_stream_)); +#else + CHECK(cudaStreamSynchronize(structural_descriptor_stream_)); + CHECK(cudaStreamSynchronize(spin_descriptor_stream_)); +#endif + constexpr int potential_lanes_per_atom = 4; + constexpr int potential_block_size = 128; + constexpr int potential_atoms_per_block = potential_block_size / potential_lanes_per_atom; + const int potential_grid_size = (N - 1) / potential_atoms_per_block + 1; + const std::size_t potential_shared_size = static_cast(potential_atoms_per_block) * + (model_.descriptor_dim + model_.hidden_neurons) * + sizeof(float); + constexpr std::size_t maximum_potential_shared_size = 48 * 1024; + if (potential_shared_size <= maximum_potential_shared_size) { + find_potential_warp + <<>>( + N, + model_.descriptor_dim, + model_.hidden_neurons, + model_.num_types, + type.data(), + data_.descriptor.data(), + data_.parameters.data(), + data_.parameters.data() + model_.model_parameter_count, + data_.energy_reference.data(), + potential.data(), + data_.Fp.data()); + } else { + find_potential<<>>( + N, + model_.descriptor_dim, + model_.hidden_neurons, + model_.num_types, + type.data(), + data_.descriptor.data(), + data_.parameters.data(), + data_.parameters.data() + model_.model_parameter_count, + data_.energy_reference.data(), + potential.data(), + data_.Fp.data()); + } + GPU_CHECK_KERNEL + + nep_internal::launch_find_force_radial( + paramb_, + ann_, + N, + 0, + N, + box, + data_.NN_radial.data(), + data_.NL_radial.data(), + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data_.Fp.data(), + force.data(), + force.data() + N, + force.data() + 2 * N, + virial.data()); + GPU_CHECK_KERNEL + + constexpr int angular_atoms_per_warp = 4; + constexpr int angular_edge_lanes = 8; + constexpr int angular_block_size = 128; + constexpr int angular_centers_per_block = angular_block_size / 32 * angular_atoms_per_warp; + const int angular_grid_size = (N - 1) / angular_centers_per_block + 1; + const int angular_state_size = paramb_.dim_angular + (paramb_.n_max_angular + 1) * NUM_OF_ABC; + const std::size_t angular_shared_size = + static_cast(angular_centers_per_block) * angular_state_size * sizeof(float); + constexpr std::size_t maximum_angular_shared_size = 48 * 1024; + const bool use_direct_angular_force = angular_shared_size <= maximum_angular_shared_size; + if (use_direct_angular_force) { + // The body-channel signature comes from l_max in the model input. Keep a + // generic instance for every signature other than the common L4+q222 case. + const bool l4_q222_only = paramb_.L_max == 4 && paramb_.has_q_222 && !paramb_.has_q_1111 && + !paramb_.has_q_112 && !paramb_.has_q_123 && !paramb_.has_q_233 && + !paramb_.has_q_134 && paramb_.num_L == 5; + if (l4_q222_only) { + launch_force_angular_warp( + paramb_, + ann_, + box, + type, + position, + data_, + force, + virial, + angular_grid_size, + angular_block_size, + angular_shared_size); + } else { + launch_force_angular_warp( + paramb_, + ann_, + box, + type, + position, + data_, + force, + virial, + angular_grid_size, + angular_block_size, + angular_shared_size); + } + } else { + nep_internal::launch_find_partial_force_angular( + paramb_, + ann_, + N, + 0, + N, + box, + data_.NN_angular.data(), + data_.NL_angular.data(), + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + data_.Fp.data(), + data_.sum_fxyz.data(), + data_.f12x.data(), + data_.f12y.data(), + data_.f12z.data()); + } + GPU_CHECK_KERNEL + if (!use_direct_angular_force) { + find_properties_many_body( + box, + data_.NN_angular.data(), + data_.NL_angular.data(), + data_.f12x.data(), + data_.f12y.data(), + data_.f12z.data(), + position, + force, + virial); + } + + launch_spin3_forces(model_, box, type, position, spin, data_, force, mforce, virial); + if (zbl_.enabled) { + find_spin_zbl_force<<>>( + paramb_, + N, + zbl_, + box, + neighbor_.NN.data(), + neighbor_.NL.data(), + type.data(), + position.data(), + position.data() + N, + position.data() + 2 * N, + nullptr, + nullptr, + nullptr, + force.data(), + force.data() + N, + force.data() + 2 * N, + virial.data(), + potential.data()); + GPU_CHECK_KERNEL + } +} diff --git a/src/force/nep_spin/nep_spin.cuh b/src/force/nep_spin/nep_spin.cuh new file mode 100644 index 000000000..791b6044e --- /dev/null +++ b/src/force/nep_spin/nep_spin.cuh @@ -0,0 +1,225 @@ +/* + Copyright 2017 Zheyong Fan and GPUMD development team + Copyright 2026 NEPAdapters contributors + This file is part of GPUMD and is distributed under LGPLv3 or later. +*/ + +#pragma once + +#include "force/neighbor.cuh" +#include "force/nep.cuh" +#include "force/potential.cuh" +#include "utilities/common.cuh" +#include "utilities/gpu_vector.cuh" +#include +#include +#include + +class NEP_Spin : public Potential +{ +private: + struct Data { + GPU_Vector parameters; + GPU_Vector descriptor_parameters_type_pair; + GPU_Vector descriptor; + GPU_Vector Fp; + GPU_Vector sum_fxyz; + GPU_Vector f12x; + GPU_Vector f12y; + GPU_Vector f12z; + GPU_Vector energy_reference; + GPU_Vector spin_type_active; + GPU_Vector spin_cutoff_pair; + GPU_Vector spin_projection_parameters; + GPU_Vector spin3_moments; + GPU_Vector NN_radial; + GPU_Vector NL_radial; + GPU_Vector NN_angular; + GPU_Vector NL_angular; + GPU_Vector NN_spin; + GPU_Vector NL_spin; + GPU_Vector r12_radial; + GPU_Vector r12_angular; + GPU_Vector r12_spin; + GPU_Vector required_small_box_capacity; + }; + + struct Body_Channels { + int l_max_3body = 0; + bool has_q_222 = false; + bool has_q_1111 = false; + bool has_q_112 = false; + bool has_q_123 = false; + bool has_q_233 = false; + bool has_q_134 = false; + + int count(void) const; + }; + + struct Model { + int num_types = 0; + int n_max_radial = 0; + int n_max_angular = 0; + int basis_size_radial = 0; + int basis_size_angular = 0; + int hidden_neurons = 0; + int max_neighbors_global = 0; + int max_neighbors_angular = 0; + int neighbor_capacity = 0; + int spin_basis_size[2] = {0, 0}; + int spin_l_max[3] = {0, 0, 0}; + int spin_compress = 0; + int spin_order = 0; + int spin_soc = 0; + int spin_projection_size = 0; + double cutoff_radial = 0.0; + double cutoff_angular = 0.0; + double spin_cutoff = 0.0; + int struct_descriptor_dim = 0; + int spin_descriptor_dim = 0; + int descriptor_dim = 0; + std::size_t ann_parameter_count = 0; + std::size_t radial_parameter_count = 0; + std::size_t angular_parameter_count = 0; + std::size_t spin_parameter_count = 0; + std::size_t spin_projection_parameter_count = 0; + std::size_t model_parameter_count = 0; + std::vector elements; + std::vector energy_reference; + std::vector spin_type_active; + std::vector spin_cutoff_by_type; + Body_Channels body; + int spin_moment_count = 0; + }; + +public: + NEP_Spin(const char* file_potential, const int num_atoms); + virtual ~NEP_Spin(void); + + using Potential::compute; + + void compute( + Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + GPU_Vector& potential, + GPU_Vector& force, + GPU_Vector& virial) override; + + void compute( + Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + GPU_Vector& potential, + GPU_Vector& force, + GPU_Vector& virial, + GPU_Vector& mforce) override; + + void resize_domain(int num_atoms); + void invalidate_neighbors() { neighbor_.invalidate(); } + + const std::vector& spin_types() const { return model_.spin_type_active; } + +private: + Model model_; + Data data_; + class Neighbors + { + public: + GPU_Vector NN, NL; + void initialize(double rc, int atoms, int capacity); + void invalidate() { valid_ = false; } + int get_capacity() const { return capacity_; } + double minimum_safe_box_thickness(double rc) const { return 2.5 * (rc + skin_); } + void find_neighbor_global( + double rc, Box& box, const GPU_Vector& type, const GPU_Vector& position); + + private: + static constexpr double skin_ = 1.0; + int capacity_ = 0; + bool valid_ = false; + double box_[9] = {}; + GPU_Vector cell_count_, cell_sum_, cell_contents_, flag_; + GPU_Vector reference_; + }; + Neighbors neighbor_; + NEP::ParaMB paramb_; + NEP::ANN ann_; + NEP::ZBL zbl_; + int num_atoms_ = 0; + int input_neighbor_capacity_ = 0; + gpuStream_t structural_descriptor_stream_ = nullptr; + gpuStream_t spin_descriptor_stream_ = nullptr; + bool small_box_initialized_ = false; + int small_box_capacity_ = 0; + NEP::ExpandedBox expanded_box_; + + template + void launch_force_angular_warp( + const NEP::ParaMB& paramb, + const NEP::ANN& ann, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const Data& data, + GPU_Vector& force, + GPU_Vector& virial, + const int grid_size, + const int block_size, + const std::size_t shared_size); + + template + void launch_spin3_descriptors_typed( + const Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + const double* slot_r12 = nullptr, + int r12_plane_size = 0, + gpuStream_t stream = nullptr); + + void launch_spin3_descriptors( + const Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + const double* slot_r12 = nullptr, + int r12_plane_size = 0, + gpuStream_t stream = nullptr); + + template + void launch_spin3_forces_typed( + const Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + GPU_Vector& force, + GPU_Vector& mforce, + GPU_Vector& virial, + const double* slot_r12 = nullptr, + int r12_plane_size = 0); + + void launch_spin3_forces( + const Model& model, + const Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + Data& data, + GPU_Vector& force, + GPU_Vector& mforce, + GPU_Vector& virial, + const double* slot_r12 = nullptr, + int r12_plane_size = 0); + + void read_model(const char* file_potential); + void initialize_small_box(const Box& box); + void resize_small_box_storage(int capacity); +}; diff --git a/src/force/nep_spin/nep_spin3_descriptor.cuh b/src/force/nep_spin/nep_spin3_descriptor.cuh index d7066e8e5..7996bdc5e 100644 --- a/src/force/nep_spin/nep_spin3_descriptor.cuh +++ b/src/force/nep_spin/nep_spin3_descriptor.cuh @@ -243,9 +243,9 @@ __device__ __forceinline__ float spin3_oc_l22_scalar_from_canonical( return value; } -template +template __global__ void build_spin3_oc_density_bank( - SpinPolynomialLayout layout, + Layout layout, int atom_count, int atom_stride, int struct_dim, @@ -385,9 +385,9 @@ __global__ void build_spin3_oc_density_bank( #undef NEP_SPIN3_OC_STORE_DYNAMIC } -template +template __global__ void contract_spin3_oc_descriptors( - SpinPolynomialLayout layout, + Layout layout, int atom_count, int atom_stride, int struct_dim, diff --git a/src/force/nep_spin/nep_spin3_force.cuh b/src/force/nep_spin/nep_spin3_force.cuh index 740cdc129..6feec0e5a 100644 --- a/src/force/nep_spin/nep_spin3_force.cuh +++ b/src/force/nep_spin/nep_spin3_force.cuh @@ -184,9 +184,9 @@ __device__ __forceinline__ void spin3_oc_reverse_l22_scalar_from_canonical( spin3_oc_add_full_matrix_pull_to_stf5(q_matrix_pull, q_pull); } -template +template __device__ __noinline__ void build_spin3_oc_center_pull_row( - SpinPolynomialLayout layout, + Layout layout, int atom, int row, unsigned active_mask, @@ -516,9 +516,14 @@ struct Spin3ForceTile { static constexpr int centers_per_block = block_size / 32 * atoms_per_warp; }; -template +template < + int C, + SpinVirialMode VirialMode, + bool AccumulateSpinTransfer, + int AtomsPerWarp, + typename Layout> __global__ void accumulate_spin3_oc_native_forces( - SpinPolynomialLayout layout, + Layout layout, int atom_count, int atom_stride, int struct_dim, diff --git a/src/force/nep_spin/nep_spin_multigpu.cu b/src/force/nep_spin/nep_spin_multigpu.cu new file mode 100644 index 000000000..0e7fd47c0 --- /dev/null +++ b/src/force/nep_spin/nep_spin_multigpu.cu @@ -0,0 +1,557 @@ +/* Single-process spatial decomposition for NEP-Spin. LGPL-3.0-or-later. */ +#include "model/box.cuh" +#include "nep_spin_multigpu.cuh" +#include "utilities/gpu_vector.cuh" +#include +#include +#include +#include +#include + +namespace +{ +static __device__ void find_cell_id( + const int partition_direction, + const Box& box, + const double x, + const double y, + const double z, + const double rc_inv, + const int nx, + const int ny, + const int nz, + int& cell_id_x, + int& cell_id_y, + int& cell_id_z, + int& cell_id) +{ + const double sx = box.cpu_h[9] * x + box.cpu_h[10] * y + box.cpu_h[11] * z; + const double sy = box.cpu_h[12] * x + box.cpu_h[13] * y + box.cpu_h[14] * z; + const double sz = box.cpu_h[15] * x + box.cpu_h[16] * y + box.cpu_h[17] * z; + cell_id_x = floor(sx * box.thickness_x * rc_inv); + cell_id_y = floor(sy * box.thickness_y * rc_inv); + cell_id_z = floor(sz * box.thickness_z * rc_inv); + + while (cell_id_x < 0) + cell_id_x += nx; + while (cell_id_x >= nx) + cell_id_x -= nx; + while (cell_id_y < 0) + cell_id_y += ny; + while (cell_id_y >= ny) + cell_id_y -= ny; + while (cell_id_z < 0) + cell_id_z += nz; + while (cell_id_z >= nz) + cell_id_z -= nz; + if (partition_direction == 0) { + cell_id = cell_id_y + ny * (cell_id_z + nz * cell_id_x); + } else if (partition_direction == 1) { + cell_id = cell_id_x + nx * (cell_id_z + nz * cell_id_y); + } else { + cell_id = cell_id_x + nx * (cell_id_y + ny * cell_id_z); + } +} + +static __device__ void find_cell_id( + const int partition_direction, + const Box& box, + const double x, + const double y, + const double z, + const double rc_inv, + const int nx, + const int ny, + const int nz, + int& cell_id) +{ + int cell_id_x, cell_id_y, cell_id_z; + find_cell_id( + partition_direction, + box, + x, + y, + z, + rc_inv, + nx, + ny, + nz, + cell_id_x, + cell_id_y, + cell_id_z, + cell_id); +} + +static __global__ void find_cell_counts( + const int partition_direction, + const Box box, + const int N, + int* cell_count, + const double* x, + const double* y, + const double* z, + const int nx, + const int ny, + const int nz, + const double rc_inv) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x; + if (n1 < N) { + int cell_id; + find_cell_id(partition_direction, box, x[n1], y[n1], z[n1], rc_inv, nx, ny, nz, cell_id); + atomicAdd(&cell_count[cell_id], 1); + } +} + +static __global__ void find_cell_contents( + const int partition_direction, + const Box box, + const int N, + int* cell_count, + const int* cell_count_sum, + int* cell_contents, + const double* x, + const double* y, + const double* z, + const int nx, + const int ny, + const int nz, + const double rc_inv) +{ + const int n1 = blockIdx.x * blockDim.x + threadIdx.x; + if (n1 < N) { + int cell_id; + find_cell_id(partition_direction, box, x[n1], y[n1], z[n1], rc_inv, nx, ny, nz, cell_id); + const int ind = atomicAdd(&cell_count[cell_id], 1); + cell_contents[cell_count_sum[cell_id] + ind] = n1; + } +} + +static void __global__ set_to_zero(int size, int* data) +{ + int n = threadIdx.x + blockIdx.x * blockDim.x; + if (n < size) { + data[n] = 0; + } +} + +static void find_cell_list( + gpuStream_t& stream, + const int partition_direction, + const double rc, + const int* num_bins, + Box& box, + const int N, + const GPU_Vector& position_per_atom, + GPU_Vector& cell_count, + GPU_Vector& cell_count_sum, + GPU_Vector& cell_contents) +{ + const int offset = position_per_atom.size() / 3; + const int block_size = 256; + const int grid_size = (N - 1) / block_size + 1; + const double rc_inv = 1.0 / rc; + const double* x = position_per_atom.data(); + const double* y = position_per_atom.data() + offset; + const double* z = position_per_atom.data() + offset * 2; + const int N_cells = num_bins[0] * num_bins[1] * num_bins[2]; + + // number of cells is allowed to be larger than the number of atoms + if (N_cells > cell_count.size()) { + cell_count.resize(N_cells); + cell_count_sum.resize(N_cells); + } + + set_to_zero<<<(cell_count.size() - 1) / 64 + 1, 64, 0, stream>>>( + cell_count.size(), cell_count.data()); + GPU_CHECK_KERNEL + + set_to_zero<<<(cell_count_sum.size() - 1) / 64 + 1, 64, 0, stream>>>( + cell_count_sum.size(), cell_count_sum.data()); + GPU_CHECK_KERNEL + + set_to_zero<<<(cell_contents.size() - 1) / 64 + 1, 64, 0, stream>>>( + cell_contents.size(), cell_contents.data()); + GPU_CHECK_KERNEL + + find_cell_counts<<>>( + partition_direction, + box, + N, + cell_count.data(), + x, + y, + z, + num_bins[0], + num_bins[1], + num_bins[2], + rc_inv); + GPU_CHECK_KERNEL + + thrust::exclusive_scan( +#ifdef USE_HIP + thrust::hip::par.on(stream), +#else + thrust::cuda::par.on(stream), +#endif + cell_count.data(), + cell_count.data() + N_cells, + cell_count_sum.data()); + + set_to_zero<<<(cell_count.size() - 1) / 64 + 1, 64, 0, stream>>>( + cell_count.size(), cell_count.data()); + GPU_CHECK_KERNEL + + find_cell_contents<<>>( + partition_direction, + box, + N, + cell_count.data(), + cell_count_sum.data(), + cell_contents.data(), + x, + y, + z, + num_bins[0], + num_bins[1], + num_bins[2], + rc_inv); + GPU_CHECK_KERNEL +} + +// The sorted index map stays on GPU 0. Geometry and spins always travel together, +// including repeated force evaluations at unchanged positions. +__global__ void pack_domain( + int n, + int local, + int start, + const int* order, + const int* type, + const double* position, + const double* spin, + int* local_type, + double* local_position, + double* local_spin) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= local) + return; + int id = order[(start + i) % n]; + local_type[i] = type[id]; + for (int d = 0; d < 3; ++d) { + local_position[i + size_t(d) * local] = position[id + size_t(d) * n]; + local_spin[i + size_t(d) * local] = spin[id + size_t(d) * n]; + } +} + +// Ownership may stay fixed while atoms move through slab faces. The 2.5*rc +// halo contains every two-hop environment while each atom moves < 0.2*rc: +// an initial separation can grow to at most 2*rc + 2*(0.2*rc) < 2.5*rc. +__global__ void check_partition_displacement( + int n, + Box box, + double limit_squared, + const double* position, + const double* reference, + int* changed) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + double dx = position[i] - reference[i]; + double dy = position[i + n] - reference[i + n]; + double dz = position[i + 2 * n] - reference[i + 2 * n]; + apply_mic(box, dx, dy, dz); + if (dx * dx + dy * dy + dz * dz >= limit_squared) + atomicExch(changed, 1); +} + +__global__ void collect_domain( + int n, + int count, + int first, + const int* order, + const double* packed, + double* energy, + double* force, + double* virial, + double* mforce) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= count) + return; + int id = order[first + i]; + energy[id] += packed[i]; + for (int d = 0; d < 3; ++d) { + force[id + size_t(d) * n] += packed[i + size_t(1 + d) * count]; + mforce[id + size_t(d) * n] += packed[i + size_t(4 + d) * count]; + } + for (int d = 0; d < 9; ++d) + virial[id + size_t(d) * n] += packed[i + size_t(7 + d) * count]; +} + +void transfer(void* dst, int dst_gpu, const void* src, int src_gpu, size_t bytes) +{ + if (!bytes) + return; + if (dst_gpu == src_gpu) { + CHECK(gpuMemcpy(dst, src, bytes, gpuMemcpyDeviceToDevice)); + } else { + // The runtime stages through host memory if direct peer transfer is unavailable. +#ifdef USE_HIP + CHECK(hipMemcpyPeer(dst, dst_gpu, src, src_gpu, bytes)); +#else + CHECK(cudaMemcpyPeer(dst, dst_gpu, src, src_gpu, bytes)); +#endif + } +} +} // namespace + +void NEP_Spin_MultiGPU::Domain::resize(int count) +{ + if (type.size() == size_t(count)) + return; + if (type.size() != (count)) { + type.resize(count); + } + if (position.size() != (size_t(count) * 3)) { + position.resize(size_t(count) * 3); + } + if (spin.size() != (size_t(count) * 3)) { + spin.resize(size_t(count) * 3); + } + if (energy.size() != (count)) { + energy.resize(count); + } + if (force.size() != (size_t(count) * 3)) { + force.resize(size_t(count) * 3); + } + if (mforce.size() != (size_t(count) * 3)) { + mforce.resize(size_t(count) * 3); + } + if (virial.size() != (size_t(count) * 9)) { + virial.resize(size_t(count) * 9); + } + potential->resize_domain(count); +} + +NEP_Spin_MultiGPU::NEP_Spin_MultiGPU(int devices, const char* model, int atoms, int direction) + : devices_(devices), atoms_(atoms), direction_(direction) +{ + if (devices < 2 || atoms <= 0 || direction < -1 || direction > 2) + PRINT_INPUT_ERROR("Invalid NEP-Spin multi-GPU partition."); + for (int gpu = 0; gpu < devices_; ++gpu) { + CHECK(gpuSetDevice(gpu)); + domains_.emplace_back(new Domain); + domains_.back()->potential.reset(new NEP_Spin(model, 1)); + } + rc = domains_[0]->potential->rc; + CHECK(gpuSetDevice(0)); + CHECK(gpuStreamCreate(&partition_stream_)); + order_.resize(atoms_); + partition_position_.resize(size_t(atoms_) * 3); + partition_changed_.resize(1); + cell_count_.resize(1); + cell_sum_.resize(1); + printf( + "NEP-Spin: spatial decomposition over %d GPUs; two-cutoff geometry/spin halos.\n", devices_); +} + +NEP_Spin_MultiGPU::~NEP_Spin_MultiGPU() +{ + for (int gpu = 0; gpu < devices_; ++gpu) { + CHECK(gpuSetDevice(gpu)); + domains_[gpu].reset(); + } + CHECK(gpuSetDevice(0)); + CHECK(gpuStreamDestroy(partition_stream_)); +} + +void NEP_Spin_MultiGPU::compute( + Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + GPU_Vector& energy, + GPU_Vector& force, + GPU_Vector& virial, + GPU_Vector& mforce) +{ + CHECK(gpuSetDevice(0)); + const int n = type.size(); + if (n != atoms_ || position.size() != size_t(n) * 3 || spin.size() != size_t(n) * 3) + PRINT_INPUT_ERROR("NEP-Spin multi-GPU atom count mismatch."); + if (!(box.pbc_x && box.pbc_y && box.pbc_z)) + PRINT_INPUT_ERROR("NEP_Spin requires periodic boundary conditions in all three directions."); + bool repartition = !partition_valid_ || !std::equal(box.cpu_h, box.cpu_h + 9, partition_box_); + if (!repartition) { + partition_changed_.fill(0); + const double displacement_limit = 0.2 * rc; + check_partition_displacement<<<(n + 127) / 128, 128>>>( + n, + box, + displacement_limit * displacement_limit, + position.data(), + partition_position_.data(), + partition_changed_.data()); + GPU_CHECK_KERNEL + int changed = 0; + partition_changed_.copy_to_host(&changed); + repartition = changed != 0; + } + if (repartition) { + int bins[3]; + const double cell_width = 0.5 * rc; + box.get_num_bins(cell_width, bins); + for (int d = 0; d < 3; ++d) + bins[d] = std::max(1, bins[d]); + const long long cells = static_cast(bins[0]) * bins[1] * bins[2]; + if (cells > std::numeric_limits::max()) + PRINT_INPUT_ERROR("NEP-Spin partition cell count exceeds the supported int range."); + if (direction_ < 0) + direction_ = std::max_element(bins, bins + 3) - bins; + find_cell_list( + partition_stream_, + direction_, + cell_width, + bins, + box, + n, + position, + cell_count_, + cell_sum_, + order_); +#ifdef USE_HIP + CHECK(hipStreamSynchronize(partition_stream_)); +#else + CHECK(cudaStreamSynchronize(partition_stream_)); +#endif + prefix_.resize(cells + 1); + cell_sum_.copy_to_host(prefix_.data(), cells); + prefix_[cells] = n; + const int longitudinal = bins[direction_]; + const int transverse = cells / longitudinal; + // Four half-cutoff cells span 2*rc; the fifth supplies displacement + // skin plus a rounding margin. If halos overlap, copy each global atom once and retain full + // PBC. + constexpr int halo = 5; + for (int gpu = 0; gpu < devices_; ++gpu) { + Domain& dom = *domains_[gpu]; + int b0 = static_cast(gpu) * longitudinal / devices_; + int b1 = static_cast(gpu + 1) * longitudinal / devices_; + dom.global_first = prefix_[b0 * transverse]; + dom.owned = prefix_[b1 * transverse] - dom.global_first; + if (!dom.owned) { + dom.size = 0; + continue; + } + if (b1 - b0 + 2 * halo >= longitudinal) { + dom.start = 0; + dom.size = n; + dom.first = dom.global_first; + } else { + int lo = (b0 - halo + longitudinal) % longitudinal; + int hi = (b1 + halo) % longitudinal; + dom.start = prefix_[lo * transverse]; + dom.size = prefix_[hi * transverse] - dom.start; + if (hi < lo) + dom.size += n; + dom.first = (dom.global_first - dom.start + n) % n; + } + } + partition_position_.copy_from_device(position.data()); + std::copy(box.cpu_h, box.cpu_h + 9, partition_box_); + partition_valid_ = true; + } + for (int gpu = 0; gpu < devices_; ++gpu) { + Domain& dom = *domains_[gpu]; + if (!dom.owned) + continue; + CHECK(gpuSetDevice(gpu)); + dom.resize(dom.size); + CHECK(gpuSetDevice(0)); + if (packed_type_.size() != (dom.size)) { + packed_type_.resize(dom.size); + } + if (packed_position_.size() != (size_t(dom.size) * 3)) { + packed_position_.resize(size_t(dom.size) * 3); + } + if (packed_spin_.size() != (size_t(dom.size) * 3)) { + packed_spin_.resize(size_t(dom.size) * 3); + } + pack_domain<<<(dom.size + 127) / 128, 128>>>( + n, + dom.size, + dom.start, + order_.data(), + type.data(), + position.data(), + spin.data(), + packed_type_.data(), + packed_position_.data(), + packed_spin_.data()); + GPU_CHECK_KERNEL + CHECK(gpuDeviceSynchronize()); + transfer(dom.type.data(), gpu, packed_type_.data(), 0, size_t(dom.size) * sizeof(int)); + transfer( + dom.position.data(), gpu, packed_position_.data(), 0, size_t(dom.size) * 3 * sizeof(double)); + transfer(dom.spin.data(), gpu, packed_spin_.data(), 0, size_t(dom.size) * 3 * sizeof(double)); + } + std::vector> workers; + for (int gpu = 0; gpu < devices_; ++gpu) { + if (!domains_[gpu]->owned) + continue; + workers.emplace_back(std::async(std::launch::async, [&, gpu] { + CHECK(gpuSetDevice(gpu)); + Domain& dom = *domains_[gpu]; + dom.energy.fill(0.0); + dom.force.fill(0.0); + dom.virial.fill(0.0); + dom.mforce.fill(0.0); + if (repartition) + dom.potential->invalidate_neighbors(); + Box local_box = box; // Neighbor/box bookkeeping is local to each host worker. + dom.potential->compute( + local_box, dom.type, dom.position, dom.spin, dom.energy, dom.force, dom.virial, dom.mforce); + CHECK(gpuDeviceSynchronize()); + })); + } + for (auto& worker : workers) + worker.get(); + CHECK(gpuSetDevice(0)); + for (int gpu = 0; gpu < devices_; ++gpu) { + const Domain& dom = *domains_[gpu]; + if (!dom.owned) + continue; + if (collected_.size() != (size_t(dom.owned) * 16)) { + collected_.resize(size_t(dom.owned) * 16); + } + auto copy_planes = [&](const GPU_Vector& src, int planes, int offset) { + for (int d = 0; d < planes; ++d) + transfer( + collected_.data() + size_t(offset + d) * dom.owned, + 0, + src.data() + size_t(d) * dom.size + dom.first, + gpu, + size_t(dom.owned) * sizeof(double)); + }; + copy_planes(dom.energy, 1, 0); + copy_planes(dom.force, 3, 1); + copy_planes(dom.mforce, 3, 4); + copy_planes(dom.virial, 9, 7); + // Only owner outputs are collected. Every owner's neighbors and their + // descriptors are complete inside 2*rc; ghost output must never be summed. + collect_domain<<<(dom.owned + 127) / 128, 128>>>( + n, + dom.owned, + dom.global_first, + order_.data(), + collected_.data(), + energy.data(), + force.data(), + virial.data(), + mforce.data()); + GPU_CHECK_KERNEL + CHECK(gpuDeviceSynchronize()); + } +} diff --git a/src/force/nep_spin/nep_spin_multigpu.cuh b/src/force/nep_spin/nep_spin_multigpu.cuh new file mode 100644 index 000000000..3384d7ba7 --- /dev/null +++ b/src/force/nep_spin/nep_spin_multigpu.cuh @@ -0,0 +1,51 @@ +/* Single-process spatial decomposition for NEP-Spin. LGPL-3.0-or-later. */ +#pragma once +#include "nep_spin.cuh" +#include + +class NEP_Spin_MultiGPU : public Potential +{ +public: + NEP_Spin_MultiGPU(int devices, const char* model, int atoms, int direction); + ~NEP_Spin_MultiGPU() override; + using Potential::compute; + void compute( + Box&, + const GPU_Vector&, + const GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&) override + { + PRINT_INPUT_ERROR("NEP_Spin requires spin inputs."); + } + void compute( + Box&, + const GPU_Vector&, + const GPU_Vector&, + const GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&) override; + const std::vector& spin_types() const { return domains_[0]->potential->spin_types(); } + +private: + struct Domain { + std::unique_ptr potential; + GPU_Vector type; + GPU_Vector position, spin, energy, force, virial, mforce; + int start = 0, size = 0, first = 0, owned = 0, global_first = 0; + void resize(int count); + }; + int devices_, atoms_, direction_; + std::vector> domains_; + GPU_Vector cell_count_, cell_sum_, order_, packed_type_; + GPU_Vector packed_position_, packed_spin_, collected_; + std::vector prefix_; + GPU_Vector partition_position_; + GPU_Vector partition_changed_; + double partition_box_[9] = {}; + bool partition_valid_ = false; + gpuStream_t partition_stream_ = nullptr; +}; diff --git a/src/force/nep_spin/nep_spin_small_box.cuh b/src/force/nep_spin/nep_spin_small_box.cuh new file mode 100644 index 000000000..99e9be681 --- /dev/null +++ b/src/force/nep_spin/nep_spin_small_box.cuh @@ -0,0 +1,564 @@ +/* + Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. + GPUMD is free software: you can redistribute it and/or modify + it under the terms of the GNU Lesser General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + GPUMD is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Lesser General Public License for more details. + You should have received a copy of the GNU Lesser General Public License + along with GPUMD. If not, see . +*/ + +#pragma once + +#include "force/nep_small_box.cuh" + +// Spin retains double displacement and force accumulation in small periodic cells. +// Ordinary NEP uses its unchanged small-box kernels. +namespace +{ +template +static __device__ __forceinline__ void +spin_complex_product(const Real a, const Real b, Real& real_part, Real& imag_part) +{ + const Real real_temp = real_part; + real_part = a * real_temp - b * imag_part; + imag_part = a * imag_part + b * real_temp; +} + +template +static __device__ __forceinline__ void spin_get_f12_4body( + const Real d12, + const Real d12inv, + const Real fn, + const Real fnp, + const Real Fp, + const Real* s, + const Real* r12, + Real* f12) +{ + Real fn_factor = Fp * fn; + Real fnp_factor = Fp * fnp * d12inv; + Real y20 = (3.0f * r12[2] * r12[2] - d12 * d12); + + // derivative wrt s[0] + Real tmp0 = C4B[0] * 3.0f * s[0] * s[0] + C4B[1] * (s[1] * s[1] + s[2] * s[2]) + + C4B[2] * (s[3] * s[3] + s[4] * s[4]); + Real tmp1 = tmp0 * y20 * fnp_factor; + Real tmp2 = tmp0 * fn_factor; + f12[0] += tmp1 * r12[0] - tmp2 * 2.0f * r12[0]; + f12[1] += tmp1 * r12[1] - tmp2 * 2.0f * r12[1]; + f12[2] += tmp1 * r12[2] + tmp2 * 4.0f * r12[2]; + + // derivative wrt s[1] + tmp0 = C4B[1] * s[0] * s[1] * 2.0f - C4B[3] * s[3] * s[1] * 2.0f + C4B[4] * s[2] * s[4]; + tmp1 = tmp0 * r12[0] * r12[2] * fnp_factor; + tmp2 = tmp0 * fn_factor; + f12[0] += tmp1 * r12[0] + tmp2 * r12[2]; + f12[1] += tmp1 * r12[1]; + f12[2] += tmp1 * r12[2] + tmp2 * r12[0]; + + // derivative wrt s[2] + tmp0 = C4B[1] * s[0] * s[2] * 2.0f + C4B[3] * s[3] * s[2] * 2.0f + C4B[4] * s[1] * s[4]; + tmp1 = tmp0 * r12[1] * r12[2] * fnp_factor; + tmp2 = tmp0 * fn_factor; + f12[0] += tmp1 * r12[0]; + f12[1] += tmp1 * r12[1] + tmp2 * r12[2]; + f12[2] += tmp1 * r12[2] + tmp2 * r12[1]; + + // derivative wrt s[3] + tmp0 = C4B[2] * s[0] * s[3] * 2.0f + C4B[3] * (s[2] * s[2] - s[1] * s[1]); + tmp1 = tmp0 * (r12[0] * r12[0] - r12[1] * r12[1]) * fnp_factor; + tmp2 = tmp0 * fn_factor; + f12[0] += tmp1 * r12[0] + tmp2 * 2.0f * r12[0]; + f12[1] += tmp1 * r12[1] - tmp2 * 2.0f * r12[1]; + f12[2] += tmp1 * r12[2]; + + // derivative wrt s[4] + tmp0 = C4B[2] * s[0] * s[4] * 2.0f + C4B[4] * s[1] * s[2]; + tmp1 = tmp0 * (2.0f * r12[0] * r12[1]) * fnp_factor; + tmp2 = tmp0 * fn_factor; + f12[0] += tmp1 * r12[0] + tmp2 * 2.0f * r12[1]; + f12[1] += tmp1 * r12[1] + tmp2 * 2.0f * r12[0]; + f12[2] += tmp1 * r12[2]; +} + +template +static __device__ __forceinline__ void spin_calculate_s_one( + const int n, const int n_max_angular_plus_1, const float* Fp, const float* sum_fxyz, Real* s) +{ + const int L_minus_1 = L - 1; + const int L_twice_plus_1 = 2 * L + 1; + const int L_square_minus_1 = L * L - 1; + Real Fp_factor = 2.0f * Fp[L_minus_1 * n_max_angular_plus_1 + n]; + s[0] = sum_fxyz[n * num_abc + L_square_minus_1] * C3B[L_square_minus_1] * Fp_factor; + Fp_factor *= 2.0f; + for (int k = 1; k < L_twice_plus_1; ++k) { + s[k] = sum_fxyz[n * num_abc + L_square_minus_1 + k] * C3B[L_square_minus_1 + k] * Fp_factor; + } +} + +template +static __device__ __forceinline__ void spin_accumulate_f12_one( + const Real d12inv, const Real fn, const Real fnp, const Real* s, const Real* r12, Real* f12) +{ + const Real dx[3] = { + (1.0f - r12[0] * r12[0]) * d12inv, -r12[0] * r12[1] * d12inv, -r12[0] * r12[2] * d12inv}; + const Real dy[3] = { + -r12[0] * r12[1] * d12inv, (1.0f - r12[1] * r12[1]) * d12inv, -r12[1] * r12[2] * d12inv}; + const Real dz[3] = { + -r12[0] * r12[2] * d12inv, -r12[1] * r12[2] * d12inv, (1.0f - r12[2] * r12[2]) * d12inv}; + + Real z_pow[L + 1] = {1.0f}; + for (int n = 1; n <= L; ++n) { + z_pow[n] = r12[2] * z_pow[n - 1]; + } + + Real real_part = 1.0f; + Real imag_part = 0.0f; + for (int n1 = 0; n1 <= L; ++n1) { + int n2_start = (L + n1) % 2 == 0 ? 0 : 1; + Real z_factor = 0.0f; + Real dz_factor = 0.0f; + for (int n2 = n2_start; n2 <= L - n1; n2 += 2) { + if (L == 1) { + z_factor += Z_COEFFICIENT_1[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_1[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 2) { + z_factor += Z_COEFFICIENT_2[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_2[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 3) { + z_factor += Z_COEFFICIENT_3[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_3[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 4) { + z_factor += Z_COEFFICIENT_4[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_4[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 5) { + z_factor += Z_COEFFICIENT_5[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_5[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 6) { + z_factor += Z_COEFFICIENT_6[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_6[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 7) { + z_factor += Z_COEFFICIENT_7[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_7[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + if (L == 8) { + z_factor += Z_COEFFICIENT_8[n1][n2] * z_pow[n2]; + if (n2 > 0) { + dz_factor += Z_COEFFICIENT_8[n1][n2] * n2 * z_pow[n2 - 1]; + } + } + } + if (n1 == 0) { + for (int d = 0; d < 3; ++d) { + f12[d] += s[0] * (z_factor * fnp * r12[d] + fn * dz_factor * dz[d]); + } + } else { + Real real_part_n1 = n1 * real_part; + Real imag_part_n1 = n1 * imag_part; + for (int d = 0; d < 3; ++d) { + Real real_part_dx = dx[d]; + Real imag_part_dy = dy[d]; + spin_complex_product(real_part_n1, imag_part_n1, real_part_dx, imag_part_dy); + f12[d] += (s[2 * n1 - 1] * real_part_dx + s[2 * n1 - 0] * imag_part_dy) * z_factor * fn; + } + spin_complex_product(r12[0], r12[1], real_part, imag_part); + const Real xy_temp = s[2 * n1 - 1] * real_part + s[2 * n1 - 0] * imag_part; + for (int d = 0; d < 3; ++d) { + f12[d] += xy_temp * (z_factor * fnp * r12[d] + fn * dz_factor * dz[d]); + } + } + } +} + +template +static __device__ __forceinline__ void accumulate_one_angular_channel_in_double( + const int n, + const int n_max_angular_plus_1, + const double d12inv, + const double fn, + const double fnp, + const float* Fp, + const float* sum_fxyz, + const double* r12unit, + double* f12) +{ + double s[2 * L + 1]; + double contribution[3] = {0.0, 0.0, 0.0}; + spin_calculate_s_one(n, n_max_angular_plus_1, Fp, sum_fxyz, s); + spin_accumulate_f12_one(d12inv, fn, fnp, s, r12unit, contribution); + for (int d = 0; d < 3; ++d) { + f12[d] += static_cast(contribution[d]); + } +} +static __device__ __forceinline__ void find_fn_and_fnp_in_double( + const int n_max, const double rc, const double distance, double* fn, double* fnp) +{ + const double rcinv = 1.0 / rc; + const double scaled = distance * rcinv; + const double fc = distance < rc ? 0.5 * cos(3.14159265358979323846 * scaled) + 0.5 : 0.0; + const double fcp = + distance < rc ? -1.57079632679489661923 * sin(3.14159265358979323846 * scaled) * rcinv : 0.0; + const double x = 2.0 * (scaled - 1.0) * (scaled - 1.0) - 1.0; + fn[0] = fc; + fnp[0] = fcp; + fn[1] = (x + 1.0) * 0.5 * fc; + fnp[1] = 2.0 * (scaled - 1.0) * rcinv * fc + (x + 1.0) * 0.5 * fcp; + double u0 = 1.0; + double u1 = 2.0 * x; + double t0 = 1.0; + double t1 = x; + for (int m = 2; m <= n_max; ++m) { + const double t2 = 2.0 * x * t1 - t0; + t0 = t1; + t1 = t2; + const double derivative_chebyshev = m * u1; + const double u2 = 2.0 * x * u1 - u0; + u0 = u1; + u1 = u2; + const double basis = (t2 + 1.0) * 0.5; + fnp[m] = derivative_chebyshev * 2.0 * (scaled - 1.0) * rcinv * fc + basis * fcp; + fn[m] = basis * fc; + } +} +static __device__ __forceinline__ void accumulate_f12_small_box_in_double( + const int L_max, + const int has_q_222, + const int has_q_1111, + const int has_q_112, + const int has_q_123, + const int has_q_233, + const int has_q_134, + const int num_L, + const int n, + const int n_max_angular_plus_1, + const double d12, + const double* r12, + const double fn, + const double fnp, + const float* Fp, + const float* sum_fxyz, + double* f12) +{ + const bool only_supported_extra = !has_q_1111 && !has_q_112 && !has_q_123 && !has_q_233 && + !has_q_134 && num_L == L_max + static_cast(has_q_222); + if (!only_supported_extra) { + float fallback[3] = {0.0f}; + const float d12_float = static_cast(d12); + const float r12_float[3] = { + static_cast(r12[0]), static_cast(r12[1]), static_cast(r12[2])}; + accumulate_f12( + L_max, + has_q_222, + has_q_1111, + has_q_112, + has_q_123, + has_q_233, + has_q_134, + num_L, + n, + n_max_angular_plus_1, + d12_float, + r12_float, + static_cast(fn), + static_cast(fnp), + Fp, + sum_fxyz, + fallback); + for (int d = 0; d < 3; ++d) { + f12[d] += static_cast(fallback[d]); + } + return; + } + const double d12inv = 1.0 / d12; + const double r12unit[3] = {r12[0] * d12inv, r12[1] * d12inv, r12[2] * d12inv}; + if (L_max >= 1) + accumulate_one_angular_channel_in_double<1>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 2) + accumulate_one_angular_channel_in_double<2>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 3) + accumulate_one_angular_channel_in_double<3>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 4) + accumulate_one_angular_channel_in_double<4>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 5) + accumulate_one_angular_channel_in_double<5>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 6) + accumulate_one_angular_channel_in_double<6>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 7) + accumulate_one_angular_channel_in_double<7>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (L_max >= 8) + accumulate_one_angular_channel_in_double<8>( + n, n_max_angular_plus_1, d12inv, fn, fnp, Fp, sum_fxyz, r12unit, f12); + if (has_q_222) { + double radial_fn = fn * d12inv; + double radial_fnp = fnp * d12inv - fn * d12inv * d12inv; + const double fn2 = radial_fn * d12inv; + const double fnp2 = radial_fnp * d12inv - radial_fn * d12inv * d12inv; + const double s2[5] = { + sum_fxyz[n * NUM_OF_ABC + 3], + sum_fxyz[n * NUM_OF_ABC + 4], + sum_fxyz[n * NUM_OF_ABC + 5], + sum_fxyz[n * NUM_OF_ABC + 6], + sum_fxyz[n * NUM_OF_ABC + 7]}; + double contribution[3] = {0.0, 0.0, 0.0}; + spin_get_f12_4body( + d12, + d12inv, + fn2, + fnp2, + static_cast(Fp[L_max * n_max_angular_plus_1 + n]), + s2, + r12, + contribution); + for (int d = 0; d < 3; ++d) { + f12[d] += static_cast(contribution[d]); + } + } +} +static __global__ void spin_find_force_radial_small_box( + NEP::ParaMB paramb, + NEP::ANN annmb, + const int N, + const int N1, + const int N2, + const int* g_NN, + const int* g_NL, + const int* __restrict__ g_type, + const double* __restrict__ g_x12, + const double* __restrict__ g_y12, + const double* __restrict__ g_z12, + const float* __restrict__ g_Fp, + double* g_fx, + double* g_fy, + double* g_fz, + double* g_virial) +{ + int n1 = blockIdx.x * blockDim.x + threadIdx.x + N1; + if (n1 < N2) { + int t1 = g_type[n1]; + for (int i1 = 0; i1 < g_NN[n1]; ++i1) { + int index = i1 * N + n1; + int n2 = g_NL[index]; + int t2 = g_type[n2]; + double r12_double[3] = { + static_cast(g_x12[index]), + static_cast(g_y12[index]), + static_cast(g_z12[index])}; + const double d12_double = sqrt( + r12_double[0] * r12_double[0] + r12_double[1] * r12_double[1] + + r12_double[2] * r12_double[2]); + double f12_double[3] = {0.0, 0.0, 0.0}; + float rc = (paramb.rc_radial[t1] + paramb.rc_radial[t2]) * 0.5f; + double fn12_double[MAX_NUM_N] = {0.0}; + double fnp12_double[MAX_NUM_N] = {0.0}; + find_fn_and_fnp_in_double( + paramb.basis_size_radial, static_cast(rc), d12_double, fn12_double, fnp12_double); + for (int n = 0; n <= paramb.n_max_radial; ++n) { + double gnp12_double = 0.0; + for (int k = 0; k <= paramb.basis_size_radial; ++k) { + int c_index = get_c_index( + t1 * paramb.num_types + t2, n, k, paramb.n_max_radial, paramb.basis_size_radial); + gnp12_double += fnp12_double[k] * static_cast(annmb.c_type_pair[c_index]); + } + const double tmp12_double = + static_cast(g_Fp[n1 + n * N]) * gnp12_double / d12_double; + for (int d = 0; d < 3; ++d) { + f12_double[d] += tmp12_double * r12_double[d]; + } + } + double s_sxx = 0.0; + double s_sxy = 0.0; + double s_sxz = 0.0; + double s_syx = 0.0; + double s_syy = 0.0; + double s_syz = 0.0; + double s_szx = 0.0; + double s_szy = 0.0; + double s_szz = 0.0; + s_sxx -= r12_double[0] * f12_double[0]; + s_syy -= r12_double[1] * f12_double[1]; + s_szz -= r12_double[2] * f12_double[2]; + s_sxy -= r12_double[0] * f12_double[1]; + s_sxz -= r12_double[0] * f12_double[2]; + s_syz -= r12_double[1] * f12_double[2]; + s_syx -= r12_double[1] * f12_double[0]; + s_szx -= r12_double[2] * f12_double[0]; + s_szy -= r12_double[2] * f12_double[1]; + atomicAdd(&g_fx[n1], f12_double[0]); + atomicAdd(&g_fy[n1], f12_double[1]); + atomicAdd(&g_fz[n1], f12_double[2]); + atomicAdd(&g_fx[n2], -f12_double[0]); + atomicAdd(&g_fy[n2], -f12_double[1]); + atomicAdd(&g_fz[n2], -f12_double[2]); + // save virial + // xx xy xz 0 3 4 + // yx yy yz 6 1 5 + // zx zy zz 7 8 2 + atomicAdd(&g_virial[n2 + 0 * N], s_sxx); + atomicAdd(&g_virial[n2 + 1 * N], s_syy); + atomicAdd(&g_virial[n2 + 2 * N], s_szz); + atomicAdd(&g_virial[n2 + 3 * N], s_sxy); + atomicAdd(&g_virial[n2 + 4 * N], s_sxz); + atomicAdd(&g_virial[n2 + 5 * N], s_syz); + atomicAdd(&g_virial[n2 + 6 * N], s_syx); + atomicAdd(&g_virial[n2 + 7 * N], s_szx); + atomicAdd(&g_virial[n2 + 8 * N], s_szy); + } + } +} +static __global__ void spin_find_force_angular_small_box( + NEP::ParaMB paramb, + NEP::ANN annmb, + const int N, + const int N1, + const int N2, + const int* g_NN_angular, + const int* g_NL_angular, + const int* __restrict__ g_type, + const double* __restrict__ g_x12, + const double* __restrict__ g_y12, + const double* __restrict__ g_z12, + const float* __restrict__ g_Fp, + const float* __restrict__ g_sum_fxyz, + double* g_fx, + double* g_fy, + double* g_fz, + double* g_virial) +{ + int n1 = blockIdx.x * blockDim.x + threadIdx.x + N1; + if (n1 < N2) { + float Fp[MAX_DIM_ANGULAR] = {0.0f}; + float sum_fxyz[NUM_OF_ABC * MAX_NUM_N]; + for (int d = 0; d < paramb.dim_angular; ++d) { + Fp[d] = g_Fp[(paramb.n_max_radial + 1 + d) * N + n1]; + } + for (int n = 0; n < paramb.n_max_angular + 1; ++n) { + for (int abc = 0; abc < (paramb.L_max + 1) * (paramb.L_max + 1) - 1; ++abc) { + sum_fxyz[n * NUM_OF_ABC + abc] = + g_sum_fxyz[(n * ((paramb.L_max + 1) * (paramb.L_max + 1) - 1) + abc) * N + n1]; + } + } + int t1 = g_type[n1]; + for (int i1 = 0; i1 < g_NN_angular[n1]; ++i1) { + int index = i1 * N + n1; + int n2 = g_NL_angular[n1 + N * i1]; + double r12_double[3] = { + static_cast(g_x12[index]), + static_cast(g_y12[index]), + static_cast(g_z12[index])}; + double f12_double[3] = {0.0, 0.0, 0.0}; + int t2 = g_type[n2]; + float rc = (paramb.rc_angular[t1] + paramb.rc_angular[t2]) * 0.5f; + const double d12_double = sqrt( + r12_double[0] * r12_double[0] + r12_double[1] * r12_double[1] + + r12_double[2] * r12_double[2]); + double fn12_double[MAX_NUM_N] = {0.0}; + double fnp12_double[MAX_NUM_N] = {0.0}; + find_fn_and_fnp_in_double( + paramb.basis_size_angular, static_cast(rc), d12_double, fn12_double, fnp12_double); + for (int n = 0; n <= paramb.n_max_angular; ++n) { + double gn12_double = 0.0; + double gnp12_double = 0.0; + for (int k = 0; k <= paramb.basis_size_angular; ++k) { + int c_index = get_c_index( + t1 * paramb.num_types + t2, + n, + k, + paramb.n_max_angular, + paramb.basis_size_angular, + paramb.num_c_radial); + gn12_double += fn12_double[k] * static_cast(annmb.c_type_pair[c_index]); + gnp12_double += fnp12_double[k] * static_cast(annmb.c_type_pair[c_index]); + } + accumulate_f12_small_box_in_double( + paramb.L_max, + paramb.has_q_222, + paramb.has_q_1111, + paramb.has_q_112, + paramb.has_q_123, + paramb.has_q_233, + paramb.has_q_134, + paramb.num_L, + n, + paramb.n_max_angular + 1, + d12_double, + r12_double, + gn12_double, + gnp12_double, + Fp, + sum_fxyz, + f12_double); + } + double s_sxx = 0.0; + double s_sxy = 0.0; + double s_sxz = 0.0; + double s_syx = 0.0; + double s_syy = 0.0; + double s_syz = 0.0; + double s_szx = 0.0; + double s_szy = 0.0; + double s_szz = 0.0; + s_sxx -= r12_double[0] * f12_double[0]; + s_syy -= r12_double[1] * f12_double[1]; + s_szz -= r12_double[2] * f12_double[2]; + s_sxy -= r12_double[0] * f12_double[1]; + s_sxz -= r12_double[0] * f12_double[2]; + s_syz -= r12_double[1] * f12_double[2]; + s_syx -= r12_double[1] * f12_double[0]; + s_szx -= r12_double[2] * f12_double[0]; + s_szy -= r12_double[2] * f12_double[1]; + atomicAdd(&g_fx[n1], f12_double[0]); + atomicAdd(&g_fy[n1], f12_double[1]); + atomicAdd(&g_fz[n1], f12_double[2]); + atomicAdd(&g_fx[n2], -f12_double[0]); + atomicAdd(&g_fy[n2], -f12_double[1]); + atomicAdd(&g_fz[n2], -f12_double[2]); + // save virial + // xx xy xz 0 3 4 + // yx yy yz 6 1 5 + // zx zy zz 7 8 2 + atomicAdd(&g_virial[n2 + 0 * N], s_sxx); + atomicAdd(&g_virial[n2 + 1 * N], s_syy); + atomicAdd(&g_virial[n2 + 2 * N], s_szz); + atomicAdd(&g_virial[n2 + 3 * N], s_sxy); + atomicAdd(&g_virial[n2 + 4 * N], s_sxz); + atomicAdd(&g_virial[n2 + 5 * N], s_syz); + atomicAdd(&g_virial[n2 + 6 * N], s_syx); + atomicAdd(&g_virial[n2 + 7 * N], s_szx); + atomicAdd(&g_virial[n2 + 8 * N], s_szy); + } + } +} +} // namespace diff --git a/src/force/nep_spin/runtime.cu b/src/force/nep_spin/runtime.cu new file mode 100644 index 000000000..a7a208352 --- /dev/null +++ b/src/force/nep_spin/runtime.cu @@ -0,0 +1,93 @@ +/* Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. SPDX-License-Identifier: LGPL-3.0-or-later */ +#include "nep_spin.cuh" +#include "nep_spin_multigpu.cuh" +#include "runtime.cuh" +#include "utilities/error.cuh" +#include "utilities/run_input.cuh" +#include + +Spin_Runtime::Spin_Runtime(const RunInput& input) +{ + for (const auto& line : input.lines()) { + if (!line.tokens.empty()) { + check_command(line.tokens); + } + } +} + +void Spin_Runtime::check_command(const std::vector& tokens) const +{ + bool unsupported = false; + for (const char* name : + {"dump_observer", + "compute_hac", + "compute_hnemd", + "compute_hnemdec", + "compute_shc", + "compute_gkma", + "compute_hnema", + "deform", + "compute_lsqt", + "minimize", + "compute_phonon", + "compute_cohesive", + "compute_elastic", + "kspace", + "dftd3"}) { + unsupported = unsupported || tokens[0] == name; + } + if (tokens[0] == "ensemble") { + unsupported = tokens.size() < 2 || (tokens[1] != "nve" && tokens[1] != "nvt_nhc"); + } + if (unsupported) { + PRINT_INPUT_ERROR(("NEP-spin does not support command " + tokens[0] + ".").c_str()); + } +} + +void Spin_Runtime::initialize(Atom& atom) const +{ + if (types_.empty()) + return; + if (!atom.has_spin) + PRINT_INPUT_ERROR("NEP-spin requires spin:R:3 in model.xyz."); + const int n = atom.number_of_atoms; + for (int i = 0; i < n; ++i) { + const int type = atom.cpu_type[i]; + for (int d = 0; d < 3; ++d) { + double& spin = atom.cpu_spin_per_atom[i + n * d]; + if (!types_[type]) + spin = 0.0; + else if (!std::isfinite(spin)) + PRINT_INPUT_ERROR("Active spin components must be finite."); + } + } + atom.spin_per_atom.copy_from_host(atom.cpu_spin_per_atom.data()); +} + +std::unique_ptr +Spin_Runtime::create_potential(const std::vector& tokens, int atoms) +{ + int direction = -1; + if (tokens.size() == 3) { + if (tokens[2] == "x") + direction = 0; + else if (tokens[2] == "y") + direction = 1; + else if (tokens[2] == "z") + direction = 2; + else + PRINT_INPUT_ERROR("partition direction for multi-GPU NEP can only be x or y or z."); + } + int devices; + CHECK(gpuGetDeviceCount(&devices)); + if (devices == 1) { + std::unique_ptr potential(new NEP_Spin(tokens[1].c_str(), atoms)); + types_ = potential->spin_types(); + return std::unique_ptr(std::move(potential)); + } + std::unique_ptr potential( + new NEP_Spin_MultiGPU(devices, tokens[1].c_str(), atoms, direction)); + types_ = potential->spin_types(); + return std::unique_ptr(std::move(potential)); +} diff --git a/src/force/nep_spin/runtime.cuh b/src/force/nep_spin/runtime.cuh new file mode 100644 index 000000000..38a4e6248 --- /dev/null +++ b/src/force/nep_spin/runtime.cuh @@ -0,0 +1,21 @@ +/* Copyright 2017 Zheyong Fan and GPUMD development team + This file is part of GPUMD. SPDX-License-Identifier: LGPL-3.0-or-later */ +#pragma once +#include "force/potential.cuh" +#include "model/atom.cuh" +#include +#include + +class RunInput; + +class Spin_Runtime +{ +public: + explicit Spin_Runtime(const RunInput& input); + std::unique_ptr create_potential(const std::vector& tokens, int atoms); + void initialize(Atom& atom) const; + +private: + void check_command(const std::vector& tokens) const; + std::vector types_; +}; diff --git a/src/force/potential.cu b/src/force/potential.cu index 7e7fe9342..bf7b7019b 100644 --- a/src/force/potential.cu +++ b/src/force/potential.cu @@ -32,6 +32,19 @@ Potential::~Potential(void) // nothing } +void Potential::compute( + Box&, + const GPU_Vector&, + const GPU_Vector&, + const GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&, + GPU_Vector&) +{ + PRINT_INPUT_ERROR("This potential does not support spin."); +} + static __global__ void gpu_find_force_many_body( const int number_of_particles, const int N1, diff --git a/src/force/potential.cuh b/src/force/potential.cuh index 2593dab07..ab9844f8a 100644 --- a/src/force/potential.cuh +++ b/src/force/potential.cuh @@ -38,6 +38,16 @@ public: GPU_Vector& force, GPU_Vector& virial) = 0; + virtual void compute( + Box& box, + const GPU_Vector& type, + const GPU_Vector& position, + const GPU_Vector& spin, + GPU_Vector& potential, + GPU_Vector& force, + GPU_Vector& virial, + GPU_Vector& mforce); + virtual void compute( const float /* temperature */, Box& /* box */, diff --git a/src/main_gpumd/replicate.cu b/src/main_gpumd/replicate.cu index d21c05f54..3950bcf77 100644 --- a/src/main_gpumd/replicate.cu +++ b/src/main_gpumd/replicate.cu @@ -56,6 +56,10 @@ void Replicate( new_atoms.cpu_atom_symbol.resize(N); new_atoms.cpu_position_per_atom.resize(N * 3); new_atoms.cpu_velocity_per_atom.resize(N * 3); + new_atoms.has_spin = atoms.has_spin; + if (atoms.has_spin) { + new_atoms.cpu_spin_per_atom.resize(N * 3); + } int cur = 0; for (int i = 0; i < r[0]; i++) { for (int j = 0; j < r[1]; j++) { @@ -72,6 +76,9 @@ void Replicate( new_atoms.cpu_position_per_atom[cur + d * N] += i * box.cpu_h[d * 3] + j * box.cpu_h[d * 3 + 1] + k * box.cpu_h[d * 3 + 2]; new_atoms.cpu_velocity_per_atom[cur + d * N] = atoms.cpu_velocity_per_atom[nn + d * n]; + if (atoms.has_spin) { + new_atoms.cpu_spin_per_atom[cur + d * N] = atoms.cpu_spin_per_atom[nn + d * n]; + } } cur++; } @@ -101,6 +108,9 @@ void Replicate( new_atoms.cpu_position_per_atom.begin(), new_atoms.cpu_position_per_atom.end()); atoms.cpu_velocity_per_atom.assign( new_atoms.cpu_velocity_per_atom.begin(), new_atoms.cpu_velocity_per_atom.end()); + atoms.has_spin = new_atoms.has_spin; + atoms.cpu_spin_per_atom.assign( + new_atoms.cpu_spin_per_atom.begin(), new_atoms.cpu_spin_per_atom.end()); for (int& i : atoms.cpu_type_size) i = i * r[0] * r[1] * r[2]; diff --git a/src/main_gpumd/run.cu b/src/main_gpumd/run.cu index 519f6dbb9..3d61e566c 100644 --- a/src/main_gpumd/run.cu +++ b/src/main_gpumd/run.cu @@ -216,17 +216,30 @@ void Run::compute_force() atom.mass); } } else { - force.compute( - box, - atom.position_per_atom, - atom.type, - group, - atom.potential_per_atom, - atom.force_per_atom, - atom.virial_per_atom, - atom.velocity_per_atom, - atom.mass, - atom.position_image.size() > 0 ? atom.position_image.data() : nullptr); + if (force.has_spin_potential()) { + force.compute_spin( + box, + atom.position_per_atom, + atom.type, + atom.potential_per_atom, + atom.force_per_atom, + atom.virial_per_atom, + atom.spin_per_atom, + atom.mforce_per_atom, + atom.position_image.size() > 0 ? atom.position_image.data() : nullptr); + } else { + force.compute( + box, + atom.position_per_atom, + atom.type, + group, + atom.potential_per_atom, + atom.force_per_atom, + atom.virial_per_atom, + atom.velocity_per_atom, + atom.mass, + atom.position_image.size() > 0 ? atom.position_image.data() : nullptr); + } } } @@ -324,6 +337,7 @@ void Run::parse_one_keyword( if (tokens[0] == "potential") { force.parse_potential(tokens, box, atom.type.size(), run_input); + force.initialize_spin(atom); } else if (tokens[0] == "minimize") { Minimize minimize; minimize.parse_minimize( diff --git a/src/main_nep/nep_spin/spin_training_kernels.cuh b/src/main_nep/nep_spin/spin_training_kernels.cuh index 764729ca6..29e23eb92 100644 --- a/src/main_nep/nep_spin/spin_training_kernels.cuh +++ b/src/main_nep/nep_spin/spin_training_kernels.cuh @@ -82,12 +82,6 @@ __global__ void mask_inactive_spin_mforce_training( } } -#ifdef NEP_SPIN_LAYOUT_JIT -using SpinPolynomialLayout = SpinConstantLayout; -#else -using SpinPolynomialLayout = NEP_Spin_Model::Spin_Polynomial_Layout; -#endif - template __device__ __forceinline__ void evaluate_spin3_edge_weights_b9_f32( const float spin_cutoff, diff --git a/src/makefile b/src/makefile index 8c2bdecfc..d3c487d37 100644 --- a/src/makefile +++ b/src/makefile @@ -40,6 +40,7 @@ LIBS = -lcublas -lcusolver -lcufft NEP_LIBS = GNEP_LIBS = ifndef OS +LIBS += -lpthread NEP_LIBS += -ldl GNEP_LIBS += -lpthread endif diff --git a/src/measure/dump_restart.cu b/src/measure/dump_restart.cu index d7b83ad66..cce526879 100644 --- a/src/measure/dump_restart.cu +++ b/src/measure/dump_restart.cu @@ -89,6 +89,9 @@ void Dump_Restart::end_of_step( atom.position_per_atom.copy_to_host(atom.cpu_position_per_atom.data()); atom.velocity_per_atom.copy_to_host(atom.cpu_velocity_per_atom.data()); + if (atom.has_spin) { + atom.spin_per_atom.copy_to_host(atom.cpu_spin_per_atom.data()); + } fprintf(fid, "%d\n", number_of_atoms); @@ -108,11 +111,14 @@ void Dump_Restart::end_of_step( box.cpu_h[5], box.cpu_h[8]); - if (group.size() == 0) { - fprintf(fid, "Properties=species:S:1:pos:R:3:mass:R:1:vel:R:3\n"); - } else { - fprintf(fid, "Properties=species:S:1:pos:R:3:mass:R:1:vel:R:3:group:I:%d\n", int(group.size())); + fprintf(fid, "Properties=species:S:1:pos:R:3:mass:R:1:vel:R:3"); + if (atom.has_spin) { + fprintf(fid, ":spin:R:3"); + } + if (!group.empty()) { + fprintf(fid, ":group:I:%d", int(group.size())); } + fprintf(fid, "\n"); for (int n = 0; n < number_of_atoms; n++) { const double natural_to_A_per_fs = 1.0 / TIME_UNIT_CONVERSION; @@ -128,6 +134,15 @@ void Dump_Restart::end_of_step( atom.cpu_velocity_per_atom[n + number_of_atoms] * natural_to_A_per_fs, atom.cpu_velocity_per_atom[n + 2 * number_of_atoms] * natural_to_A_per_fs); + if (atom.has_spin) { + fprintf( + fid, + "%.17g %.17g %.17g ", + atom.cpu_spin_per_atom[n], + atom.cpu_spin_per_atom[n + number_of_atoms], + atom.cpu_spin_per_atom[n + 2 * number_of_atoms]); + } + for (int m = 0; m < group.size(); ++m) { fprintf(fid, "%d ", group[m].cpu_label[n]); } diff --git a/src/measure/dump_xyz.cu b/src/measure/dump_xyz.cu index 82a582460..8552eb51a 100644 --- a/src/measure/dump_xyz.cu +++ b/src/measure/dump_xyz.cu @@ -117,6 +117,15 @@ void Dump_XYZ::parse( bool split_seen = false; for (int m = 3; m < num_param; ++m) { + if (tokens[m] == "spin") { + has_spin_ = true; + continue; + } + if (tokens[m] == "mforce") { + has_mforce_ = true; + continue; + } + if (tokens[m] == "group") { if (group_seen) { PRINT_INPUT_ERROR("Option 'group' is specified more than once in dump_xyz.\n"); @@ -195,6 +204,18 @@ void Dump_XYZ::pre_run( if (quantities.has_force_) { cpu_force_per_atom_.resize(atom.number_of_atoms * 3); } + if ((has_spin_ || has_mforce_) && + !atom.has_spin) { + PRINT_INPUT_ERROR("dump_xyz spin quantities require spin:R:3 in model.xyz."); + } + + if (has_spin_) { + cpu_spin_per_atom_.resize(atom.number_of_atoms * 3); + } + if (has_mforce_) { + cpu_mforce_per_atom_.resize(atom.number_of_atoms * 3); + } + if (quantities.has_potential_) { cpu_potential_per_atom_.resize(atom.number_of_atoms); } @@ -301,6 +322,13 @@ void Dump_XYZ::output_line2( if (quantities.has_force_) { fprintf(fid_, ":forces:R:3"); } + if (has_spin_) { + fprintf(fid_, ":spin:R:3"); + } + if (has_mforce_) { + fprintf(fid_, ":mforce:R:3"); + } + if (quantities.has_potential_) { fprintf(fid_, ":energy_atom:R:1"); } @@ -363,6 +391,13 @@ void Dump_XYZ::end_of_step( if (quantities.has_force_) { atom.force_per_atom.copy_to_host(cpu_force_per_atom_.data()); } + if (has_spin_) { + atom.spin_per_atom.copy_to_host(cpu_spin_per_atom_.data()); + } + if (has_mforce_) { + atom.mforce_per_atom.copy_to_host(cpu_mforce_per_atom_.data()); + } + if (quantities.has_potential_) { atom.potential_per_atom.copy_to_host(cpu_potential_per_atom_.data()); } @@ -428,6 +463,17 @@ void Dump_XYZ::end_of_step( fprintf(fid_, fmt_.c_str(), cpu_force_per_atom_[m + atom.number_of_atoms * d]); } } + if (has_spin_) { + for (int d = 0; d < 3; ++d) { + fprintf(fid_, fmt_.c_str(), cpu_spin_per_atom_[m + atom.number_of_atoms * d]); + } + } + if (has_mforce_) { + for (int d = 0; d < 3; ++d) { + fprintf(fid_, fmt_.c_str(), cpu_mforce_per_atom_[m + atom.number_of_atoms * d]); + } + } + if (quantities.has_potential_) { fprintf(fid_, fmt_.c_str(), cpu_potential_per_atom_[m]); } diff --git a/src/measure/dump_xyz.cuh b/src/measure/dump_xyz.cuh index 54ec24d16..87678251d 100644 --- a/src/measure/dump_xyz.cuh +++ b/src/measure/dump_xyz.cuh @@ -71,6 +71,8 @@ private: int dump_interval_ = 1; int precision_ = 1; // 1 = single precision, 2 = double DumpQuantities quantities; + bool has_spin_ = false; + bool has_mforce_ = false; int separated_ = 0; int split_frames_ = 0; XYZFrameSplitter splitter_; @@ -80,6 +82,8 @@ private: std::vector cpu_unwrapped_position_; std::vector cpu_force_per_atom_; + std::vector cpu_spin_per_atom_; + std::vector cpu_mforce_per_atom_; std::vector cpu_potential_per_atom_; std::vector cpu_virial_per_atom_; GPU_Vector gpu_total_virial_; diff --git a/src/measure/measure.cu b/src/measure/measure.cu index ce9e72209..811e8d257 100644 --- a/src/measure/measure.cu +++ b/src/measure/measure.cu @@ -20,6 +20,7 @@ The driver class dealing with measurement. #include "measure.cuh" #include "active.cuh" #include "add_efield.cuh" +#include "force/nep_spin/mfield.cuh" #include "add_force.cuh" #include "add_random_force.cuh" #include "add_spring.cuh" @@ -267,6 +268,9 @@ bool Measure::parse_action( std::unique_ptr action; action.reset(new Add_Spring(tokens, group, atom)); actions_.emplace_back(std::move(action)); + } else if (tokens[0] == "add_mfield") { + std::unique_ptr action(new Spin_Magnetic_Field(tokens, group)); + actions_.emplace_back(std::move(action)); } else if (tokens[0] == "add_efield") { std::unique_ptr action; action.reset(new Add_Efield( @@ -306,13 +310,13 @@ void Measure::pre_run( printf(" Please set the action name you developed.\n"); exit(1); } - // dump_xyz, dump_netcdf, compute_chunk, add_force, add_spring, and add_efield are allowed - // to be called multiple times; others are not + // dump_xyz, dump_netcdf, compute_chunk, add_force, add_spring, add_efield, and + // add_mfield are allowed to be called multiple times; others are not if ( action->action_name != "dump_xyz" && action->action_name != "dump_netcdf" && action->action_name != "compute_chunk" && action->action_name != "add_force" && action->action_name != "add_spring" && - action->action_name != "add_efield") { + action->action_name != "add_efield" && action->action_name != "add_mfield") { for (auto& action_name : action_names) { if (action_name == action->action_name) { std::cout << "There are multiple " << action->action_name << " keywords within one run.\n"; diff --git a/src/model/atom.cuh b/src/model/atom.cuh index b9551bda8..989f4a2d2 100644 --- a/src/model/atom.cuh +++ b/src/model/atom.cuh @@ -30,6 +30,7 @@ public: std::vector cpu_charge; // float is sufficient for charge std::vector cpu_position_per_atom; std::vector cpu_velocity_per_atom; + std::vector cpu_spin_per_atom; std::vector cpu_atom_symbol; GPU_Vector type; // per-atom type (1 component) GPU_Vector mass; // per-atom mass (1 component) @@ -38,10 +39,13 @@ public: GPU_Vector position_image; // per-atom periodic image (3 components) GPU_Vector unwrapped_position; // unwrapped per-atom position (3 components) GPU_Vector velocity_per_atom; // per-atom velocity (3 components) + GPU_Vector spin_per_atom; // Cartesian magnetic moment (3 components) + GPU_Vector mforce_per_atom; // -dE/dS (3 components) GPU_Vector force_per_atom; // per-atom force (3 components) GPU_Vector heat_per_atom; // per-atom heat current (5 components) GPU_Vector virial_per_atom; // per-atom virial (9 components) GPU_Vector potential_per_atom; // per-atom potential energy (1 component) + bool has_spin = false; // for beads in PIMD int number_of_beads = 0; std::vector> position_beads; diff --git a/src/model/read_xyz.cu b/src/model/read_xyz.cu index 6f9aa911f..cff34805c 100644 --- a/src/model/read_xyz.cu +++ b/src/model/read_xyz.cu @@ -25,6 +25,7 @@ The class defining the simulation model. #include "utilities/compact_nep.cuh" #include "utilities/error.cuh" #include "utilities/run_input.cuh" +#include "utilities/nep_spin_model_file.cuh" #include "utilities/gpu_macro.cuh" #include #include @@ -160,6 +161,8 @@ static void read_xyz_line_2( int& has_velocity_in_xyz, bool& has_mass, bool& has_charge, + bool& has_spin, + const bool enable_spin, int& num_columns, int* property_offset, std::vector& group) @@ -241,10 +244,11 @@ static void read_xyz_line_2( } // properties - const int num_properties = 6; - std::string property_name[num_properties] = {"species", "pos", "mass", "charge", "vel", "group"}; + const int num_properties = 7; + std::string property_name[num_properties] = { + "species", "pos", "mass", "charge", "vel", "group", "spin"}; int property_position[num_properties] = { - -1, -1, -1, -1, -1, -1}; // species,pos,mass,charge,vel,group + -1, -1, -1, -1, -1, -1, -1}; for (int n = 0; n < tokens.size(); ++n) { const std::string properties_string = "properties="; if (tokens[n].substr(0, properties_string.length()) == properties_string) { @@ -257,12 +261,23 @@ static void read_xyz_line_2( std::vector sub_tokens = get_tokens(line); for (int k = 0; k < sub_tokens.size() / 3; ++k) { for (int prop = 0; prop < num_properties; ++prop) { - if (sub_tokens[k * 3] == property_name[prop]) { + if (sub_tokens[k * 3] == property_name[prop] && (prop != 6 || enable_spin)) { + if (prop == 6 && property_position[prop] >= 0) { + PRINT_INPUT_ERROR("Duplicate spin property."); + } property_position[prop] = k; } } } + if (property_position[6] >= 0) { + const int base = property_position[6] * 3; + if (sub_tokens[base + 1] != "r" || + get_int_from_token(sub_tokens[base + 2], __FILE__, __LINE__) != 3) { + PRINT_INPUT_ERROR("spin must be R:3."); + } + } + if (property_position[4] < 0) { has_velocity_in_xyz = 0; printf("Do not specify initial velocities here.\n"); @@ -309,6 +324,7 @@ static void read_xyz_line_2( } else { has_charge = true; } + has_spin = property_position[6] >= 0; } void read_xyz_in_line_3( @@ -317,6 +333,7 @@ void read_xyz_in_line_3( const int has_velocity_in_xyz, const bool has_mass, const bool has_charge, + const bool has_spin, const int num_columns, const int* property_offset, int& number_of_types, @@ -327,6 +344,7 @@ void read_xyz_in_line_3( std::vector& cpu_charge, std::vector& cpu_position_per_atom, std::vector& cpu_velocity_per_atom, + std::vector& cpu_spin_per_atom, std::vector& group) { cpu_atom_symbol.resize(N); @@ -335,6 +353,7 @@ void read_xyz_in_line_3( cpu_charge.resize(N, 0.0); cpu_position_per_atom.resize(N * 3); cpu_velocity_per_atom.resize(N * 3); + cpu_spin_per_atom.resize(has_spin ? N * 3 : 0); number_of_types = atom_symbols.size(); for (int m = 0; m < group.size(); ++m) { @@ -388,6 +407,13 @@ void read_xyz_in_line_3( } } + if (has_spin) { + for (int d = 0; d < 3; ++d) { + cpu_spin_per_atom[n + N * d] = + get_double_from_token(tokens[property_offset[6] + d], __FILE__, __LINE__); + } + } + for (int m = 0; m < group.size(); ++m) { group[m].cpu_label[n] = get_int_from_token(tokens[property_offset[5] + m], __FILE__, __LINE__); @@ -493,12 +519,21 @@ void initialize_position( atom_symbols = get_atom_symbols(filename_potential); read_xyz_line_1(input, atom.number_of_atoms); - int property_offset[6] = {0, 0, 0, 0, 0, 0}; // species,pos,mass,vel,group + int property_offset[7] = {0, 0, 0, 0, 0, 0, 0}; int num_columns = 0; bool has_mass = true; bool has_charge = true; read_xyz_line_2( - input, box, has_velocity_in_xyz, has_mass, has_charge, num_columns, property_offset, group); + input, + box, + has_velocity_in_xyz, + has_mass, + has_charge, + atom.has_spin, + NEP_Spin_Model_File::is_spin_model(filename_potential), + num_columns, + property_offset, + group); read_xyz_in_line_3( input, @@ -506,6 +541,7 @@ void initialize_position( has_velocity_in_xyz, has_mass, has_charge, + atom.has_spin, num_columns, property_offset, number_of_types, @@ -516,6 +552,7 @@ void initialize_position( atom.cpu_charge, atom.cpu_position_per_atom, atom.cpu_velocity_per_atom, + atom.cpu_spin_per_atom, group); input.close(); @@ -560,6 +597,11 @@ void allocate_memory_gpu(std::vector& group, Atom& atom, GPU_Vector @@ -419,7 +420,7 @@ void prepare_compact_nep_files( } std::vector first = get_first_line(potential_files.back()); - if (!is_compactable_nep(first)) { + if (!is_compactable_nep(first) || NEP_Spin_Model_File::is_spin_model(potential_files.back())) { return; } for (const std::string& filename : potential_files) { diff --git a/src/utilities/nep_spin_model_file.cu b/src/utilities/nep_spin_model_file.cu index 3e71d765c..fbde91988 100644 --- a/src/utilities/nep_spin_model_file.cu +++ b/src/utilities/nep_spin_model_file.cu @@ -3,6 +3,7 @@ #include "nep_spin_model_file.cuh" #include #include +#include #include namespace @@ -235,3 +236,15 @@ int NEP_Spin_Model_File::spin_dimension(int channels, int l_max, int order, int } return dimension; } + +bool NEP_Spin_Model_File::is_spin_model(const std::string& filename) +{ + std::ifstream input(filename); + NEP_File file; + std::string error; + if (!file.read(input, error)) { + return false; + } + const NEP_File::Field* model = file.find_field("model"); + return model && model->values.size() == 1 && model->values[0] == "nep_spin"; +} diff --git a/src/utilities/nep_spin_model_file.cuh b/src/utilities/nep_spin_model_file.cuh index 349b54ee9..e160dd1ae 100644 --- a/src/utilities/nep_spin_model_file.cuh +++ b/src/utilities/nep_spin_model_file.cuh @@ -23,6 +23,7 @@ struct NEP_Spin_Model_File { std::vector energy_reference; std::vector parameters, q_scaler; + static bool is_spin_model(const std::string& filename); bool read(std::istream& input, std::string& error); static int spin_dimension(int channels, int l_max, int order, int soc); }; diff --git a/tests_pytest/spin_fixtures.py b/tests_pytest/spin_fixtures.py index 57e8cf6e3..c1219d160 100644 --- a/tests_pytest/spin_fixtures.py +++ b/tests_pytest/spin_fixtures.py @@ -5,6 +5,19 @@ FIXTURE = Path(__file__).parent / "fixtures" / "nep_spin3" +def write_xyz(path, case): + lattice = " ".join(f"{value:.17g}" for value in case["cell"]) + lines = [ + str(len(case["types"])), + f'Lattice="{lattice}" ' + 'Properties=species:S:1:pos:R:3:spin:R:3 pbc="T T T"', + ] + for symbol, position, spin in zip( + case["types"], case["positions"], case["spins"]): + values = position + spin + lines.append(symbol + " " + " ".join(f"{value:.17g}" for value in values)) + path.write_text("\n".join(lines) + "\n") + def reshape(values, width): if len(values) % width: raise ValueError(f"cannot reshape {len(values)} values with width {width}") diff --git a/tests_pytest/spin_runtime_helpers.py b/tests_pytest/spin_runtime_helpers.py new file mode 100644 index 000000000..66f242e8e --- /dev/null +++ b/tests_pytest/spin_runtime_helpers.py @@ -0,0 +1,26 @@ +"""Fixed-spin model evaluation helpers.""" +import os, shutil, subprocess +from pathlib import Path +from spin_fixtures import FIXTURE +ROOT = Path(__file__).resolve().parents[1] +GPUMD = Path(os.environ.get("GPUMD_COMMAND", ROOT / "src/gpumd")) + +def base_model(): + lines = (FIXTURE / "model_large_box.xyz").read_text().splitlines() + lines[1] = lines[1].replace( + "Properties=species:S:1:pos:R:3:spin:R:3", + "Properties=species:S:1:pos:R:3:vel:R:3:spin:R:3") + for index in range(2, len(lines)): + fields = lines[index].split() + lines[index] = " ".join(fields[:4] + ["0", "0", "0"] + fields[4:]) + return "\n".join(lines) + "\n" + +def run_case(root, name, run_input, model_text=None): + case = root / name + case.mkdir() + shutil.copy(FIXTURE / "nep_spin_o3c2_uniform.nep", case / "nep.txt") + (case / "model.xyz").write_text(model_text or base_model()) + (case / "run.in").write_text(run_input) + result = subprocess.run( + [str(GPUMD)], cwd=case, capture_output=True, text=True, check=False) + return case, result diff --git a/tests_pytest/test_spin_model_format.py b/tests_pytest/test_spin_model_format.py index 67accc960..9681cbf1d 100644 --- a/tests_pytest/test_spin_model_format.py +++ b/tests_pytest/test_spin_model_format.py @@ -26,6 +26,9 @@ def codec(tmp_path_factory): NEP_Spin_Model_File model; std::string error; if (!model.read(input, error)) { std::cerr << error; return 1; } + if (!NEP_Spin_Model_File::is_spin_model(argv[1])) { + std::cerr << "Spin model was not recognized"; return 2; + } std::cout << std::setprecision(17); auto dump = [](const char* name, const auto& values) { std::cout << name; @@ -142,3 +145,12 @@ def test_spin_cutoff_bound(codec, tmp_path, cutoff, spin_first): else: assert result.returncode == 1 assert 'spin_cutoff must not exceed the NEP radial cutoff' in result.stderr + + +@pytest.mark.parametrize('record', ['MODEL NEP_SPIN', 'model nep_spin#comment', None]) +def test_model_selection_uses_file_parser(codec, tmp_path, record): + model, _ = load_fixture('o3c2_uniform') + changed = model.replace('model nep_spin\n', record + '\n' if record else '') + if record is None: + changed += '\nmodel nep_spin\n' + assert parse(codec, tmp_path / 'model.nep', changed)['spin_type'] == [1, 0] diff --git a/tests_pytest/validate_mfield_runtime.py b/tests_pytest/validate_mfield_runtime.py new file mode 100644 index 000000000..3b66b11be --- /dev/null +++ b/tests_pytest/validate_mfield_runtime.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""GPU runtime checks for group-selected Zeeman coupling; standard library only.""" +import json +import math +import tempfile +from pathlib import Path +import spin_runtime_helpers as runtime + +CB = 5.7883818060e-5 +FIELD = [12.0, -7.0, 31.0] + +def grouped(model): + lines = model.splitlines() + lines[1] = lines[1].replace('spin:R:3', 'spin:R:3:group:I:2') + for i in range(2, len(lines)): + lines[i] += f' 0 {(i-2)%2}' + return '\n'.join(lines)+'\n' + +def frame(path): + lines = path.read_text().splitlines() + n = int(lines[0]) + values = [[float(x) for x in line.split()[1:]] for line in lines[-n:]] + assert all(math.isfinite(x) for row in values for x in row) + return values + +def main(): + original_run = runtime.run_case + with tempfile.TemporaryDirectory(prefix='gpumd-mfield-') as tmp: + root = Path(tmp) + model = grouped(runtime.base_model()) + def static(name, command, model_text=None): + case, result = original_run(root, name, + 'potential nep.txt\n'+command+'\nensemble nve\ntime_step 0\n' + 'dump_xyz 1 state.xyz precision double mass spin mforce potential\nrun 1\n', model_text or model) + assert result.returncode == 0, result.stdout+result.stderr + return frame(case/'state.xyz') + baseline = static('baseline', '') + applied = static('field', 'add_mfield 1 1 12 -7 31') + for i, (before, after) in enumerate(zip(baseline, applied)): + h = [CB*b if i%2 else 0.0 for b in FIELD] + for k in range(3): + assert abs((after[7+k]-before[7+k])-h[k]) < 2e-12 + expected = -sum(before[4+k]*h[k] for k in range(3)) + assert abs(after[10]-before[10]-expected) < 2e-12 + excluded = model.splitlines() + fields = excluded[3].split() + fields[0] = "Ge" + excluded[3] = " ".join(fields) + excluded = "\n".join(excluded) + "\n" + masked = static("excluded-off", "", excluded) + masked_field = static("excluded-on", "add_mfield 1 1 12 -7 31", excluded) + assert masked_field[1][4:10] == [0.0] * 6 + assert abs(masked_field[1][10] - masked[1][10]) < 2e-12 + zero = static('zero', 'add_mfield 1 1 0 0 0') + assert zero == baseline + with_zero = static('with-zero', 'add_mfield 1 1 12 -7 31\nadd_mfield 0 0 0 0 0') + assert with_zero == applied + summed = static('summed', 'add_mfield 1 1 5 -2 10\nadd_mfield 1 1 7 -5 21') + for left, right in zip(summed, applied): + assert max(abs(a-b) for a, b in zip(left, right)) < 2e-12 + cancelled = static('cancelled', 'add_mfield 1 1 12 -7 31\nadd_mfield 1 1 -12 7 -31') + for left, right in zip(cancelled, baseline): + assert max(abs(a-b) for a, b in zip(left, right)) < 2e-12 + commands = 'potential nep.txt\n' + for name, command in [('first', 'add_mfield 1 1 12 -7 31\n'), + ('second', ''), ('third', 'add_mfield 1 1 12 -7 31\n')]: + commands += (command + 'ensemble nve\ntime_step 0\n' + f'dump_xyz 1 {name}.xyz precision double mass spin mforce potential\nrun 1\n') + case, result = original_run(root, 'multiple-runs', commands, model) + assert result.returncode == 0, result.stdout + result.stderr + for name, expected in [('first', applied), ('second', baseline), ('third', applied)]: + actual = frame(case / (name + '.xyz')) + for left, right in zip(actual, expected): + assert max(abs(a-b) for a, b in zip(left, right)) < 2e-12, name + negative = ['add_mfield 0 0 1 2', 'add_mfield -1 0 1 2 3', + 'add_mfield 0 99 1 2 3', 'add_mfield 0 0 nan 0 0', + 'add_mfield 0 0 1 2 3 block 0 1 0 1 0 1'] + for i, command in enumerate(negative): + _, result = original_run(root, f'negative{i}', 'potential nep.txt\n'+command+'\n', model) + assert result.returncode != 0, command + region_command = 'add_mfield region 0.25 0.75 0 1 0 1 12 -7 31' + for scale, skew, wrap in [(20, 0, False), (20, 4, False), (40, 8, False), (20, 0, True)]: + lines = runtime.base_model().splitlines() # region mode needs no groups + import re + lines[1] = re.sub(r'Lattice="[^"]+"', + f'Lattice="{scale} 0 0 {skew} {scale} 0 0 0 {scale}"', lines[1]) + for i, fraction in enumerate([0.25, 0.5, 0.75, 0.875]): + fields = lines[i+2].split() + fields[1:4] = [str(scale*(fraction + (1 if wrap else 0)) + skew*0.25), + str(scale*0.25), str(scale*0.25)] + lines[i+2] = ' '.join(fields) + text = '\n'.join(lines)+'\n' + tag = f'region-{scale}-{skew}-{wrap}' + before = static(tag+'-off', '', text) + after = static(tag+'-on', region_command, text) + for i, (left, right) in enumerate(zip(before, after)): + h = [CB*b if i < 2 else 0.0 for b in FIELD] + for k in range(3): + assert abs(right[7+k]-left[7+k]-h[k]) < 2e-12, tag + assert abs(right[10]-left[10]+sum(left[4+k]*h[k] for k in range(3))) < 2e-12, tag + # A moving atom enters and then leaves the region; the field must be + # recomputed from current positions, not cached as a fixed group. + lines = runtime.base_model().splitlines() + fields = lines[2].split() + fields[4] = '2' # vx = 2 Angstrom/fs + lines[2] = ' '.join(fields) + moving = '\n'.join(lines)+'\n' + states = [] + for tag, command in [('off', ''), ('on', 'add_mfield region 0.015 0.025 0 1 0 1 12 -7 31\n')]: + case, result = original_run(root, 'moving-'+tag, + 'potential nep.txt\n'+command+'ensemble nve\ntime_step 0.1\n' + 'dump_xyz 1 state.xyz precision double mass spin mforce potential\nrun 2\n', moving) + assert result.returncode == 0, result.stdout+result.stderr + rows = (case/'state.xyz').read_text().splitlines() + states.append([[float(x) for x in rows[offset+2].split()[1:]] for offset in (0, 6)]) + assert 0.3 <= states[1][0][0] < 0.5 and states[1][1][0] >= 0.5 + for step in range(2): + for k in range(3): + expected = CB*FIELD[k] if step == 0 else 0.0 + assert abs(states[1][step][7+k]-states[0][step][7+k]-expected) < 2e-12 + for i, bounds in enumerate(['-0.1 1 0 1 0 1', '0 1.1 0 1 0 1', + '0.5 0.5 0 1 0 1', '0 1 nan 1 0 1']): + _, result = original_run(root, f'region-negative{i}', + 'potential nep.txt\nadd_mfield region '+bounds+' 0 0 10\n', runtime.base_model()) + assert result.returncode != 0 + print(json.dumps({'group_energy_force': 'passed', 'zero_field': 'passed', + 'parser_negatives': len(negative)+4, 'region_boundary_skew_wrap': 'passed', + 'region_enter_exit': 'passed', 'additive_fields': 'passed', + 'per_run_lifecycle': 'passed'})) + +if __name__ == '__main__': + main() diff --git a/tests_pytest/validate_nep_spin3_runtime.py b/tests_pytest/validate_nep_spin3_runtime.py new file mode 100644 index 000000000..81a7b2b0f --- /dev/null +++ b/tests_pytest/validate_nep_spin3_runtime.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""Validate the GPUMD NEP-spin force path against the frozen FP64 oracle.""" + +import argparse +import math +import os +import re +import subprocess +import tempfile +from pathlib import Path +from spin_fixtures import write_xyz, load_fixture +from validate_main_nep_spin3_training import maximum_error, finite_values + + +ROOT = Path(__file__).resolve().parents[1] +FIXTURE = Path(__file__).parent / "fixtures" / "nep_spin3" +GPUMD = Path(os.environ.get("GPUMD_COMMAND", ROOT / "src" / "gpumd")) + + +def read_frame(path): + lines = path.read_text().splitlines() + atom_count = int(lines[0]) + rows = [line.split() for line in lines[-atom_count:]] + values = [[float(value) for value in row[1:]] for row in rows] + finite_values(values) + energy = float( + re.search(r"\benergy=([-+0-9.eE]+)", lines[-atom_count - 1]).group(1)) + finite_values([energy]) + virial = [] + for row in values: + # dump_xyz already serializes the internal 9-plane tensor in row-major + # order (xx, xy, xz, yx, yy, yz, zx, zy, zz). + virial.append(row[13:22]) + return { + "spin": [row[6:9] for row in values], + "energy": energy, + "force": [row[3:6] for row in values], + "mforce": [row[9:12] for row in values], + "potential": [row[12] for row in values], + "atom_virial": virial, + "virial": [sum(row[index] for row in virial) for index in range(9)], + } + + +def flatten(values): + if not isinstance(values, list): + return [values] + result = [] + for value in values: + result.extend(flatten(value)) + return result + + +max_error = maximum_error + + +def worst_difference(candidate, reference): + left = flatten(candidate) + right = flatten(reference) + index = max(range(len(left)), key=lambda item: abs(left[item] - right[item])) + finite_values(left) + finite_values(right) + return index, left[index], right[index] + + +def run_case(root, name, case, model): + directory = root / name + directory.mkdir() + write_xyz(directory / "model.xyz", case) + (directory / "nep.txt").write_text(model) + (directory / "run.in").write_text( + "potential nep.txt\n" + "velocity 1\n" + "ensemble nve\n" + "time_step 0\n" + "dump_xyz 1 result.xyz force potential spin mforce virial precision double\n" + "run 1\n") + result = subprocess.run( + [str(GPUMD)], cwd=directory, capture_output=True, text=True, check=False) + if result.returncode: + raise RuntimeError( + f"{name}: GPUMD failed with {result.returncode}\n" + f"{result.stdout}{result.stderr}") + return read_frame(directory / "result.xyz") + + +def validate_energy_reference(root): + model, oracle = load_fixture("o3c2_uniform") + case = oracle["cases"]["noncollinear"] + base = run_case(root, "reference-base", case, model) + old = next(line for line in model.splitlines() if line.startswith("energy_reference ")) + references = list(map(float, old.split()[1:])) + shifts = [100000.123456789, -200000.987654321] + changed = "energy_reference " + " ".join(format(a + b, '.17g') for a, b in zip(references, shifts)) + shifted = run_case(root, "reference-shift", case, model.replace(old, changed)) + expected = base["energy"] + sum(shifts[0 if t == "Fe" else 1] for t in case["types"]) + if abs(shifted["energy"] - expected) > 2e-6: + raise AssertionError(f"energy reference mismatch: actual={shifted['energy']:.17g}, expected={expected:.17g}") + for field in ("force", "mforce", "virial"): + if max_error(base[field], shifted[field]) > 1e-10: + raise AssertionError(f"constant energy reference changed {field}") + print("Spin3 double energy-reference shift and derivative invariance passed") + + + +def validate_io(root): + import copy + import shutil + model, oracle = load_fixture("o3c2_uniform") + case = copy.deepcopy(oracle["cases"]["noncollinear"]) + # Include an excluded type with deliberately nonzero input spin. + case["types"][-1] = "Ge" + case["spins"][-1] = [0.8, -0.2, 0.4] + raw = run_case(root, "inactive-raw", case, model) + case["spins"][-1] = [0.0, 0.0, 0.0] + zero = run_case(root, "inactive-zero", case, model) + for field in ("energy", "potential", "force", "mforce", "virial", "spin"): + assert max_error(raw[field], zero[field]) < 1e-10, field + assert raw["spin"][-1] == [0.0] * 3 and raw["mforce"][-1] == [0.0] * 3 + + # Selection and validation must share the same lexical and ordering rules. + variants = { + "uppercase": model.replace("model nep_spin", "MODEL NEP_SPIN"), + "inline-comment": model.replace("model nep_spin", "model nep_spin#comment"), + "model-after-blocks": model.replace("model nep_spin\n", "") + "\nmodel nep_spin\n", + } + for name, changed in variants.items(): + alternate = run_case(root, name, case, changed) + for field in ("energy", "force", "mforce", "virial"): + assert max_error(alternate[field], zero[field]) < 1e-10, (name, field) + + for prop, message in [("spin:R:1", "spin must be R:3"), + ("spin:R:3:spin:R:3", "Duplicate spin property")]: + from spin_runtime_helpers import base_model, run_case as input_case + _, result = input_case(root, "invalid-spin-" + str(len(prop)), + "potential nep.txt\nensemble nve\nrun 1\n", + base_model().replace("spin:R:3", prop)) + assert result.returncode != 0 and message in result.stdout + result.stderr + + # Small-cell lists must count and grow rather than truncate a short initial allocation. + reference = oracle["cases"]["noncollinear"] + grown = run_case(root, "small-capacity-growth", reference, + model.replace("cutoff 6 5 128 128", "cutoff 6 5 1 1")) + regular = run_case(root, "small-capacity-reference", reference, model) + for field in ("energy", "force", "mforce", "virial"): + assert max_error(grown[field], regular[field]) < 1e-10, field + + # Restart and replication preserve the fixed-spin model, including type masking. + source = root / "io-source" + source.mkdir() + (source / "nep.txt").write_text(model) + write_xyz(source / "model.xyz", case) + (source / "run.in").write_text( + "replicate 2 1 1\npotential nep.txt\nvelocity 1\nensemble nve\ntime_step 0\n" + "dump_restart 1\ndump_xyz 1 result.xyz force potential spin mforce virial precision double\nrun 1\n") + result = subprocess.run([str(GPUMD)], cwd=source, capture_output=True, text=True) + assert result.returncode == 0, result.stdout + result.stderr + before = read_frame(source / "result.xyz") + expected_spin = case["spins"] * 2 + assert max_error(before["spin"], expected_spin) < 1e-10 + resume = root / "io-resume" + resume.mkdir() + shutil.copy(source / "nep.txt", resume / "nep.txt") + shutil.copy(source / "restart.xyz", resume / "model.xyz") + (resume / "run.in").write_text( + "potential nep.txt\nensemble nve\ntime_step 0\n" + "dump_xyz 1 result.xyz force potential spin mforce virial precision double\nrun 1\n") + result = subprocess.run([str(GPUMD)], cwd=resume, capture_output=True, text=True) + assert result.returncode == 0, result.stdout + result.stderr + after = read_frame(resume / "result.xyz") + for field in ("energy", "potential", "force", "mforce", "virial", "spin"): + assert max_error(before[field], after[field]) < 2e-4, field + # Use upstream split output with the same Spin quantities and reader. + (resume / "run.in").write_text( + "potential nep.txt\nensemble nve\ntime_step 0\n" + "dump_xyz 1 split.xyz split 1 force potential spin mforce virial precision double\nrun 2\n") + result = subprocess.run([str(GPUMD)], cwd=resume, capture_output=True, text=True) + assert result.returncode == 0, result.stdout + result.stderr + for path in (resume / "split_0001.xyz", resume / "split_0002.xyz"): + split = read_frame(path) + for field in ("energy", "force", "mforce", "virial", "spin"): + assert max_error(split[field], after[field]) < 2e-4, field + print("model selection, neighbor growth, masking, replication, restart and split output passed") + +def validate_thin_cell(root): + # 10,000 atoms and two periodic images previously reserved ~19.2 GB. + model, _ = load_fixture("o3c2_uniform") + positions = [[1.5 + 3 * i, 2 + 4 * j, 1 + 2 * k] + for i in range(2) for j in range(50) for k in range(100)] + case = {"types": ["Fe"] * len(positions), "positions": positions, + "spins": [[0.2, -0.1, 2.0]] * len(positions), + "cell": [6., 0., 0., 0., 200., 0., 0., 0., 200.]} + small = run_case(root, "thin-cell", case, model) + expanded = dict(case) + expanded["cell"] = [18., 0., 0., 0., 200., 0., 0., 0., 200.] + expanded["positions"] = [[x + 6 * shift, y, z] for shift in range(3) for x, y, z in positions] + expanded["types"] = case["types"] * 3 + expanded["spins"] = case["spins"] * 3 + large = run_case(root, "thin-cell-expanded", expanded, model) + for field in ("potential", "force", "mforce"): + assert max_error(small[field], large[field][:len(positions)]) < 2e-4, field + assert max_error(small["virial"], [value / 3 for value in large["virial"]]) < 2e-2 + print("10000-atom thin-cell allocation and equivalent expanded-cell comparison passed") + + +def validate_ordinary_spin_properties(root): + # Unknown properties must remain ignored for an ordinary model, even if named spin. + model = (ROOT / "tests_pytest/fixtures/models/nep_C.txt").read_text() + variants = [("none", "", ""), ("scalar", ":spin:R:1", " 7"), + ("vector", ":spin:R:3", " 1 2 3"), + ("nonfinite", ":spin:R:3", " nan inf -inf"), + ("string", ":spin:S:1", " ignored"), + ("duplicate", ":spin:R:1:spin:R:1", " 1 2")] + reference = None + for name, prop, values in variants: + case = root / ("ordinary-spin-" + name) + case.mkdir() + (case / "nep.txt").write_text(model) + (case / "model.xyz").write_text( + '2\nLattice="20 0 0 0 20 0 0 0 20" ' + 'Properties=species:S:1:pos:R:3' + prop + ':vel:R:3 pbc="T T T"\n' + + 'C 1 1 1' + values + ' 0 0 0\nC 2.5 1 1' + values + ' 0 0 0\n') + (case / "run.in").write_text( + "potential nep.txt\nensemble nve\ntime_step 0\n" + "dump_xyz 1 state.xyz force potential precision double\ndump_restart 1\nrun 1\n") + result = subprocess.run([str(GPUMD)], cwd=case, capture_output=True, text=True) + assert result.returncode == 0, (name, result.stdout, result.stderr) + outputs = [(case / file).read_bytes() for file in ("state.xyz", "restart.xyz")] + assert all(b"spin:" not in data and b"mforce:" not in data for data in outputs) + if reference is None: + reference = outputs + else: + assert outputs == reference, name + print("ordinary potentials ignore scalar/vector/nonfinite/string/duplicate spin properties") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--cases", nargs="*", default=None, + help="oracle cases to run; default runs every case") + parser.add_argument( + "--fixtures", nargs="*", + choices=("o3c2_uniform", "o3c2_typewise", "o3c3_uniform"), + default=("o3c2_uniform", "o3c2_typewise", "o3c3_uniform"), + help="frozen Spin3 fixtures to run") + parser.add_argument("--tolerance", type=float, default=2.0e-4) + args = parser.parse_args() + if not math.isfinite(args.tolerance) or args.tolerance <= 0: + parser.error("tolerance must be finite and positive") + + errors = {} + with tempfile.TemporaryDirectory(prefix="gpumd-spin3-") as temp: + root = Path(temp) + validate_ordinary_spin_properties(root) + validate_energy_reference(root) + validate_io(root) + validate_thin_cell(root) + for fixture in args.fixtures: + model, oracle = load_fixture(fixture) + selected = args.cases or list(oracle["cases"]) + fixture_root = root / fixture + fixture_root.mkdir() + for name in selected: + if name not in oracle["cases"]: + continue + reference = oracle["cases"][name] + candidate = run_case(fixture_root, name, reference, model) + case_errors = { + "energy": max_error(candidate["energy"], sum(reference["potential"])), + "potential": max_error(candidate["potential"], reference["potential"]), + "force": max_error(candidate["force"], reference["force"]), + "mforce": max_error(candidate["mforce"], reference["mforce"]), + "virial": max_error( + candidate["virial"], + [sum(row[index] for row in reference["atom_virial"]) + for index in range(9)]), + } + key = f"{fixture}/{name}" + errors[key] = case_errors + print(key + ": " + " ".join( + f"{field}={error:.3e}" for field, error in case_errors.items())) + failed = { + field: error for field, error in case_errors.items() + if error > args.tolerance + } + if failed: + for field in failed: + if field == "energy": + target = sum(reference["potential"]) + elif field == "virial": + target = [sum(row[index] for row in reference["atom_virial"]) + for index in range(9)] + else: + target = reference[field] + index, actual, expected = worst_difference(candidate[field], target) + print( + f"{key} {field} worst[{index}]: " + f"actual={actual:.10e} expected={expected:.10e}", + flush=True) + raise AssertionError( + f"{key}: tolerance {args.tolerance:.3e} exceeded: {failed}") + + worst = max(error for case in errors.values() for error in case.values()) + print( + f"spin3 oracle validation passed: cases={len(errors)} " + f"max_error={worst:.3e}") + + +if __name__ == "__main__": + main() diff --git a/tests_pytest/validate_spin_multigpu.py b/tests_pytest/validate_spin_multigpu.py new file mode 100644 index 000000000..ab4480c88 --- /dev/null +++ b/tests_pytest/validate_spin_multigpu.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +"""Same-system 1/2/4 GPU NEP-Spin checks; run inside one four-GPU allocation. + +Requires numpy. Keeps inputs, per-case logs, and a JSON error ledger under --output. +The single-GPU executable is an independent pre-decomposition build when --reference is supplied. +""" +import argparse, hashlib, json, os, re, shutil, subprocess, time +from pathlib import Path +import numpy as np + +# Componentwise gates: |candidate-reference| <= atol + rtol*|reference|. +# Energy tolerance is per atom and does not scale with the elemental baseline. +GATES = {'potential': (5e-6, 0), 'force': (5e-5, 1e-5), + 'mforce': (5e-6, 1e-5), 'virial': (5e-5, 1e-5), + 'pos': (2e-5, 0), 'vel': (2e-5, 0), 'spin': (3e-5, 0)} + +def sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + +def frame(path): + with path.open() as f: + while True: + line = f.readline() + if not line: break + n = int(line); header = f.readline() + props = re.search(r'Properties=([^\s]+)', header).group(1).split(':') + raw = [f.readline().split() for _ in range(n)] + offset = 0; result = {} + for i in range(0, len(props), 3): + key, kind, width = props[i], props[i+1], int(props[i+2]) + key = {'forces': 'force', 'energy_atom': 'potential'}.get(key, key) + if kind != 'S': + result[key] = np.array([r[offset:offset+width] for r in raw], dtype=float) + offset += width + return result + +def compare(left, right, keys): + results = {} + for key in keys: + x, y = left[key], right[key] + assert x.shape == y.shape + assert np.isfinite(x).all() and np.isfinite(y).all(), key + err = np.abs(x-y); atol, rtol = GATES[key] + flat = int(np.argmax(err)); idx = np.unravel_index(flat, err.shape) + result = {'max_abs': float(err.max()), 'rms': float(np.sqrt(np.mean(err**2))), + 'max_rel_floor_1e-3': float(np.max(err / np.maximum(np.abs(y), 1e-3))), + 'worst_index': list(map(int, idx)), 'reference': float(y[idx]), + 'candidate': float(x[idx]), 'atol': atol, 'rtol': rtol, + 'exceedances': int(np.sum(err > atol+rtol*np.abs(y)))} + if key in ('potential','virial','force','mforce'): + result['sum_delta'] = (x.sum(axis=0)-y.sum(axis=0)).tolist() + results[key] = result + return results + +def write_xyz(path, positions, spins, cell, velocity=None, species=None): + if velocity is None: velocity = np.zeros_like(positions) + n = len(positions) + with path.open('w') as f: + f.write(str(n)+'\nLattice="'+' '.join(map(str,cell.ravel()))+'" ' + 'Properties=species:S:1:pos:R:3:vel:R:3:spin:R:3:group:I:1 pbc="T T T"\n') + for i,(r,v,s) in enumerate(zip(positions,velocity,spins)): + symbol = species[i] if species is not None else 'Fe' + f.write(symbol+' '+' '.join(f'{x:.15g}' for x in np.r_[r,v,s])+f' {i%2}\n') + +def structure(dims=(48,12,12), variant='bcc'): + rng = np.random.default_rng(813) + a = 2.8402441439 + ijk = np.indices(dims).reshape(3,-1).T + pos = np.vstack((ijk,ijk+0.5))*a + cell = np.diag(np.array(dims)*a) + pos += rng.uniform(-0.012,0.012,pos.shape) + pos %= np.diag(cell) + if variant == 'nonuniform': + keep = (pos[:,0] < cell[0,0]*.42) | (rng.random(len(pos)) < .32) + pos = pos[keep] + if variant == 'triclinic': + frac = pos / np.diag(cell) + cell[1,0] = 5.4; cell[2,1] = 2.8 + pos = frac @ cell + spin = np.tile([0.,0.,2.4],(len(pos),1)) + rng.uniform(-.25,.25,pos.shape) + velocity = rng.normal(0,.0012,pos.shape) + velocity -= velocity.mean(axis=0) + return pos,spin,cell,velocity + +def execute(args, name, xyz, gpus, ensemble='nve', steps=1, dt=0, field='', direction='', reference=False): + path = args.output/name; path.mkdir() + shutil.copyfile(args.model,path/'nep.txt'); shutil.copyfile(xyz,path/'model.xyz') + run = (f'potential nep.txt{(" "+direction) if direction else ""}\n' + f'ensemble {ensemble}\ntime_step {dt}\n'+field+ + f'dump_xyz {steps} state.xyz precision double velocity force potential spin mforce virial\nrun {steps}\n') + (path/'run.in').write_text(run) + env = dict(os.environ); env['CUDA_VISIBLE_DEVICES'] = ','.join(args.devices[:gpus]) + exe = args.reference if reference and args.reference else args.binary + start=time.perf_counter() + with (path/'run.log').open('w') as log: + proc=subprocess.run([str(exe)],cwd=path,env=env,stdout=log,stderr=subprocess.STDOUT) + if proc.returncode: raise RuntimeError(f'{name}: exit {proc.returncode}: {path}/run.log') + return frame(path/'state.xyz'),time.perf_counter()-start + +def main(): + p=argparse.ArgumentParser();p.add_argument('--binary',type=Path,required=True) + p.add_argument('--reference',type=Path);p.add_argument('--model',type=Path,required=True) + p.add_argument('--output',type=Path,required=True) + a=p.parse_args() + for key in ['binary','reference','model','output']: + if getattr(a,key):setattr(a,key,getattr(a,key).resolve()) + a.output.mkdir(parents=True,exist_ok=False) + a.devices=os.environ.get('CUDA_VISIBLE_DEVICES','0,1,2,3').split(',') + if len(a.devices)<4:raise RuntimeError('Requires four allocated GPUs') + report={'gates':GATES,'binary_sha256':sha(a.binary),'model_sha256':sha(a.model), + 'reference_sha256':sha(a.reference) if a.reference else None, + 'cases':{}} + def save(): + (a.output/'report.json').write_text(json.dumps(report,indent=2)) + def record(name,candidate,ref,keys): + result=compare(candidate,ref,keys);report['cases'][name]=result;save() + bad={k:v for k,v in result.items() if v['exceedances']} + if bad:raise AssertionError((name,bad)) + print(name,'PASS',flush=True) + xyz=a.output/'input.xyz' + outputs=['potential','force','mforce','virial'] + original_model = a.model + model = original_model.read_text() + converted = original_model + # Both elements appear on every slab; use a large cell to exercise the + # actual decomposition rather than its small-box fallback. + pos,spin,cell,vel=structure((48,8,8)) + symbols=original_model.read_text().splitlines()[0].split()[2:] + species=np.array([symbols[i % len(symbols)] for i in range(len(pos))]) + write_xyz(xyz,pos,spin,cell,vel,species) + ref,_=execute(a,'format-reference',xyz,1) + reference_line=next(line for line in model.splitlines() if line.startswith('energy_reference ')) + original=np.array(list(map(float,reference_line.split()[1:]))) + shifts=np.array([(-1)**i*(100000.123456789+i*100000.864197532) + for i in range(len(symbols))]) + shifted=a.output/'shifted.nep' + shifted.write_text(model.replace(reference_line,'energy_reference '+ + ' '.join(f'{v:.17g}' for v in original+shifts),1)) + atom_shift=np.array([shifts[i % len(symbols)] for i in range(len(pos))])[:,None] + for g in [1,2,4]: + a.model=converted + out,_=execute(a,f'format-block-g{g}',xyz,g) + record(f'format-block-g{g}',out,ref,outputs) + if g > 1: + repeated,_=execute(a,f'format-repeat-g{g}',xyz,g) + record(f'format-repeat-g{g}',repeated,out,outputs) + a.model=shifted + shifted_out,_=execute(a,f'energy-reference-g{g}',xyz,g) + shifted_out['potential']-=atom_shift + record(f'energy-reference-g{g}',shifted_out,out,outputs) + # Multi-GPU cell-list ordering varies between launches. Keep the + # established componentwise gates above and record a repeat control; + # the single-GPU path also supports a strict derivative check. + if g == 1: + for key in ['force','mforce','virial']: + if np.max(np.abs(shifted_out[key]-out[key]))>1e-10: + raise AssertionError(f'Energy reference changed {key} on {g} GPUs') + a.model = converted + field = 'add_mfield 0 0 100 -200 300\nadd_mfield 0 0 -20 50 10\n' + field_ref, _ = execute(a, 'field-reference', xyz, 1, field=field) + assert np.max(np.abs(field_ref['mforce'] - ref['mforce'])) > 1e-3 + for g in [1, 2, 4]: + out, _ = execute(a, f'field-g{g}', xyz, g, field=field) + record(f'field-g{g}', out, field_ref, outputs) + for variant in ['bcc','nonuniform','triclinic','small','cutoff']: + if variant=='cutoff': + # Isolated pairs just inside/outside the 6 A cutoff, including slab/PBC faces. + cell=np.diag([144.,24.,24.]);pos=[] + for k,x in enumerate([0.,33.,69.,105.]): + for j,eps in enumerate([-1e-5,0,1e-5]): + pos.extend([[x,4+8*j,4],[((x+6+eps)%144),4+8*j,4]]) + pos=np.array(pos);spin=np.tile([.3,.4,2.35],(len(pos),1));vel=np.zeros_like(pos) + else: + pos,spin,cell,vel=structure((2,2,2) if variant=='small' else (48,12,12),variant) + write_xyz(xyz,pos,spin,cell,vel) + ref,_=execute(a,variant+'-reference',xyz,1,reference=True) + for g in [1,2,4]: + out,_=execute(a,f'{variant}-g{g}',xyz,g) + record(f'{variant}-g{g}',out,ref,outputs) + if variant=='bcc': + for d in ['x','y','z']: + out,_=execute(a,f'{variant}-{d}',xyz,4,direction=d) + record(f'{variant}-{d}',out,ref,outputs) + # n_max_angular=8 and ten angular channels need 16*(90+9*80)*4 = 51,840 shared bytes, + # exceeding the 48 KiB direct-force limit and selecting the generic fallback. + # The large cell also keeps every actual multi-GPU domain out of the small-box path. + dim = 5 + 9 * 10 + 55 + count = (dim + 2) * 30 * 2 + 1 + 4 * (5 * 9 + 9 * 9 + 2 * 9) + 16 + rng = np.random.default_rng(9127) + header = model.split("@ parameters", 1)[0].replace("n_max 4 4", "n_max 4 8").replace("l_max 4 2 0", "l_max 8 2 1") + fallback = a.output / "angular-fallback.nep" + fallback.write_text(header + f"@ parameters {count}\n" + + "\n".join(f"{v:.9g}" for v in rng.normal(0, .03, count)) + + f"\n@ q_scaler {dim}\n" + "1\n" * dim) + a.model = fallback + pos, spin, cell, vel = structure((48, 12, 12)) + write_xyz(xyz, pos, spin, cell, vel) + ref, _ = execute(a, 'angular-fallback-reference', xyz, 1) + assert np.max(np.abs(ref['force'])) > 1e-8 + for g in [1, 2, 4]: + out, _ = execute(a, f'angular-fallback-g{g}', xyz, g) + record(f'angular-fallback-g{g}', out, ref, outputs) + a.model = original_model + + # Uniform translation across slab/PBC faces exceeds the cached-partition + # displacement limit twice. Spin and force outputs must follow atom identity. + pos,spin,cell,vel=structure((48,8,8));vel[:]=[.75,0,0] + write_xyz(xyz,pos,spin,cell,vel) + ref,_=execute(a,'migration-reference',xyz,1,'nve',40,.1,reference=True) + for g in [1,2,4]: + out,_=execute(a,f'migration-g{g}',xyz,g,'nve',40,.1) + record(f'migration-g{g}',out,ref,outputs+['pos','vel','spin']) + delta=out['pos']-pos;delta-=np.rint(delta/np.diag(cell))*np.diag(cell) + assert delta[:,0].min()>2.5, 'Migration probe must exceed the partition rebuild limit' + report['passed']=True;save() + +if __name__=='__main__':main()