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             : #ifndef __PLUMED_matrixtools_MatrixTimesVectorBase_h
      23             : #define __PLUMED_matrixtools_MatrixTimesVectorBase_h
      24             : 
      25             : #include "core/ActionWithVector.h"
      26             : #include "core/ParallelTaskManager.h"
      27             : 
      28             : namespace PLMD {
      29             : namespace matrixtools {
      30             : 
      31             : namespace helpers {
      32             : ///true if T has a public declared "isRowSum" type
      33             : template <typename T, typename=void>
      34             : constexpr bool isRowSum=false;
      35             : template <typename T>
      36             : constexpr bool isRowSum<T,std::void_t<typename T::isRowSum>> =true;
      37             : }
      38             : 
      39        1185 : class MatrixTimesVectorData {
      40             : public:
      41             :   std::size_t fshift;
      42             :   Matrix<std::size_t> pairs;
      43             : #ifdef __PLUMED_HAS_OPENACC
      44             :   void toACCDevice() const {
      45             : #pragma acc enter data copyin(this[0:1],fshift)
      46             :     pairs.toACCDevice();
      47             :   }
      48             :   void removeFromACCDevice() const {
      49             :     pairs.removeFromACCDevice();
      50             : #pragma acc exit data delete(fshift,this[0:1])
      51             :   }
      52             : #endif //__PLUMED_HAS_OPENACC
      53             : };
      54             : 
      55             : class MatrixForceIndexInput {
      56             : public:
      57             :   std::size_t rowlen;
      58             :   View<const std::size_t> indices;
      59             : //temporary template for compilation
      60             :   template <typename ParallelActionsInput>
      61      124096 :   MatrixForceIndexInput( std::size_t task_index,
      62             :                          std::size_t ipair,
      63             :                          const MatrixTimesVectorData& actiondata,
      64             :                          const ParallelActionsInput& input ):
      65      124096 :     rowlen(input.bookeeping[input.bookstarts[actiondata.pairs[ipair][0]]
      66      124096 :                             + (1+input.ncols[actiondata.pairs[ipair][0]])*task_index]),
      67      124096 :     indices(input.bookeeping + input.bookstarts[actiondata.pairs[ipair][0]]
      68      124096 :             + (1+input.ncols[actiondata.pairs[ipair][0]])*task_index + 1,
      69      124096 :             rowlen) {}
      70             : };
      71             : 
      72             : template <typename precision>
      73             : struct MatrixTimesVectorInput {
      74             :   bool noderiv;
      75             :   std::size_t rowlen;
      76             :   View<const std::size_t> indices;
      77             :   View<const precision> matrow;
      78             :   View<const precision> vector;
      79             : //temporary template for compilation
      80             :   template <typename ParallelActionsInput>
      81    10157202 :   MatrixTimesVectorInput( std::size_t task_index,
      82             :                           std::size_t ipair,
      83             :                           const MatrixTimesVectorData& actiondata,
      84             :                           const ParallelActionsInput& input,
      85             :                           precision* argdata ):
      86    10157202 :     noderiv(input.noderiv),
      87    10157202 :     rowlen(input.bookeeping[input.bookstarts[actiondata.pairs[ipair][0]]
      88    10157202 :                             + (1+input.ncols[actiondata.pairs[ipair][0]])*task_index]),
      89    10157202 :     indices(input.bookeeping + input.bookstarts[actiondata.pairs[ipair][0]]
      90    10157202 :             + (1+input.ncols[actiondata.pairs[ipair][0]])*task_index + 1,rowlen),
      91    10157202 :     matrow(argdata + input.argstarts[actiondata.pairs[ipair][0]]
      92    10157202 :            + task_index*input.ncols[actiondata.pairs[ipair][0]],rowlen),
      93    10157202 :     vector(argdata + input.argstarts[actiondata.pairs[ipair][1]], input.shapedata[1]) {
      94    10157202 :   }
      95             : };
      96             : 
      97             : template <typename precision>
      98             : class MatrixTimesVectorOutput {
      99             : public:
     100             :   std::size_t rowlen;
     101             :   View<precision,1> values;
     102             :   View<precision> matrow_deriv;
     103             :   View<precision> vector_deriv;
     104             : //temporary template for compilation
     105             :   template <typename ParallelActionsInput, typename ParallelActionsOutput>
     106    10157202 :   MatrixTimesVectorOutput( std::size_t task_index,
     107             :                            std::size_t ipair,
     108             :                            std::size_t nder,
     109             :                            const MatrixTimesVectorData& actiondata,
     110             :                            const ParallelActionsInput& input,
     111             :                            ParallelActionsOutput& output ):
     112    10157202 :     rowlen(input.bookeeping[input.bookstarts[actiondata.pairs[ipair][0]]
     113    10157202 :                             + (1+input.ncols[actiondata.pairs[ipair][0]])*task_index]),
     114    10157202 :     values(output.values.data()+ipair),
     115    10157202 :     matrow_deriv(output.derivatives.data()+ipair*nder,rowlen),
     116    10157202 :     vector_deriv(output.derivatives.data()+ipair*nder+rowlen,rowlen) {
     117    10157202 :   }
     118             : };
     119             : 
     120             : template <class CV, typename myPTM=defaultPTM>
     121             : class MatrixTimesVectorBase : public ActionWithVector {
     122             : public:
     123             :   using input_type = MatrixTimesVectorData;
     124             :   typedef cvprecision_t<CV> precision;
     125             :   typedef MatrixTimesVectorBase<CV,myPTM> mytype;
     126             :   using PTM = typename myPTM::template PTM<mytype>;
     127             :   typedef typename PTM::ParallelActionsInput ParallelActionsInput;
     128             :   typedef typename PTM::ParallelActionsOutput ParallelActionsOutput;
     129             : private:
     130             : /// The parallel task manager
     131             :   PTM taskmanager;
     132             : public:
     133             :   static void registerKeywords( Keywords& keys );
     134             :   static void registerLocalKeywords( Keywords& keys );
     135             :   explicit MatrixTimesVectorBase(const ActionOptions&);
     136             :   std::string getOutputComponentDescription( const std::string& cname, const Keywords& keys ) const override ;
     137             :   unsigned getNumberOfDerivatives() override ;
     138             :   void prepare() override ;
     139             :   void calculate() override ;
     140             :   void applyNonZeroRankForces( std::vector<double>& outforces ) override ;
     141             :   int checkTaskIsActive( const unsigned& itask ) const override ;
     142             :   /// Override this so we write the graph properly
     143          10 :   std::string writeInGraph() const override {
     144          10 :     return "MATRIX_VECTOR_PRODUCT";
     145             :   }
     146             :   static void performTask( std::size_t task_index,
     147             :                            const MatrixTimesVectorData& actiondata,
     148             :                            ParallelActionsInput& input,
     149             :                            ParallelActionsOutput& output );
     150             :   static int getNumberOfValuesPerTask( std::size_t task_index,
     151             :                                        const MatrixTimesVectorData& actiondata );
     152             :   static void getForceIndices( std::size_t task_index,
     153             :                                std::size_t colno,
     154             :                                std::size_t ntotal_force,
     155             :                                const MatrixTimesVectorData& actiondata,
     156             :                                const ParallelActionsInput& input,
     157             :                                ForceIndexHolder force_indices );
     158             : };
     159             : 
     160             : template <class CV, typename myPTM>
     161         802 : void MatrixTimesVectorBase<CV, myPTM>::registerKeywords( Keywords& keys ) {
     162         802 :   ActionWithVector::registerKeywords(keys);
     163         802 :   keys.setDisplayName("MATRIX_VECTOR_PRODUCT");
     164         802 :   registerLocalKeywords( keys );
     165         802 :   ActionWithValue::useCustomisableComponents(keys);
     166         802 :   keys.add("hidden","NO_ACTION_LOG","suppresses printing from action on the log");
     167         802 : }
     168             : 
     169             : template <class CV, typename myPTM>
     170        1489 : void MatrixTimesVectorBase<CV, myPTM>::registerLocalKeywords( Keywords& keys ) {
     171        1489 :   PTM::registerKeywords( keys );
     172        2978 :   keys.addInputKeyword("compulsory","ARG","matrix/vector/scalar","the label for the matrix and the vector/scalar that are being multiplied.  Alternatively, you can provide labels for multiple matrices and a single vector or labels for a single matrix and multiple vectors. In these cases multiple matrix vector products will be computed.");
     173        1489 :   keys.add("hidden","MASKED_INPUT_ALLOWED","turns on that you are allowed to use masked inputs ");
     174        2978 :   keys.setValueDescription("vector","the vector that is obtained by taking the product between the matrix and the vector that were input");
     175        1489 :   ActionWithValue::useCustomisableComponents(keys);
     176        1489 : }
     177             : 
     178             : template <class CV, typename myPTM>
     179           6 : std::string MatrixTimesVectorBase<CV, myPTM>::getOutputComponentDescription( const std::string& cname,
     180             :     const Keywords& keys ) const {
     181           6 :   if( getPntrToArgument(1)->getRank()==1 ) {
     182           0 :     for(unsigned i=1; i<getNumberOfArguments(); ++i) {
     183           0 :       if( getPntrToArgument(i)->getName().find(cname)!=std::string::npos ) {
     184           0 :         return "the product of the matrix " + getPntrToArgument(0)->getName() + " and the vector " + getPntrToArgument(i)->getName();
     185             :       }
     186             :     }
     187             :   }
     188          21 :   for(unsigned i=0; i<getNumberOfArguments()-1; ++i) {
     189          21 :     if( getPntrToArgument(i)->getName().find(cname)!=std::string::npos ) {
     190          18 :       return "the product of the matrix " + getPntrToArgument(i)->getName() + " and the vector " + getPntrToArgument(getNumberOfArguments()-1)->getName();
     191             :     }
     192             :   }
     193           0 :   plumed_merror( "could not understand request for component " + cname );
     194             :   return "";
     195             : }
     196             : 
     197             : template <class CV, typename myPTM>
     198         395 : MatrixTimesVectorBase<CV, myPTM>::MatrixTimesVectorBase(const ActionOptions&ao):
     199             :   Action(ao),
     200             :   ActionWithVector(ao),
     201         395 :   taskmanager(this) {
     202         395 :   if( getNumberOfArguments()<2 ) {
     203           0 :     error("Not enough arguments specified");
     204             :   }
     205             :   bool vectormask=false, derivbool = true;
     206        2057 :   for(unsigned i=0; i<getNumberOfArguments(); ++i) {
     207        1662 :     if( getPntrToArgument(i)->hasDerivatives() ) {
     208           0 :       error("arguments should be vectors or matrices");
     209             :     }
     210        1662 :     if( getPntrToArgument(i)->getRank()<=1 ) {
     211         580 :       ActionWithVector* av=dynamic_cast<ActionWithVector*>( getPntrToArgument(i)->getPntrToAction() );
     212         580 :       if( av && av->getNumberOfMasks()>=0 ) {
     213             :         vectormask=true;
     214             :       }
     215             :     }
     216        1662 :     if( !getPntrToArgument(i)->isDerivativeZeroWhenValueIsZero() ) {
     217             :       derivbool = false;
     218             :     }
     219             :   }
     220         395 :   if( !vectormask ) {
     221         395 :     ignoreMaskArguments();
     222             :   }
     223             : 
     224         395 :   std::vector<std::size_t> shape(1);
     225         395 :   shape[0]=getPntrToArgument(0)->getShape()[0];
     226         395 :   if( getNumberOfArguments()==2 ) {
     227         341 :     addValue( shape );
     228         341 :     setNotPeriodic();
     229         341 :     if( derivbool ) {
     230         260 :       getPntrToComponent(0)->setDerivativeIsZeroWhenValueIsZero();
     231             :     }
     232             :   } else {
     233             :     unsigned namestart=1, nameend=getNumberOfArguments();
     234          54 :     if( getPntrToArgument(1)->getRank()==2 ) {
     235             :       namestart = 0;
     236          45 :       nameend = getNumberOfArguments()-1;
     237             :     }
     238             : 
     239         980 :     for(unsigned i=namestart; i<nameend; ++i) {
     240         926 :       std::string name = getPntrToArgument(i)->getName();
     241         926 :       if( name.find_first_of(".")!=std::string::npos ) {
     242             :         std::size_t dot=name.find_first_of(".");
     243        1464 :         name = name.substr(dot+1);
     244             :       }
     245         926 :       addComponent( name, shape );
     246         926 :       componentIsNotPeriodic( name );
     247         926 :       if( derivbool ) {
     248        1844 :         copyOutput( getLabel() + "." + name )->setDerivativeIsZeroWhenValueIsZero();
     249             :       }
     250             :     }
     251             :   }
     252             :   // This sets up an array in the parallel task manager to hold all the indices
     253             :   // Sets up the index list in the task manager
     254         395 :   std::size_t nder = getPntrToArgument(getNumberOfArguments()-1)->getNumberOfStoredValues();
     255             :   MatrixTimesVectorData input;
     256         395 :   input.pairs.resize( getNumberOfArguments()-1, 2 );
     257         395 :   if( getPntrToArgument(1)->getRank()==2 ) {
     258         777 :     for(unsigned i=1; i<getNumberOfArguments(); ++i) {
     259         732 :       input.pairs[i-1][0] = i-1;
     260         732 :       input.pairs[i-1][1] = getNumberOfArguments()-1;
     261             :     }
     262          45 :     input.fshift=0;
     263             :   } else {
     264         885 :     for(unsigned i=1; i<getNumberOfArguments(); ++i) {
     265         535 :       input.pairs[i-1][0] = 0;
     266         535 :       input.pairs[i-1][1] = i;
     267             :     }
     268         350 :     input.fshift=nder;
     269             :   }
     270         395 :   taskmanager.setActionInput( input );
     271         395 : }
     272             : 
     273             : template <class CV, typename myPTM>
     274       28908 : unsigned MatrixTimesVectorBase<CV, myPTM>::getNumberOfDerivatives() {
     275             :   unsigned nderivatives=0;
     276      701252 :   for(unsigned i=0; i<getNumberOfArguments(); ++i) {
     277      672344 :     nderivatives += getPntrToArgument(i)->getNumberOfStoredValues();
     278             :   }
     279       28908 :   return nderivatives;
     280             : }
     281             : 
     282             : template <class CV, typename myPTM>
     283       13798 : void MatrixTimesVectorBase<CV, myPTM>::prepare() {
     284       13798 :   ActionWithVector::prepare();
     285       13798 :   Value* myval = getPntrToComponent(0);
     286       13798 :   if( myval->getShape()[0]==getPntrToArgument(0)->getShape()[0] ) {
     287       13788 :     return;
     288             :   }
     289          10 :   std::vector<std::size_t> shape(1);
     290          10 :   shape[0] = getPntrToArgument(0)->getShape()[0];
     291          10 :   myval->setShape(shape);
     292          10 : }
     293             : 
     294             : template <class CV, typename myPTM>
     295       13791 : void MatrixTimesVectorBase<CV, myPTM>::calculate() {
     296       13791 :   std::size_t nder = getPntrToArgument(getNumberOfArguments()-1)->getNumberOfStoredValues();
     297             :   std::size_t nvectors;
     298       13791 :   if( getPntrToArgument(1)->getRank()==2 ) {
     299             :     nvectors = 1;
     300             :   } else {
     301       13663 :     nvectors = getNumberOfArguments()-1;
     302             :   }
     303             :   // Each product has a derivative for every element of the matrix row and for
     304             :   // the vector element it multiplies: the number of columns of the sparse
     305             :   // matrix bounds the length of a row. Sizing this with the length of the
     306             :   // vector instead made the per-task derivative buffer of the GPU task manager
     307             :   // quadratic in the number of atoms
     308             :   std::size_t maxcols = 0;
     309       44424 :   for(unsigned i=0; i<getNumberOfArguments(); ++i) {
     310       30633 :     if( getPntrToArgument(i)->getRank()==2 && getPntrToArgument(i)->getNumberOfColumns()>maxcols ) {
     311             :       maxcols = getPntrToArgument(i)->getNumberOfColumns();
     312             :     }
     313             :   }
     314             :   if constexpr ( helpers::isRowSum<CV>) { //getName()=="MATRIX_VECTOR_PRODUCT_ROWSUMS" ) {
     315             :     taskmanager.setupParallelTaskManager( maxcols, 0 );
     316             :   } else {
     317       13791 :     taskmanager.setupParallelTaskManager( 2*maxcols, nvectors*nder );
     318             :   }
     319       13791 :   taskmanager.runAllTasks();
     320       13791 : }
     321             : 
     322             : template <class CV, typename myPTM>
     323     1827125 : int MatrixTimesVectorBase<CV, myPTM>::checkTaskIsActive( const unsigned& itask ) const {
     324     2518770 :   for(unsigned i=0; i<getNumberOfArguments(); ++i) {
     325             :     Value* myarg = getPntrToArgument(i);
     326     2518770 :     if( myarg->getRank()==1 && !myarg->hasDerivatives() ) {
     327             :       return 0;
     328     1827201 :     } else if( myarg->getRank()==2 && !myarg->hasDerivatives() ) {
     329     1827201 :       if (myarg->checkValueIsActiveForMMul(itask)) {
     330             :         return 1;
     331             :       }
     332             :     } else {
     333           0 :       plumed_merror("should not be in action " + getName() );
     334             :     }
     335             :   }
     336             :   return 0;
     337             : }
     338             : 
     339             : template <class CV, typename myPTM>
     340     1516567 : void MatrixTimesVectorBase<CV, myPTM>::performTask( std::size_t task_index,
     341             :     const MatrixTimesVectorData& actiondata,
     342             :     ParallelActionsInput& input,
     343             :     ParallelActionsOutput& output ) {
     344    11673769 :   for(unsigned i=0; i<actiondata.pairs.nrows(); ++i) {
     345    10157202 :     MatrixTimesVectorOutput<precision> doutput( task_index,
     346             :         i,
     347             :         input.nderivatives_per_scalar,
     348             :         actiondata,
     349             :         input,
     350             :         output );
     351    10157202 :     CV::performTask( MatrixTimesVectorInput<precision>( task_index,
     352             :                      i,
     353             :                      actiondata,
     354             :                      input,
     355             :                      input.inputdata ),
     356             :                      doutput );
     357             :   }
     358     1516567 : }
     359             : 
     360             : template <class CV, typename myPTM>
     361       11950 : void MatrixTimesVectorBase<CV, myPTM>::applyNonZeroRankForces( std::vector<double>& outforces ) {
     362       11950 :   taskmanager.applyForces( outforces );
     363       11950 : }
     364             : 
     365             : template <class CV, typename myPTM>
     366      444675 : int MatrixTimesVectorBase<CV, myPTM>::getNumberOfValuesPerTask( std::size_t task_index,
     367             :     const MatrixTimesVectorData& actiondata ) {
     368      444675 :   return 1;
     369             : }
     370             : 
     371             : template <class CV, typename myPTM>
     372      444675 : void MatrixTimesVectorBase<CV, myPTM>::getForceIndices( std::size_t task_index,
     373             :     std::size_t colno,
     374             :     std::size_t ntotal_force,
     375             :     const MatrixTimesVectorData& actiondata,
     376             :     const ParallelActionsInput& input,
     377             :     ForceIndexHolder force_indices ) {
     378     2267455 :   for(unsigned i=0; i<actiondata.pairs.nrows(); ++i) {
     379     1822780 :     std::size_t base = input.argstarts[actiondata.pairs[i][0]]
     380     1822780 :                        + task_index*input.ncols[actiondata.pairs[i][0]];
     381     1822780 :     std::size_t n = input.bookeeping[input.bookstarts[actiondata.pairs[i][0]]
     382     1822780 :                                      + (1+input.ncols[actiondata.pairs[i][0]])*task_index];
     383    17109241 :     for(unsigned j=0; j<n; ++j) {
     384    15286461 :       force_indices.indices[i][j] = base + j;
     385             :     }
     386     1822780 :     force_indices.threadsafe_derivatives_end[i] = n;
     387     1822780 :     force_indices.tot_indices[i] = CV::getAdditionalIndices( n,
     388      124096 :                                    input.argstarts[actiondata.pairs[i][1]],
     389      248192 :                                    MatrixForceIndexInput( task_index, i, actiondata, input ),
     390             :                                    force_indices.indices[i] );
     391             :   }
     392      444675 : }
     393             : 
     394             : template <typename T>
     395             : struct MatrixTimesVectorRowSums {
     396             :   typedef void isRowSums;
     397             :   typedef T precision;
     398     9894129 :   static void performTask( const MatrixTimesVectorInput<T>& input,
     399             :                            MatrixTimesVectorOutput<T>& output ) {
     400    84626337 :     for(unsigned i=0; i<input.rowlen; ++i) {
     401    74732208 :       output.values[0] += input.matrow[i];
     402    74732208 :       if( input.noderiv ) {
     403    60268915 :         continue;
     404             :       }
     405    14463293 :       output.matrow_deriv[i] = 1;
     406             :     }
     407     9894129 :   }
     408             :   static std::size_t getAdditionalIndices( std::size_t n,
     409             :       std::size_t vecstart,
     410             :       const MatrixForceIndexInput& fin,
     411             :       View<std::size_t> indices ) {
     412             :     return n;
     413             :   }
     414             : };
     415             : 
     416             : 
     417             : template <typename T>
     418             : struct MatrixTimesVectorProper {
     419             :   typedef T precision;
     420      263073 :   static void performTask( const MatrixTimesVectorInput<T>& input,
     421             :                            MatrixTimesVectorOutput<T>& output ) {
     422     2839372 :     for(unsigned i=0; i<input.rowlen; ++i) {
     423     2576299 :       std::size_t ind = input.indices[i];
     424     2576299 :       output.values[0] += input.matrow[i]*input.vector[ind];
     425     2576299 :       if( input.noderiv ) {
     426     1753131 :         continue;
     427             :       }
     428      823168 :       output.matrow_deriv[i] = input.vector[ind];
     429      823168 :       output.vector_deriv[i] = input.matrow[i];
     430             :     }
     431      263073 :   }
     432      124096 :   static std::size_t getAdditionalIndices( std::size_t n,
     433             :       std::size_t vecstart,
     434             :       const MatrixForceIndexInput& fin,
     435             :       View<std::size_t> indices ) {
     436      947264 :     for(unsigned i=0; i<fin.rowlen; ++i) {
     437      823168 :       indices[n+i] = vecstart + fin.indices[i];
     438             :     }
     439      124096 :     return n + fin.rowlen;
     440             :   }
     441             : };
     442             : }
     443             : }
     444             : #endif