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/idef.h"
  68 #include "gromacs/topology/ifunc.h"
  69 #include "gromacs/utility/alignedallocator.h"
  70 #include "gromacs/utility/arrayref.h"
  71 #include "gromacs/utility/bitmask.h"
  72 #include "gromacs/utility/classhelpers.h"
  73
  74 namespace gmx
  75 {
  76
  77 class ForceWithShiftForces;
  78 class StepWorkload;
  79
  80 /*! \internal
  81  * \brief Object that holds force and energies buffers plus a mask for a thread
  82  *
  83  * \tparam ForceBufferElementType  The type for components of the normal force buffer: rvec or rvec4
  84  */
  85 template<typename ForceBufferElementType>
  86 class ThreadForceBuffer
  87 {
  88 public:
  89     /* We reduce the force array in blocks of 2^5 atoms. This is large enough
  90      * to not cause overhead and 32*sizeof(rvec) is a multiple of the cache-line
  91      * size on all systems.
  92      */
  93     //! The log2 of the reduction block size
  94     static constexpr int s_numReductionBlockBits = 5;
  95     //! Force buffer block size in atoms
  96     static constexpr int s_reductionBlockSize = (1 << s_numReductionBlockBits);
  97
  98     //! Constructor
  99     ThreadForceBuffer(int threadIndex, int numEnergyGroups);
 100
 101     //! Resizes the buffer to \p numAtoms and clears the mask
 102     void resizeBufferAndClearMask(int numAtoms);
 103
 104     //! Adds atom with index \p atomIndex for reduction
 105     void addAtomToMask(const int atomIndex)
 106     {
 107         bitmask_set_bit(&reductionMask_[atomIndex >> s_numReductionBlockBits], threadIndex_);
 108     }
 109
 110     void processMask();
 111
 112     //! Returns the size of the force buffer in number of atoms
 113     index size() const { return numAtoms_; }
 114
 115     //! Clears all force and energy buffers
 116     void clearForcesAndEnergies();
 117
 118     //! Returns a plain pointer to the force buffer
 119     ForceBufferElementType* forceBuffer()
 120     {
 121         return reinterpret_cast<ForceBufferElementType*>(forceBuffer_.data());
 122     }
 123
 124     //! Returns a view of the shift force buffer
 125     ArrayRef<RVec> shiftForces() { return shiftForces_; }
 126
 127     //! Returns a view of the energy terms, size F_NRE
 128     ArrayRef<real> energyTerms() { return energyTerms_; }
 129
 130     //! Returns a reference to the energy group pair energies
 131     gmx_grppairener_t& groupPairEnergies() { return groupPairEnergies_; }
 132
 133     //! Returns a reference to the dvdl terms
 134     EnumerationArray<FreeEnergyPerturbationCouplingType, real>& dvdl() { return dvdl_; }
 135
 136     //! Returns a const view to the reduction masks
 137     ArrayRef<const gmx_bitmask_t> reductionMask() const { return reductionMask_; }
 138
 139 private:
 140     //! Force array buffer
 141     std::vector<real, AlignedAllocator<real>> forceBuffer_;
 142     //! Mask for marking which parts of f are filled, working array for constructing mask in bonded_threading_t
 143     std::vector<gmx_bitmask_t> reductionMask_;
 144     //! Index to touched blocks
 145     std::vector<int> usedBlockIndices_;
 146     //! The index of our thread
 147     int threadIndex_;
 148     //! The number of atoms in the buffer
 149     int numAtoms_ = 0;
 150
 151     //! Shift force array, size c_numShiftVectors
 152     std::vector<RVec> shiftForces_;
 153     //! Energy array
 154     std::array<real, F_NRE> energyTerms_;
 155     //! Group pair energy data for pairs
 156     gmx_grppairener_t groupPairEnergies_;
 157     //! Free-energy dV/dl output
 158     gmx::EnumerationArray<FreeEnergyPerturbationCouplingType, real> dvdl_;
 159
 160     // Disallow copy and assign, remove this we we get rid of f_
 161     GMX_DISALLOW_COPY_MOVE_AND_ASSIGN(ThreadForceBuffer);
 162 };
 163
 164 /*! \internal
 165  * \brief Class for accumulating and reducing forces and energies on threads in parallel
 166  *
 167  * \tparam ForceBufferElementType  The type for components of the normal force buffer: rvec or rvec4
 168  */
 169 template<typename ForceBufferElementType>
 170 class ThreadedForceBuffer
 171 {
 172 public:
 173     //! Constructor
 174     ThreadedForceBuffer(int numThreads, int numEnergyGroups);
 175
 176     //! Returns the number of thread buffers
 177     int numThreadBuffers() const { return threadForceBuffers_.size(); }
 178
 179     //! Returns a reference to the buffer object for the thread with index \p bufferIndex
 180     ThreadForceBuffer<ForceBufferElementType>& threadForceBuffer(int bufferIndex)
 181     {
 182         return *threadForceBuffers_[bufferIndex];
 183     }
 184
 185     //! Sets up the reduction, should be called after generating the masks on each thread
 186     void setupReduction();
 187
 188     /*! \brief Reduces forces and energies, as requested by \p stepWork
 189      *
 190      * The reduction of all output starts at the output from thread \p reductionBeginIndex,
 191      * except for the normal force buffer, which always starts at 0.
 192      */
 193     void reduce(gmx::ForceWithShiftForces* forceWithShiftForces,
 194                 real*                      ener,
 195                 gmx_grppairener_t*         grpp,
 196                 gmx::ArrayRef<real>        dvdl,
 197                 const gmx::StepWorkload&   stepWork,
 198                 int                        reductionBeginIndex);
 199
 200 private:
 201     //! Force/energy data per thread, size nthreads, stored in unique_ptr to allow thread local allocation
 202     std::vector<std::unique_ptr<ThreadForceBuffer<ForceBufferElementType>>> threadForceBuffers_;
 203     //! Indices of blocks that are used, i.e. have force contributions.
 204     std::vector<int> usedBlockIndices_;
 205     //! 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
 206     std::vector<gmx_bitmask_t> reductionMask_;
 207     //! The number of atoms forces are computed for
 208     int numAtomsForce_ = 0;
 209
 210     // Disallow copies to avoid sub-optimal ownership of allocated memory
 211     GMX_DISALLOW_COPY_MOVE_AND_ASSIGN(ThreadedForceBuffer);
 212 };
 213
 214 // Instantiate for rvec4. Can also be instantiated for rvec.
 215 extern template class ThreadForceBuffer<rvec4>;
 216 extern template class ThreadedForceBuffer<rvec4>;
 217
 218 } // namespace gmx
 219
 220 #endif