From a83aaefedaf27140af420599183cbb2a98300be3 Mon Sep 17 00:00:00 2001 From: Berk Hess Date: Sun, 12 Nov 2017 13:10:54 +0100 Subject: [PATCH] Add SIMD for AWH When using the convolved potential with AWH, a large number of double precision exp() functions need to be evaluated at every step. These are now SIMD accelerated. Change-Id: If1e3a916469c4fd7e26740123009ae59b7927667 --- docs/user-guide/mdp-options.rst | 4 +- src/gromacs/awh/bias.cpp | 12 ++--- src/gromacs/awh/bias.h | 4 +- src/gromacs/awh/biasstate.cpp | 102 +++++++++++++++++++++++++--------------- src/gromacs/awh/biasstate.h | 24 +++++----- src/gromacs/awh/coordstate.cpp | 12 ++--- src/gromacs/awh/coordstate.h | 14 +++--- 7 files changed, 103 insertions(+), 69 deletions(-) diff --git a/docs/user-guide/mdp-options.rst b/docs/user-guide/mdp-options.rst index f708d8eba9..f659803776 100644 --- a/docs/user-guide/mdp-options.rst +++ b/docs/user-guide/mdp-options.rst @@ -1858,8 +1858,8 @@ AWH adaptive biasing using Monte-Carlo sampling. The force constant is set with :mdp:`awh1-dim1-force-constant`. The umbrella location is sampled using Monte-Carlo every :mdp:`awh-nstsample` steps. - This option can be useful for cases when calculating the convolved force for each step becomes - computationally expensive. + There are no advantages to using an umbrella. + This option is mainly for comparison and testing purposes. .. mdp:: awh-share-multisim diff --git a/src/gromacs/awh/bias.cpp b/src/gromacs/awh/bias.cpp index 8a314fdee6..36f3af0e42 100644 --- a/src/gromacs/awh/bias.cpp +++ b/src/gromacs/awh/bias.cpp @@ -113,7 +113,7 @@ void Bias::calcForceAndUpdateBias(const awh_dvec coordValue, state_.setCoordValue(grid_, coordValue); - std::vector *probWeightNeighbor = &tempWorkSpace_; + std::vector < double, AlignedAllocator < double>> &probWeightNeighbor = alignedTempWorkSpace_; /* If the convolved force is needed or this is a sampling step, * the bias in the current neighborhood needs to be up-to-date @@ -129,11 +129,11 @@ void Bias::calcForceAndUpdateBias(const awh_dvec coordValue, state_.doSkippedUpdatesInNeighborhood(params_, grid_); } - convolvedBias = state_.updateProbabilityWeightsAndConvolvedBias(dimParams_, grid_, probWeightNeighbor); + convolvedBias = state_.updateProbabilityWeightsAndConvolvedBias(dimParams_, grid_, &probWeightNeighbor); if (sampleCoord) { - state_.sampleCoordAndPmf(grid_, *probWeightNeighbor, convolvedBias); + state_.sampleCoordAndPmf(grid_, probWeightNeighbor, convolvedBias); } } @@ -149,7 +149,7 @@ void Bias::calcForceAndUpdateBias(const awh_dvec coordValue, double potential; if (params_.convolveForce) { - state_.calcConvolvedForce(dimParams_, grid_, *probWeightNeighbor, + state_.calcConvolvedForce(dimParams_, grid_, probWeightNeighbor, biasForce); potential = -convolvedBias*params_.invBeta; @@ -169,7 +169,7 @@ void Bias::calcForceAndUpdateBias(const awh_dvec coordValue, */ if (moveUmbrella) { - double newPotential = state_.moveUmbrella(dimParams_, grid_, *probWeightNeighbor, biasForce, step, seed, params_.biasIndex); + double newPotential = state_.moveUmbrella(dimParams_, grid_, probWeightNeighbor, biasForce, step, seed, params_.biasIndex); *potentialJump = newPotential - potential; } } @@ -243,7 +243,7 @@ Bias::Bias(int biasIndexInCollection, params_(awhParams, awhBiasParams, dimParams_, beta, mdTimeStep, disableUpdateSkips, numSharingSimulations, grid_.axis(), biasIndexInCollection), state_(awhBiasParams, params_.initialHistogramSize, dimParams_, grid_), thisRankDoesIO_(thisRankWillDoIO == ThisRankWillDoIO::Yes), - tempWorkSpace_(), + alignedTempWorkSpace_(), numWarningsIssued_(0) { /* For a global update updateList covers all points, so reserve that */ diff --git a/src/gromacs/awh/bias.h b/src/gromacs/awh/bias.h index 4d5625b0b0..95db7d4c61 100644 --- a/src/gromacs/awh/bias.h +++ b/src/gromacs/awh/bias.h @@ -55,6 +55,7 @@ #include #include "gromacs/math/vectypes.h" +#include "gromacs/utility/alignedallocator.h" #include "gromacs/utility/basedefinitions.h" #include "gromacs/utility/gmxassert.h" @@ -303,8 +304,9 @@ class Bias /* Temporary working vector used during the update. * This only here to avoid allocation at every MD step. */ - std::vector tempWorkSpace_; /**< Working vector of doubles. */ + std::vector < double, AlignedAllocator < double>> alignedTempWorkSpace_; /**< Working vector of doubles. */ + /* Run-local counter to avoid flooding log with warnings. */ int numWarningsIssued_; /**< The number of warning issued in the current run. */ }; diff --git a/src/gromacs/awh/biasstate.cpp b/src/gromacs/awh/biasstate.cpp index 5cbf18a069..a57efb3c60 100644 --- a/src/gromacs/awh/biasstate.cpp +++ b/src/gromacs/awh/biasstate.cpp @@ -62,6 +62,8 @@ #include "gromacs/mdtypes/awh-history.h" #include "gromacs/mdtypes/awh-params.h" #include "gromacs/mdtypes/commrec.h" +#include "gromacs/simd/simd.h" +#include "gromacs/simd/simd_math.h" #include "gromacs/utility/arrayref.h" #include "gromacs/utility/exceptions.h" #include "gromacs/utility/gmxassert.h" @@ -152,7 +154,7 @@ double freeEnergyMinimumValue(gmx::ArrayRef pointState) } /*! \brief - * Find the probability weight of a point given a coordinate value. + * Find and return the log of the probability weight of a point given a coordinate value. * * The unnormalized weight is given by * w(point|value) = exp(bias(point) - U(value,point)), @@ -164,21 +166,21 @@ double freeEnergyMinimumValue(gmx::ArrayRef pointState) * \param[in] pointIndex Point to evaluate probability weight for. * \param[in] pointBias Bias for the point (as a log weight). * \param[in] value Coordinate value. - * \returns the biased probability weight. + * \returns the log of the biased probability weight. */ -double biasedWeightFromPoint(const std::vector &dimParams, - const std::vector &points, - const Grid &grid, - int pointIndex, - double pointBias, - const awh_dvec value) +double biasedLogWeightFromPoint(const std::vector &dimParams, + const std::vector &points, + const Grid &grid, + int pointIndex, + double pointBias, + const awh_dvec value) { - double weight = 0; + double logWeight = c_largeNegativeExponent; /* Only points in the target reigon have non-zero weight */ if (points[pointIndex].inTargetRegion()) { - double logWeight = pointBias; + logWeight = pointBias; /* Add potential for all parameter dimensions */ for (size_t d = 0; d < dimParams.size(); d++) @@ -186,11 +188,9 @@ double biasedWeightFromPoint(const std::vector &dimParams, double dev = getDeviationFromPointAlongGridAxis(grid, d, pointIndex, value[d]); logWeight -= 0.5*dimParams[d].betak*dev*dev; } - - weight = std::exp(logWeight); } - return weight; + return logWeight; } } // namespace @@ -219,9 +219,10 @@ void BiasState::calcConvolvedPmf(const std::vector &dimParams, /* Add the convolved PMF weights for the neighbors of this point. Note that this function only adds point within the target > 0 region. Sum weights, take the logarithm last to get the free energy. */ - freeEnergyWeights += biasedWeightFromPoint(dimParams, points_, grid, - neighbor, biasNeighbor, - point.coordValue); + double logWeight = biasedLogWeightFromPoint(dimParams, points_, grid, + neighbor, biasNeighbor, + point.coordValue); + freeEnergyWeights += std::exp(logWeight); } GMX_RELEASE_ASSERT(freeEnergyWeights > 0, "Attempting to do log(<= 0) in AWH convolved PMF calculation."); @@ -392,7 +393,7 @@ double BiasState::calcUmbrellaForceAndPotential(const std::vector &di void BiasState::calcConvolvedForce(const std::vector &dimParams, const Grid &grid, - const std::vector &probWeightNeighbor, + gmx::ArrayRef probWeightNeighbor, awh_dvec force) const { for (size_t d = 0; d < dimParams.size(); d++) @@ -422,7 +423,7 @@ void BiasState::calcConvolvedForce(const std::vector &dimParams, double BiasState::moveUmbrella(const std::vector &dimParams, const Grid &grid, - const std::vector &probWeightNeighbor, + gmx::ArrayRef probWeightNeighbor, awh_dvec biasForce, gmx_int64_t step, gmx_int64_t seed, @@ -1095,24 +1096,50 @@ void BiasState::updateFreeEnergyAndAddSamplesToHistogram(const std::vector &dimParams, - const Grid &grid, - std::vector *weight) const +double BiasState::updateProbabilityWeightsAndConvolvedBias(const std::vector &dimParams, + const Grid &grid, + std::vector < double, AlignedAllocator < double>> *weight) const { /* Only neighbors of the current coordinate value will have a non-negligible chance of getting sampled */ const std::vector &neighbors = grid.point(coordState_.gridpointIndex()).neighbor; - weight->resize(neighbors.size()); +#if GMX_SIMD_HAVE_DOUBLE + typedef SimdDouble PackType; + constexpr int packSize = GMX_SIMD_DOUBLE_WIDTH; +#else + typedef double PackType; + constexpr int packSize = 1; +#endif + /* Round the size of the weight array up to packSize */ + const int weightSize = ((neighbors.size() + packSize - 1)/packSize)*packSize; + weight->resize(weightSize); - double weightSum = 0; - for (size_t n = 0; n < neighbors.size(); n++) + double * gmx_restrict weightData = weight->data(); + PackType weightSumPack(0.0); + for (size_t i = 0; i < neighbors.size(); i += packSize) { - int neighbor = neighbors[n]; - (*weight)[n] = biasedWeightFromPoint(dimParams, points_, grid, - neighbor, points_[neighbor].bias(), - coordState_.coordValue()); - weightSum += (*weight)[n]; + for (size_t n = i; n < i + packSize; n++) + { + if (n < neighbors.size()) + { + const int neighbor = neighbors[n]; + (*weight)[n] = biasedLogWeightFromPoint(dimParams, points_, grid, + neighbor, points_[neighbor].bias(), + coordState_.coordValue()); + } + else + { + /* Pad with values that don't affect the result */ + (*weight)[n] = c_largeNegativeExponent; + } + } + PackType weightPack = load(weightData + i); + weightPack = gmx::exp(weightPack); + weightSumPack = weightSumPack + weightPack; + store(weightData + i, weightPack); } + /* Sum of probability weights */ + double weightSum = reduce(weightSumPack); GMX_RELEASE_ASSERT(weightSum > 0, "zero probability weight when updating AWH probability weights."); /* Normalize probabilities to sum to 1 */ @@ -1137,17 +1164,18 @@ double BiasState::calcConvolvedBias(const std::vector &dimParams, double weightSum = 0; for (int neighbor : gridPoint.neighbor) { - weightSum += biasedWeightFromPoint(dimParams, points_, grid, - neighbor, points_[neighbor].bias(), - coordValue); + double logWeight = biasedLogWeightFromPoint(dimParams, points_, grid, + neighbor, points_[neighbor].bias(), + coordValue); + weightSum += std::exp(logWeight); } /* Returns -GMX_DOUBLE_MAX if no neighboring points were in the target region. */ return (weightSum > 0) ? std::log(weightSum) : -GMX_DOUBLE_MAX; } -void BiasState::sampleProbabilityWeights(const Grid &grid, - const std::vector &probWeightNeighbor) +void BiasState::sampleProbabilityWeights(const Grid &grid, + gmx::ArrayRef probWeightNeighbor) { const std::vector &neighbor = grid.point(coordState_.gridpointIndex()).neighbor; @@ -1201,9 +1229,9 @@ void BiasState::sampleProbabilityWeights(const Grid &grid, } } -void BiasState::sampleCoordAndPmf(const Grid &grid, - const std::vector &probWeightNeighbor, - double convolvedBias) +void BiasState::sampleCoordAndPmf(const Grid &grid, + gmx::ArrayRef probWeightNeighbor, + double convolvedBias) { /* Sampling-based deconvolution extracting the PMF. * Update the PMF histogram with the current coordinate value. diff --git a/src/gromacs/awh/biasstate.h b/src/gromacs/awh/biasstate.h index 27ceb0d3de..1a541c606c 100644 --- a/src/gromacs/awh/biasstate.h +++ b/src/gromacs/awh/biasstate.h @@ -57,6 +57,8 @@ #include #include "gromacs/math/vectypes.h" +#include "gromacs/utility/alignedallocator.h" +#include "gromacs/utility/arrayref.h" #include "gromacs/utility/basedefinitions.h" #include "gromacs/utility/gmxassert.h" @@ -239,7 +241,7 @@ class BiasState */ void calcConvolvedForce(const std::vector &dimParams, const Grid &grid, - const std::vector &probWeightNeighbor, + gmx::ArrayRef probWeightNeighbor, awh_dvec force) const; /*! \brief @@ -262,7 +264,7 @@ class BiasState */ double moveUmbrella(const std::vector &dimParams, const Grid &grid, - const std::vector &probWeightNeighbor, + gmx::ArrayRef probWeightNeighbor, awh_dvec biasForce, gmx_int64_t step, gmx_int64_t seed, @@ -420,13 +422,13 @@ class BiasState * * \param[in] dimParams The bias dimensions parameters * \param[in] grid The grid. - * \param[out] weight Probability weights of the neighbors. + * \param[out] weight Probability weights of the neighbors, SIMD aligned. * \returns the convolved bias. */ - double updateProbabilityWeightsAndConvolvedBias(const std::vector &dimParams, - const Grid &grid, - std::vector *weight) const; + double updateProbabilityWeightsAndConvolvedBias(const std::vector &dimParams, + const Grid &grid, + std::vector < double, AlignedAllocator < double>> *weight) const; /*! \brief * Take samples of the current probability weights for future updates and analysis. @@ -438,8 +440,8 @@ class BiasState * \param[in] grid The grid. * \param[in] probWeightNeighbor Probability weights of the neighbors. */ - void sampleProbabilityWeights(const Grid &grid, - const std::vector &probWeightNeighbor); + void sampleProbabilityWeights(const Grid &grid, + gmx::ArrayRef probWeightNeighbor); /*! \brief * Sample the reaction coordinate and PMF for future updates or analysis. @@ -451,9 +453,9 @@ class BiasState * \param[in] probWeightNeighbor Probability weights of the neighbors. * \param[in] convolvedBias The convolved bias. */ - void sampleCoordAndPmf(const Grid &grid, - const std::vector &probWeightNeighbor, - double convolvedBias); + void sampleCoordAndPmf(const Grid &grid, + gmx::ArrayRef probWeightNeighbor, + double convolvedBias); /*! \brief * Calculates the convolved bias for a given coordinate value. * diff --git a/src/gromacs/awh/coordstate.cpp b/src/gromacs/awh/coordstate.cpp index 108b50c596..e161ba4edc 100644 --- a/src/gromacs/awh/coordstate.cpp +++ b/src/gromacs/awh/coordstate.cpp @@ -126,12 +126,12 @@ int getSampleFromDistribution(ArrayRef distr, } // namespace void -CoordState::sampleUmbrellaGridpoint(const Grid &grid, - int gridpointIndex, - const std::vector &probWeightNeighbor, - gmx_int64_t step, - gmx_int64_t seed, - int indexSeed) +CoordState::sampleUmbrellaGridpoint(const Grid &grid, + int gridpointIndex, + gmx::ArrayRef probWeightNeighbor, + gmx_int64_t step, + gmx_int64_t seed, + int indexSeed) { /* Sample new umbrella reference value from the probability distribution * which is defined for the neighboring points of the current coordinate. diff --git a/src/gromacs/awh/coordstate.h b/src/gromacs/awh/coordstate.h index b6a9c8f07e..3df240ac1d 100644 --- a/src/gromacs/awh/coordstate.h +++ b/src/gromacs/awh/coordstate.h @@ -54,6 +54,8 @@ #include +#include "gromacs/utility/arrayref.h" + #include "dimparams.h" namespace gmx @@ -92,12 +94,12 @@ class CoordState * \param[in] indexSeed Second random seed, should be the bias Index. * \returns the index of the sampled point. */ - void sampleUmbrellaGridpoint(const Grid &grid, - int gridpointIndex, - const std::vector &probWeightNeighbor, - gmx_int64_t step, - gmx_int64_t seed, - int indexSeed); + void sampleUmbrellaGridpoint(const Grid &grid, + int gridpointIndex, + gmx::ArrayRef probWeightNeighbor, + gmx_int64_t step, + gmx_int64_t seed, + int indexSeed); /*! \brief Update the coordinate value with coordValue. * -- 2.11.4.GIT