Line data    Source code

       1             : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
       2             :    Copyright (c) 2011-2023 The plumed team
       3             :    (see the PEOPLE file at the root of the distribution for a list of names)
       4             : 
       5             :    See http://www.plumed.org for more information.
       6             : 
       7             :    This file is part of plumed, version 2.
       8             : 
       9             :    plumed is free software: you can redistribute it and/or modify
      10             :    it under the terms of the GNU Lesser General Public License as published by
      11             :    the Free Software Foundation, either version 3 of the License, or
      12             :    (at your option) any later version.
      13             : 
      14             :    plumed 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
      17             :    GNU Lesser General Public License for more details.
      18             : 
      19             :    You should have received a copy of the GNU Lesser General Public License
      20             :    along with plumed.  If not, see <http://www.gnu.org/licenses/>.
      21             : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
      22             : #include "ActionWithMatrix.h"
      23             : #include "tools/Communicator.h"
      24             : #include "tools/OpenMP.h"
      25             : 
      26             : namespace PLMD {
      27             : 
      28        1711 : void ActionWithMatrix::registerKeywords( Keywords& keys ) {
      29        1711 :   ActionWithVector::registerKeywords( keys );
      30        1711 :   keys.addFlag("AVOID_THREAD_GATHER",false,"avoid race conditions in shared memory by doing more calculations rather than using more memory");
      31        1711 :   keys.addFlag("DO_THREAD_GATHER",false,"do fewer calculations and use more memory");
      32        1711 : }
      33             : 
      34         544 : ActionWithMatrix::ActionWithMatrix(const ActionOptions&ao):
      35             :   Action(ao),
      36             :   ActionWithVector(ao),
      37         544 :   maxcolsize(0),
      38         544 :   diagzero(false),
      39         544 :   gatherForceOnColumns(false) {
      40             : 
      41         544 :   if( keywords.exists("ELEMENTS_ON_DIAGONAL_ARE_ZERO") ) {
      42         176 :     parseFlag("ELEMENTS_ON_DIAGONAL_ARE_ZERO",diagzero);
      43         176 :     if( diagzero ) {
      44          28 :       log.printf("  setting diagonal elements equal to zero\n");
      45             :     }
      46             :   }
      47         544 :   if( getName().find("ACC")!=std::string::npos ) {
      48             :     bool ignore;
      49           0 :     parseFlag("AVOID_THREAD_GATHER",ignore);
      50           0 :     if( ignore ) {
      51           0 :       warning("AVOID_THREAD_GATHER flag is always activated when the calculation is run on the GPU as the calculation is much faster with this option activated");
      52             :     }
      53             :     bool dothreads;
      54           0 :     parseFlag("DO_THREAD_GATHER",dothreads);
      55           0 :     if( !dothreads ) {
      56           0 :       no_thread_gather=true;
      57             :     } else {
      58           0 :       no_thread_gather=false;
      59             :     }
      60             :   } else {
      61         544 :     parseFlag("AVOID_THREAD_GATHER",no_thread_gather);
      62         544 :     if( OpenMP::getNumThreads()==1 ) {
      63           0 :       no_thread_gather = false;
      64             :     }
      65             :     bool ignore;
      66         544 :     parseFlag("DO_THREAD_GATHER",ignore);
      67         544 :     if( ignore ) {
      68           0 :       warning("NO_THREAD_GATHER flag is always activated when the calculation is run on the CPU as the calculation is faster with this option activated");
      69             :     }
      70             :   }
      71         544 :   if( no_thread_gather && comm.Get_size()>1 ) {
      72           0 :     error("AVOID_THREAD_GATHER keyword is incompatible with MPI - it should be possible to fix this. Email: gareth.tribello@gmail.com if you are interested");
      73             :   }
      74         544 :   if( no_thread_gather ) {
      75          15 :     log.printf("  turning on low memory implementation for force gathering\n");
      76             :   }
      77         544 : }
      78             : 
      79             : 
      80             : class RequiredMatrixElementsUpdater {
      81             :   RequiredMatrixElements& outmat;
      82             : public:
      83             :   RequiredMatrixElementsUpdater( RequiredMatrixElements& mat ) : outmat(mat) {}
      84             :   ~RequiredMatrixElementsUpdater() {
      85             :     outmat.update();
      86             :   }
      87             : };
      88             : 
      89       13058 : void ActionWithMatrix::copyMatrixBookeepingFromFirstComponent( RequiredMatrixElements& outmat ) {
      90             :   RequiredMatrixElementsUpdater updater(outmat);
      91       13058 :   Value* mycomp = getPntrToComponent(0);
      92       13058 :   outmat.ncols = mycomp->getNumberOfColumns();
      93             :   outmat.resize( mycomp->matrix_bookeeping.size() );
      94    48633647 :   for(unsigned i=0; i<outmat.size(); ++i) {
      95    48620589 :     outmat[i] = mycomp->matrix_bookeeping[i];
      96             :   }
      97       13058 :   if( !no_thread_gather || doNotCalculateDerivatives() ) {
      98             :     return;
      99             :   }
     100          62 :   std::vector<unsigned> column_totals( getConstPntrToComponent(0)->getShape()[1], 0 );
     101         698 :   for(unsigned i=0; i<mycomp->getShape()[0]; ++i) {
     102       17280 :     for(unsigned j=0; j<mycomp->getRowLength(i); ++j) {
     103       16644 :       column_totals[ mycomp->getRowIndex(i,j) ]++;
     104             :     }
     105             :   }
     106             :   unsigned n_nonzerocols = 0;
     107          62 :   maxcolsize = column_totals[0];
     108          62 :   if( maxcolsize>0 ) {
     109             :     n_nonzerocols = 1;
     110             :   }
     111        2232 :   for(unsigned i=1; i<column_totals.size(); ++i) {
     112        2170 :     if( column_totals[i]>maxcolsize ) {
     113           5 :       maxcolsize = column_totals[i];
     114             :     }
     115        2170 :     if( column_totals[i]>0 ) {
     116         873 :       n_nonzerocols++;
     117             :     }
     118             :   }
     119          62 :   column_list.resize( n_nonzerocols );
     120             :   n_nonzerocols = 0;
     121        2294 :   for(unsigned i=0; i<column_totals.size(); ++i) {
     122        2232 :     if( column_totals[i]>0 ) {
     123         935 :       column_list[n_nonzerocols]=i;
     124         935 :       n_nonzerocols++;
     125             :     }
     126             :   }
     127             : }
     128             : 
     129         465 : void ActionWithMatrix::updateBookeepingArrays( RequiredMatrixElements& outmat ) {
     130         465 :   Value* myval = getPntrToComponent(0);
     131         465 :   unsigned lstart = myval->getShape()[0];
     132         465 :   if( getNumberOfMasks()>0 ) {
     133         184 :     Value* maskarg = getPntrToArgument(getNumberOfArguments()-getNumberOfMasks());
     134         449 :     for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     135         265 :       getPntrToComponent(i)->reshapeMatrixStore( maskarg->getNumberOfColumns() );
     136             :     }
     137      262759 :     for(unsigned i=0; i<lstart; ++i) {
     138      262575 :       unsigned rstart = i*(1+myval->getNumberOfColumns());
     139      262575 :       myval->setMatrixBookeepingElement( rstart, maskarg->getRowLength(i) );
     140    10879714 :       for(unsigned j=0; j<maskarg->getRowLength(i); ++j) {
     141    10354564 :         myval->setMatrixBookeepingElement( rstart + 1 + j, maskarg->getRowIndex(i, j) );
     142             :       }
     143             :     }
     144         281 :   } else if ( diagzero ) {
     145         420 :     for(unsigned i=0; i<getNumberOfArguments(); ++i) {
     146         280 :       if( getPntrToArgument(i)->getRank()==2 && getPntrToArgument(i)->getShape()[1]!=getPntrToArgument(i)->getNumberOfColumns() ) {
     147           0 :         error("DIAGZERO flag is incompatible with sparse matrices");
     148             :       }
     149             :     }
     150         280 :     for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     151         140 :       getPntrToComponent(i)->reshapeMatrixStore( myval->getShape()[1]-1 );
     152             :     }
     153        2900 :     for(unsigned i=0; i<lstart; ++i) {
     154        2760 :       unsigned k=0, rstart = i*(1+myval->getNumberOfColumns());
     155        2760 :       myval->setMatrixBookeepingElement( rstart, myval->getShape()[1]-1 );
     156      204120 :       for(unsigned j=0; j<myval->getShape()[1]; ++j) {
     157      201360 :         if( i!=j ) {
     158      198600 :           myval->setMatrixBookeepingElement( rstart + 1 + k, j );
     159      198600 :           k++;
     160             :         }
     161             :       }
     162             :     }
     163             :   } else {
     164         300 :     for(unsigned i=0; i<getNumberOfComponents(); ++i) {
     165         159 :       getPntrToComponent(i)->reshapeMatrixStore( myval->getShape()[1] );
     166             :     }
     167             :   }
     168         465 :   copyMatrixBookeepingFromFirstComponent( outmat );
     169         564 :   for(unsigned i=1; i<getNumberOfComponents(); ++i) {
     170          99 :     getPntrToComponent(i)->copyBookeepingArrayFromArgument( myval );
     171             :   }
     172         465 : }
     173             : 
     174          61 : void ActionWithMatrix::getColumnBookeepingArrays( RequiredMatrixElements& outmat ) {
     175          61 :   outmat.ncols = maxcolsize;
     176          61 :   Value* mycomp = getPntrToComponent(0);
     177          61 :   outmat.resize( mycomp->getShape()[1]*(1+outmat.ncols) );
     178        2267 :   for(unsigned i=0; i<mycomp->getShape()[1]; ++i) {
     179        2206 :     outmat[i*(1+outmat.ncols)] = 0;
     180             :   }
     181         633 :   for(unsigned i=0; i<mycomp->getShape()[0]; ++i) {
     182       15552 :     for(unsigned j=0; j<mycomp->getRowLength(i); ++j) {
     183       14980 :       unsigned cstart = mycomp->getRowIndex(i,j)*( 1 + outmat.ncols );
     184       14980 :       outmat[cstart + 1 + outmat[cstart]] = i;
     185       14980 :       outmat[cstart]++;
     186             :     }
     187             :   }
     188          61 : }
     189             : 
     190       36905 : std::vector<unsigned>& ActionWithMatrix::getListOfActiveTasks() {
     191       36905 :   if( gatherForceOnColumns ) {
     192          61 :     return column_list;
     193             :   } else {
     194       36844 :     return ActionWithVector::getListOfActiveTasks();
     195             :   }
     196             : }
     197             : 
     198       27898 : unsigned ActionWithMatrix::getSizeOfValueInStash( const Value* myval ) {
     199       27898 :   unsigned nstored = myval->getNumberOfStoredValues();
     200       27898 :   if( no_thread_gather ) {
     201         184 :     unsigned ncolstore = maxcolsize*myval->getShape()[1];
     202         184 :     if( ncolstore>nstored ) {
     203             :       return ncolstore;
     204             :     }
     205             :   }
     206             :   return nstored;
     207             : }
     208             : 
     209             : // The stash holds the elements of a row contiguously, with the components of
     210             : // an element interleaved ([row][column][component]); each component Value
     211             : // stores its rows contiguously ([row][column]). The copies write the storage
     212             : // of the Values directly (ActionWithMatrix is a friend of Value), the rows
     213             : // are independent so the loop over them is parallel.
     214             : template <typename T>
     215       13058 : void ActionWithMatrix::matrixStashToValuesImpl( const std::vector<unsigned>& partialTaskList, const std::vector<T>& stash ) {
     216             :   const std::size_t ncomp = getNumberOfComponents();
     217       13058 :   const std::size_t ncols = getPntrToComponent(0)->getNumberOfColumns();
     218       13058 :   const unsigned nrows = partialTaskList.size();
     219       13058 :   std::vector<double*> data( ncomp );
     220       13058 :   std::vector<Value*> periodic;
     221       27007 :   for(std::size_t k=0; k<ncomp; ++k) {
     222       13949 :     Value* myval = getPntrToComponent(k);
     223             :     plumed_dbg_assert( !myval->hasDeriv );
     224       13949 :     myval->value_set=true;
     225       13949 :     data[k] = myval->data.data();
     226       13949 :     if( myval->isPeriodic() ) {
     227          55 :       periodic.push_back( myval );
     228             :     }
     229             :   }
     230       13058 :   const Value* first = getPntrToComponent(0);
     231             : #ifdef _OPENMP
     232             :   // without OpenMP nt would be unused (an error with -Werror)
     233       13058 :   const unsigned nt = OpenMP::getGoodNumThreads( stash );
     234             : #endif
     235       13058 :   #pragma omp parallel for num_threads(nt) schedule(static)
     236             :   for(unsigned i=0; i<nrows; ++i) {
     237             :     const std::size_t row = partialTaskList[i];
     238             :     const unsigned ncr = first->getRowLength(row);
     239             : #ifndef NDEBUG
     240             :     for(std::size_t k=1; k<ncomp; ++k) {
     241             :       plumed_assert( ncr == getPntrToComponent(k)->getRowLength(row) );
     242             :     }
     243             : #endif
     244             :     const T* src = stash.data() + ncomp*ncols*row;
     245             :     const std::size_t dst = row*ncols;
     246             :     for(std::size_t j=0; j<ncr; ++j) {
     247             :       for(std::size_t k=0; k<ncomp; ++k) {
     248             :         data[k][dst+j] = src[j*ncomp+k];
     249             :       }
     250             :     }
     251             :     // e.g. the matrix of torsions is periodic
     252             :     for(Value* myval : periodic) {
     253             :       for(std::size_t j=0; j<ncr; ++j) {
     254             :         myval->applyPeriodicity( dst+j );
     255             :       }
     256             :     }
     257             :   }
     258       13058 : }
     259             : 
     260             : template <typename T>
     261       12019 : void ActionWithMatrix::matrixForcesToStashImpl( const std::vector<unsigned>& partialTaskList, std::vector<T>& stash ) const {
     262             :   const std::size_t ncomp = getNumberOfComponents();
     263       12019 :   const std::size_t ncols = getConstPntrToComponent(0)->getNumberOfColumns();
     264       12019 :   const Value* first = getConstPntrToComponent(0);
     265       12019 :   std::vector<const double*> force( ncomp );
     266       24539 :   for(std::size_t k=0; k<ncomp; ++k) {
     267       12520 :     force[k] = getConstPntrToComponent(k)->inputForce.data();
     268             :   }
     269       12019 :   if( gatherForceOnColumns ) {
     270             :     // forces on the columns: element (i,j) of the matrix goes to the slot of
     271             :     // column j; the slots of a column are filled in the order of the rows
     272             :     const std::vector<unsigned>& rowTasks( getConstListOfActiveTasks() );
     273          61 :     const unsigned nrows = rowTasks.size();
     274          61 :     std::vector<unsigned> column_totals( first->getShape()[1], 0 );
     275         633 :     for(unsigned i=0; i<nrows; ++i) {
     276         572 :       const std::size_t row = rowTasks[i];
     277             :       const unsigned ncr = first->getRowLength(row);
     278       15552 :       for(std::size_t j=0; j<ncr; ++j) {
     279       14980 :         const std::size_t colno = first->getRowIndex(row,j);
     280       14980 :         T* dst = stash.data() + ncomp*maxcolsize*colno + column_totals[colno]*ncomp;
     281       42218 :         for(std::size_t k=0; k<ncomp; ++k) {
     282             :           plumed_dbg_assert( colno==getConstPntrToComponent(k)->getRowIndex(row,j) );
     283       27238 :           dst[k] = force[k][row*ncols+j];
     284             :         }
     285       14980 :         column_totals[colno]++;
     286             :       }
     287             :     }
     288             :     return;
     289             :   }
     290       11958 :   const unsigned nrows = partialTaskList.size();
     291             : #ifdef _OPENMP
     292             :   // without OpenMP nt would be unused (an error with -Werror)
     293       11958 :   const unsigned nt = OpenMP::getGoodNumThreads( stash );
     294             : #endif
     295       11958 :   #pragma omp parallel for num_threads(nt) schedule(static)
     296             :   for(unsigned i=0; i<nrows; ++i) {
     297             :     const std::size_t row = partialTaskList[i];
     298             :     const unsigned ncr = first->getRowLength(row);
     299             : #ifndef NDEBUG
     300             :     for(std::size_t k=1; k<ncomp; ++k) {
     301             :       plumed_assert( ncr == getConstPntrToComponent(k)->getRowLength(row) );
     302             :     }
     303             : #endif
     304             :     T* dst = stash.data() + ncomp*ncols*row;
     305             :     const std::size_t src = row*ncols;
     306             :     for(std::size_t j=0; j<ncr; ++j) {
     307             :       for(std::size_t k=0; k<ncomp; ++k) {
     308             :         dst[j*ncomp+k] = force[k][src+j];
     309             :       }
     310             :     }
     311             :   }
     312             : }
     313             : 
     314       13058 : void ActionWithMatrix::transferStashToValues( const std::vector<unsigned>& partialTaskList, const std::vector<double>& stash ) {
     315       13058 :   matrixStashToValuesImpl( partialTaskList, stash );
     316       13058 : }
     317             : 
     318           0 : void ActionWithMatrix::transferStashToValues( const std::vector<unsigned>& partialTaskList, const std::vector<float>& stash ) {
     319           0 :   matrixStashToValuesImpl( partialTaskList, stash );
     320           0 : }
     321             : 
     322       12019 : void ActionWithMatrix::transferForcesToStash( const std::vector<unsigned>& partialTaskList, std::vector<double>& stash ) const {
     323       12019 :   matrixForcesToStashImpl( partialTaskList, stash );
     324       12019 : }
     325             : 
     326           0 : void ActionWithMatrix::transferForcesToStash( const std::vector<unsigned>& partialTaskList, std::vector<float>& stash ) const {
     327           0 :   matrixForcesToStashImpl( partialTaskList, stash );
     328           0 : }
     329             : 
     330             : } //namespace PLMD