src/gromacs/mdtypes/threaded_force_buffer.h

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2021, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  *
  37  * \brief This file contains the declaration of ThreadForceBuffer and ThreadedForceBuffer.
  38  *
  39  * These classes provides thread-local force, shift force and energy buffers
  40  * for kernels. These kernels can then run completely independently on
  41  * multiple threads. Their output can be reduced thread-parallel afterwards.
  42  *
  43  * Usage:
  44  *
  45  * At domain decomposition time:
  46  * Each thread calls: ThreadForceBuffer.resizeBufferAndClearMask()
  47  * Each thread calls: ThreadForceBuffer.addAtomToMask() for all atoms used in the buffer
  48  * Each thread calls: ThreadForceBuffer.processMask()
  49  * After that ThreadedForceBuffer.setupReduction() is called
  50  *
  51  * At force computation time:
  52  * Each thread calls: ThreadForceBuffer.clearForcesAndEnergies().
  53  * Each thread can then accumulate forces and energies into the buffers in ThreadForceBuffer.
  54  * After that ThreadedForceBuffer.reduce() is called for thread-parallel reduction.
  55  *
  56  * \author Berk Hess <hess@kth.se>
  57  * \ingroup mdtypes
  58  */
  59 #ifndef GMX_MDTYPES_THREADED_FORCE_BUFFER_H
  60 #define GMX_MDTYPES_THREADED_FORCE_BUFFER_H
  61
  62 #include <memory>
  63
  64 #include "gromacs/math/vectypes.h"
  65 #include "gromacs/mdtypes/enerdata.h"
  66 #include "gromacs/mdtypes/simulation_workload.h"
  67 #include "gromacs/topology/ifunc.h"
  68 #include "gromacs/utility/alignedallocator.h"
  69 #include "gromacs/utility/arrayref.h"
  70 #include "gromacs/utility/bitmask.h"
  71 #include "gromacs/utility/classhelpers.h"
  72
  73 namespace gmx
  74 {
  75
  76 class ForceWithShiftForces;
  77 class StepWorkload;
  78
  79 /*! \internal
  80  * \brief Object that holds force and energies buffers plus a mask for a thread
  81  *
  82  * \tparam ForceBufferElementType  The type for components of the normal force buffer: rvec or rvec4
  83  */
  84 template<typename ForceBufferElementType>
  85 class ThreadForceBuffer
  86 {
  87 public:
  88     /* We reduce the force array in blocks of 2^5 atoms. This is large enough
  89      * to not cause overhead and 32*sizeof(rvec) is a multiple of the cache-line
  90      * size on all systems.
  91      */
  92     //! The log2 of the reduction block size
  93     static constexpr int s_numReductionBlockBits = 5;
  94     //! Force buffer block size in atoms
  95     static constexpr int s_reductionBlockSize = (1 << s_numReductionBlockBits);
  96
  97     /*! \brief Constructor
  98      * \param[in] threadIndex  The index of the thread that will fill the buffers in this object
  99      * \param[in] useEnergyTerms   Whether the list of energy terms will be used
 100      * \param[in] numEnergyGroups  The number of non-bonded energy groups
 101      */
 102     ThreadForceBuffer(int threadIndex, bool useEnergyTerms, int numEnergyGroups);
 103
 104     //! Resizes the buffer to \p numAtoms and clears the mask
 105     void resizeBufferAndClearMask(int numAtoms);
 106
 107     //! Adds atom with index \p atomIndex for reduction
 108     void addAtomToMask(const int atomIndex)
 109     {
 110         bitmask_set_bit(&reductionMask_[atomIndex >> s_numReductionBlockBits], threadIndex_);
 111     }
 112
 113     void processMask();
 114
 115     //! Returns the size of the force buffer in number of atoms
 116     index size() const { return numAtoms_; }
 117
 118     //! Clears all force and energy buffers
 119     void clearForcesAndEnergies();
 120
 121     //! Returns a plain pointer to the force buffer
 122     ForceBufferElementType* forceBuffer()
 123     {
 124         return reinterpret_cast<ForceBufferElementType*>(forceBuffer_.data());
 125     }
 126
 127     //! Returns a view of the shift force buffer
 128     ArrayRef<RVec> shiftForces() { return shiftForces_; }
 129
 130     //! Returns a view of the energy terms, size F_NRE
 131     ArrayRef<real> energyTerms() { return energyTerms_; }
 132
 133     //! Returns a reference to the energy group pair energies
 134     gmx_grppairener_t& groupPairEnergies() { return groupPairEnergies_; }
 135
 136     //! Returns a reference to the dvdl terms
 137     EnumerationArray<FreeEnergyPerturbationCouplingType, real>& dvdl() { return dvdl_; }
 138
 139     //! Returns a const view to the reduction masks
 140     ArrayRef<const gmx_bitmask_t> reductionMask() const { return reductionMask_; }
 141
 142 private:
 143     //! Force array buffer
 144     std::vector<real, AlignedAllocator<real>> forceBuffer_;
 145     //! Mask for marking which parts of f are filled, working array for constructing mask in bonded_threading_t
 146     std::vector<gmx_bitmask_t> reductionMask_;
 147     //! Index to touched blocks
 148     std::vector<int> usedBlockIndices_;
 149     //! The index of our thread
 150     int threadIndex_;
 151     //! The number of atoms in the buffer
 152     int numAtoms_ = 0;
 153
 154     //! Shift force array, size c_numShiftVectors
 155     std::vector<RVec> shiftForces_;
 156     //! Energy array, can be empty
 157     std::vector<real> energyTerms_;
 158     //! Group pair energy data for pairs
 159     gmx_grppairener_t groupPairEnergies_;
 160     //! Free-energy dV/dl output
 161     gmx::EnumerationArray<FreeEnergyPerturbationCouplingType, real> dvdl_;
 162
 163     // Disallow copy and assign, remove this we we get rid of f_
 164     GMX_DISALLOW_COPY_MOVE_AND_ASSIGN(ThreadForceBuffer);
 165 };
 166
 167 /*! \internal
 168  * \brief Class for accumulating and reducing forces and energies on threads in parallel
 169  *
 170  * \tparam ForceBufferElementType  The type for components of the normal force buffer: rvec or rvec4
 171  */
 172 template<typename ForceBufferElementType>
 173 class ThreadedForceBuffer
 174 {
 175 public:
 176     /*! \brief Constructor
 177      * \param[in] numThreads       The number of threads that will use the buffers and reduce
 178      * \param[in] useEnergyTerms   Whether the list of energy terms will be used
 179      * \param[in] numEnergyGroups  The number of non-bonded energy groups
 180      */
 181     ThreadedForceBuffer(int numThreads, bool useEnergyTerms, int numEnergyGroups);
 182
 183     //! Returns the number of thread buffers
 184     int numThreadBuffers() const { return threadForceBuffers_.size(); }
 185
 186     //! Returns a reference to the buffer object for the thread with index \p bufferIndex
 187     ThreadForceBuffer<ForceBufferElementType>& threadForceBuffer(int bufferIndex)
 188     {
 189         return *threadForceBuffers_[bufferIndex];
 190     }
 191
 192     //! Sets up the reduction, should be called after generating the masks on each thread
 193     void setupReduction();
 194
 195     /*! \brief Reduces forces and energies, as requested by \p stepWork
 196      *
 197      * The reduction of all output starts at the output from thread \p reductionBeginIndex,
 198      * except for the normal force buffer, which always starts at 0.
 199      *
 200      * Buffers that will not be used as indicated by the flags in \p stepWork
 201      * are allowed to be nullptr or empty.
 202      */
 203     void reduce(gmx::ForceWithShiftForces* forceWithShiftForces,
 204                 real*                      ener,
 205                 gmx_grppairener_t*         grpp,
 206                 gmx::ArrayRef<real>        dvdl,
 207                 const gmx::StepWorkload&   stepWork,
 208                 int                        reductionBeginIndex);
 209
 210 private:
 211     //! Whether the energy buffer is used
 212     bool useEnergyTerms_;
 213     //! Force/energy data per thread, size nthreads, stored in unique_ptr to allow thread local allocation
 214     std::vector<std::unique_ptr<ThreadForceBuffer<ForceBufferElementType>>> threadForceBuffers_;
 215     //! Indices of blocks that are used, i.e. have force contributions.
 216     std::vector<int> usedBlockIndices_;
 217     //! Mask array, one element corresponds to a block of reduction_block_size atoms of the force array, bit corresponding to thread indices set if a thread writes to that block
 218     std::vector<gmx_bitmask_t> reductionMask_;
 219     //! The number of atoms forces are computed for
 220     int numAtomsForce_ = 0;
 221
 222     // Disallow copies to avoid sub-optimal ownership of allocated memory
 223     GMX_DISALLOW_COPY_MOVE_AND_ASSIGN(ThreadedForceBuffer);
 224 };
 225
 226 // Instantiate for RVec
 227 extern template class ThreadForceBuffer<RVec>;
 228 extern template class ThreadedForceBuffer<RVec>;
 229
 230 // Instantiate for rvec4
 231 extern template class ThreadForceBuffer<rvec4>;
 232 extern template class ThreadedForceBuffer<rvec4>;
 233
 234 } // namespace gmx
 235
 236 #endif