From b2a66a5db2d5e145eeec78eab0ef5c4049650c42 Mon Sep 17 00:00:00 2001
From: Giovanni Bussi <giovanni.bussi@gmail.com>
Date: Wed, 13 Mar 2024 12:35:51 +0100
Subject: [PATCH] Group indexes to make memory access faster

---
 src/core/ActionAtomistic.cpp       | 80 ++++++++++++++++++++++++------
 src/core/ActionAtomistic.h         |  1 +
 src/core/ActionWithVirtualAtom.cpp | 23 +++++++--
 3 files changed, 85 insertions(+), 19 deletions(-)

diff --git a/src/core/ActionAtomistic.cpp b/src/core/ActionAtomistic.cpp
index a896d69fcf..2947fd4c29 100644
--- a/src/core/ActionAtomistic.cpp
+++ b/src/core/ActionAtomistic.cpp
@@ -98,6 +98,24 @@ void ActionAtomistic::requestAtoms(const std::vector<AtomNumber> & a, const bool
     if( atom_value_ind[i].first==0 ) unique.push_back(indexes[i]);
     else if( atom_value_ind[i].second>0 ) error("action atomistic is not set up to deal with multiple vectors in position input");
   }
+
+  atom_value_ind_grouped.clear();
+
+  if(atom_value_ind.size()>0) {
+    auto nn = atom_value_ind[0].first;
+    auto kk = atom_value_ind[0].second;
+    atom_value_ind_grouped.push_back(std::pair<std::size_t,std::vector<std::size_t>>(nn, {}));
+    atom_value_ind_grouped.back().second.push_back(kk);
+    auto prev_nn=nn;
+    for(unsigned i=1; i<atom_value_ind.size(); i++) {
+      auto nn = atom_value_ind[i].first;
+      auto kk = atom_value_ind[i].second;
+      if(nn!=prev_nn) atom_value_ind_grouped.push_back(std::pair<std::size_t,std::vector<std::size_t>>(nn, {}));
+      atom_value_ind_grouped.back().second.push_back(kk);
+      prev_nn=nn;
+    }
+  }
+
   // Add the dependencies to the actions that we require
   Tools::removeDuplicates(unique); value_depends.resize(0);
   for(unsigned i=0; i<requirements.size(); ++i ) {
@@ -271,15 +289,34 @@ void ActionAtomistic::retrieveAtoms( const bool& force ) {
   ActionToPutData* cv = chargev[0]->getPntrToAction()->castToActionToPutData();
   if(cv) chargesWereSet=cv->hasBeenSet();
   unsigned j = 0;
-  for(const auto & a : atom_value_ind) {
-    std::size_t nn = a.first, kk = a.second;
-    positions[j][0] = xpos[nn]->data[kk];
-    positions[j][1] = ypos[nn]->data[kk];
-    positions[j][2] = zpos[nn]->data[kk];
-    charges[j] = chargev[nn]->data[kk];
-    masses[j] = masv[nn]->data[kk];
-    j++;
+
+// for(const auto & a : atom_value_ind) {
+//   std::size_t nn = a.first, kk = a.second;
+//   positions[j][0] = xpos[nn]->data[kk];
+//   positions[j][1] = ypos[nn]->data[kk];
+//   positions[j][2] = zpos[nn]->data[kk];
+//   charges[j] = chargev[nn]->data[kk];
+//   masses[j] = masv[nn]->data[kk];
+//   j++;
+// }
+
+  for(const auto & a : atom_value_ind_grouped) {
+    const auto nn=a.first;
+    auto & xp=xpos[nn]->data;
+    auto & yp=ypos[nn]->data;
+    auto & zp=zpos[nn]->data;
+    auto & ch=chargev[nn]->data;
+    auto & ma=masv[nn]->data;
+    for(const auto & kk : a.second) {
+      positions[j][0] = xp[kk];
+      positions[j][1] = yp[kk];
+      positions[j][2] = zp[kk];
+      charges[j] = ch[kk];
+      masses[j] = ma[kk];
+      j++;
+    }
   }
+
 }
 
 void ActionAtomistic::setForcesOnAtoms(const std::vector<double>& forcesToApply, unsigned& ind) {
@@ -289,13 +326,28 @@ void ActionAtomistic::setForcesOnAtoms(const std::vector<double>& forcesToApply,
     ypos[value_depends[i]]->hasForce = true;
     zpos[value_depends[i]]->hasForce = true;
   }
-  for(const auto & a : atom_value_ind) {
-    plumed_dbg_massert( ind<forcesToApply.size(), "problem setting forces in " + getLabel() );
-    std::size_t nn = a.first, kk = a.second;
-    xpos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
-    ypos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
-    zpos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
+
+// for(const auto & a : atom_value_ind) {
+//   plumed_dbg_massert( ind<forcesToApply.size(), "problem setting forces in " + getLabel() );
+//   std::size_t nn = a.first, kk = a.second;
+//   xpos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
+//   ypos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
+//   zpos[nn]->inputForce[kk] += forcesToApply[ind]; ind++;
+// }
+
+  for(const auto & a : atom_value_ind_grouped) {
+    const auto nn=a.first;
+    plumed_dbg_assert(ind<forcesToApply.size()) << "problem setting forces in " << getLabel();
+    auto & xp=xpos[nn]->inputForce;
+    auto & yp=ypos[nn]->inputForce;
+    auto & zp=zpos[nn]->inputForce;
+    for(const auto & kk : a.second) {
+      xp[kk] += forcesToApply[ind]; ind++;
+      yp[kk] += forcesToApply[ind]; ind++;
+      zp[kk] += forcesToApply[ind]; ind++;
+    }
   }
+
   setForcesOnCell( forcesToApply, ind );
 }
 
diff --git a/src/core/ActionAtomistic.h b/src/core/ActionAtomistic.h
index 09ec132730..606af7b75b 100644
--- a/src/core/ActionAtomistic.h
+++ b/src/core/ActionAtomistic.h
@@ -52,6 +52,7 @@ class ActionAtomistic :
   std::vector<AtomNumber> indexes;         // the set of needed atoms
   std::vector<std::size_t>   value_depends;   // The list of values that are being used
   std::vector<std::pair<std::size_t, std::size_t > > atom_value_ind;  // The list of values and indices for the atoms that are being used
+  std::vector<std::pair<std::size_t,std::vector<std::size_t>>> atom_value_ind_grouped;
 /// unique should be an ordered set since we later create a vector containing the corresponding indexes
   std::vector<AtomNumber>  unique;
 /// unique_local should be an ordered set since we later create a vector containing the corresponding indexes
diff --git a/src/core/ActionWithVirtualAtom.cpp b/src/core/ActionWithVirtualAtom.cpp
index b103d71d05..afd1a6e156 100644
--- a/src/core/ActionWithVirtualAtom.cpp
+++ b/src/core/ActionWithVirtualAtom.cpp
@@ -21,6 +21,7 @@
 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
 #include "ActionWithVirtualAtom.h"
 #include <array>
+#include <iostream>
 
 namespace PLMD {
 
@@ -71,12 +72,24 @@ void ActionWithVirtualAtom::apply() {
   double xf = xval->inputForce[0];
   double yf = yval->inputForce[0];
   double zf = zval->inputForce[0];
-  for(const auto & a : atom_value_ind) {
-    std::size_t nn = a.first, kk = a.second;
-    xpos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
-    ypos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
-    zpos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+// for(const auto & a : atom_value_ind) {
+//   std::size_t nn = a.first, kk = a.second;
+//   xpos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+//   ypos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+//   zpos[nn]->inputForce[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+// }
+  for(const auto & a : atom_value_ind_grouped) {
+    const auto nn=a.first;
+    auto & xp=xpos[nn]->inputForce;
+    auto & yp=ypos[nn]->inputForce;
+    auto & zp=zpos[nn]->inputForce;
+    for(const auto & kk : a.second) {
+      xp[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+      yp[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+      zp[kk] += xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++;
+    }
   }
+
   std::array<double,9> virial;
   for(unsigned i=0; i<9; ++i) { virial[i] = xf*xval->data[1+k] + yf*yval->data[1+k] + zf*zval->data[1+k]; k++; }
   unsigned ind = 0; setForcesOnCell( virial.data(), virial.size(), ind );