Reuse the neighbour search between Get and the neighbour getters

Get (and GetHost) now accept a LookupTableNeighbours structure in which
they store the search they performed. Giving that same structure to
GetNeighbours or GetNeighboursIndx makes them reuse this search instead
of walking the table a second time, and they search it again as usual as
soon as different coordinates are requested.

The structure is held by the caller rather than by the lookup table, so
that it is thread-private when it is declared inside an idefix_for: a
lookup table is shared by all of the threads of a loop, and can hence not
cache anything itself. The existing Get is left untouched.

Author: dutta-alankar

doc/source/programmingguide.rst

@@ -676,6 +676,46 @@ For instance, the 2D table ``example2D.csv`` above, interpolated in (x=2.1, y=3.
 The same methods without the ``Host`` suffix can be called from within an ``idefix_for`` loop, the arrays being then
 local to the loop.
 
+Both ``Get`` and the methods above search the table for the cell surrounding the requested coordinates. When the
+interpolated value *and* its neighbours are needed for the same coordinates, this search can be performed only once,
+by handing a ``LookupTableNeighbours<nDim>`` structure to ``Get``: the search it stores is then reused by
+``GetNeighbours`` and ``GetNeighboursIndx`` instead of being done a second time.
+
+.. code-block:: c++
+
+  idefix_for("loop",0, 10, KOKKOS_LAMBDA (int i) {
+    real x[2];
+    x[0] = 2.1;
+    x[1] = 3.5;
+
+    // The search performed by Get is stored in neighbours...
+    LookupTableNeighbours<2> neighbours;
+    real value = csv.Get(x, neighbours);
+
+    // ... and is reused by the two methods below, which do not search the table again
+    real xN[4];
+    real dataN[4];
+    int idx[2];
+    int dataIdx[4];
+    csv.GetNeighbours(x, neighbours, xN, dataN);
+    csv.GetNeighboursIndx(x, neighbours, idx, dataIdx);
+  });
+
+The very same methods are available on the host, as ``GetHost``, ``GetNeighboursHost`` and ``GetNeighboursIndxHost``.
+The structure gives access to the search itself, ``neighbours.idx[n]`` being the index of the left neighbour along
+the dimension ``n``, and ``neighbours.delta[n]`` the elementary ratio used to weight the two neighbours of that
+dimension.
+
+.. note::
+  The stored search is only reused when the coordinates it was computed for are the ones being requested. Calling
+  ``GetNeighbours`` or ``GetNeighboursIndx`` with different coordinates simply searches the table again, and updates
+  the structure accordingly, so that a ``LookupTableNeighbours`` can be reused from one point to the next.
+
+.. warning::
+  The structure is declared by the caller, and not by the lookup table itself: when it is used inside an
+  ``idefix_for`` loop, it should be declared *inside* the loop, so that each thread has its own copy. A lookup table
+  is shared by all of the threads of a loop, and can therefore not store anything of the sort itself.
+
 .. note::
   When the table is interpolated in function space (see above), ``xN`` and ``dataN`` are returned in the original
   space of the table, i.e. ``invFunc`` is applied to the values which are stored in ``xin`` and ``data``. The

src/utils/lookupTable.hpp

@@ -31,6 +31,34 @@ struct LookupTableExp {
   }
 };
 
+// Neighbours of a point of a lookup table, as they are found by the search performed by Get:
+// idx[n] is the index of the left neighbour along the dimension n, and delta[n] is the elementary
+// ratio used to weight the two neighbours of that dimension.
+// Get (and GetHost) fill this structure when it is given as an argument, and GetNeighbours and
+// GetNeighboursIndx then reuse the search it contains instead of performing it a second time
+// (which is only done when they are called with the same coordinates: they search the table again
+// as usual when a different x is requested).
+// This structure is kept by the caller, so that it is thread-private when it is declared inside
+// an idefix_for loop. The lookup table itself is shared by all of the threads of a loop, and can
+// therefore not be used to store anything of the sort.
+template <const int kDim>
+struct LookupTableNeighbours {
+  real x[kDim];        // coordinates for which the neighbours below were computed
+  int idx[kDim];       // index of the left neighbour along each dimension
+  real delta[kDim];    // elementary ratio between the two neighbours of each dimension
+  bool valid{false};   // whether the neighbours above have been successfully computed
+
+  // Check whether this structure already holds the neighbours of the coordinates xIn
+  KOKKOS_INLINE_FUNCTION
+  bool Matches(const real xIn[kDim]) const {
+    if(!valid) return(false);
+    for(int n = 0 ; n < kDim ; n++) {
+      if(x[n] != xIn[n]) return(false);
+    }
+    return(true);
+  }
+};
+
 template <const int kDim, class TFunc = LookupTableLog, class TInvFunc = LookupTableExp>
 class LookupTable {
  public:
@@ -276,6 +304,59 @@ class LookupTable {
     return(value);
   }
 
+  // Fill "neighbours" with the neighbours of x, unless it already holds them (in which case the
+  // table is not searched again). This is what allows Get, GetNeighbours and GetNeighboursIndx to
+  // share a single search when they are called successively with the same coordinates.
+  template<typename Tint, typename Treal>
+  KOKKOS_INLINE_FUNCTION
+  void SearchNeighbours(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin,
+                        LookupTableNeighbours<kDim> &neighbours) const {
+    // Nothing to do if the neighbours of these very coordinates are already known
+    if(neighbours.Matches(x)) return;
+
+    neighbours.valid = GetIndices(x, dimensions, offset, xin, neighbours.idx, neighbours.delta);
+    for(int n = 0 ; n < kDim ; n++) {
+      neighbours.x[n] = x[n];
+    }
+  }
+
+  // Generic getter which stores in "neighbours" the elements of the table it used, so that a
+  // subsequent call to GetNeighbours or GetNeighboursIndx with the same coordinates does not
+  // search the table again
+  template<typename Tint, typename Treal>
+  KOKKOS_INLINE_FUNCTION
+  real Get(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin, Treal &data,
+           LookupTableNeighbours<kDim> &neighbours) const {
+    SearchNeighbours(x, dimensions, offset, xin, neighbours);
+
+    if(!neighbours.valid) return(NAN);
+
+    // Do a linear interpolation from the neightbouring points to get our value.
+    real value = 0;
+
+    // loop on all of the vertices of the neighbours
+    for(unsigned int n = 0 ; n < (1 << kDim) ; n++) {
+      real weight = 1.0;
+      for(unsigned int m = 0 ; m < kDim ; m++) {
+        unsigned int myBit = 1 << m;
+        // If bit is set, we're doing the right vertex, otherwise we're doing the left vertex
+        if((n & myBit) > 0) {
+          // We're on the right
+          weight = weight*neighbours.delta[m];
+        } else {
+          // We're on the left
+          weight = weight*(1-neighbours.delta[m]);
+        }
+      }
+      value = value + weight*data(GetDataIndex(dimensions, neighbours.idx, n));
+    }
+
+    // The interpolation was performed on func(data), so we transform the result back
+    if(interpolateInFuncSpace) value = invFunc(value);
+
+    return(value);
+  }
+
   // Generic getter for the neighbours used by the interpolation, for all kinds of input arrays.
   // On output, xN[2*n] and xN[2*n+1] are the coordinates bracketing x[n] along the dimension n,
   // and dataN[v] is the data at the vertex v of these neighbours (see GetDataIndex for the
@@ -286,19 +367,30 @@ class LookupTable {
   KOKKOS_INLINE_FUNCTION
   void GetNeighbours(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin, Treal &data,
                      real xN[2*kDim], real dataN[1 << kDim]) const {
-    int idx[kDim];
-    real delta[kDim];
+    LookupTableNeighbours<kDim> neighbours;
+    GetNeighbours(x, dimensions, offset, xin, data, neighbours, xN, dataN);
+  }
+
+  // Same as above, but the search is stored in (and reused from) "neighbours": the table is only
+  // searched again when "neighbours" does not already hold the neighbours of x, e.g. because it
+  // was filled by a previous call to Get with these very same coordinates.
+  template<typename Tint, typename Treal>
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighbours(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin, Treal &data,
+                     LookupTableNeighbours<kDim> &neighbours,
+                     real xN[2*kDim], real dataN[1 << kDim]) const {
+    SearchNeighbours(x, dimensions, offset, xin, neighbours);
 
-    if(!GetIndices(x, dimensions, offset, xin, idx, delta)) {
+    if(!neighbours.valid) {
       for(int n = 0 ; n < 2*kDim ; n++) xN[n] = NAN;
       for(unsigned int n = 0 ; n < (1 << kDim) ; n++) dataN[n] = NAN;
       return;
     }
 
     // Coordinates of the neighbours along each dimension
     for(int n = 0 ; n < kDim ; n++) {
-      xN[2*n] = xin(offset(n) + idx[n]);
-      xN[2*n+1] = xin(offset(n) + idx[n]+1);
+      xN[2*n] = xin(offset(n) + neighbours.idx[n]);
+      xN[2*n+1] = xin(offset(n) + neighbours.idx[n]+1);
       if(interpolateInFuncSpace) {
         xN[2*n] = invFunc(xN[2*n]);
         xN[2*n+1] = invFunc(xN[2*n+1]);
@@ -307,7 +399,7 @@ class LookupTable {
 
     // Data on each vertex of the neighbours
     for(unsigned int n = 0 ; n < (1 << kDim) ; n++) {
-      dataN[n] = data(GetDataIndex(dimensions, idx, n));
+      dataN[n] = data(GetDataIndex(dimensions, neighbours.idx, n));
       if(interpolateInFuncSpace) dataN[n] = invFunc(dataN[n]);
     }
   }
@@ -321,16 +413,30 @@ class LookupTable {
   KOKKOS_INLINE_FUNCTION
   void GetNeighboursIndx(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin,
                          int idx[kDim], int dataIdx[1 << kDim]) const {
-    real delta[kDim];
+    LookupTableNeighbours<kDim> neighbours;
+    GetNeighboursIndx(x, dimensions, offset, xin, neighbours, idx, dataIdx);
+  }
+
+  // Same as above, but the search is stored in (and reused from) "neighbours", exactly like the
+  // GetNeighbours variant above
+  template<typename Tint, typename Treal>
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighboursIndx(const real x[kDim], Tint &dimensions, Tint &offset, Treal &xin,
+                         LookupTableNeighbours<kDim> &neighbours,
+                         int idx[kDim], int dataIdx[1 << kDim]) const {
+    SearchNeighbours(x, dimensions, offset, xin, neighbours);
 
-    if(!GetIndices(x, dimensions, offset, xin, idx, delta)) {
+    if(!neighbours.valid) {
       for(int n = 0 ; n < kDim ; n++) idx[n] = -1;
       for(unsigned int n = 0 ; n < (1 << kDim) ; n++) dataIdx[n] = -1;
       return;
     }
 
+    for(int n = 0 ; n < kDim ; n++) {
+      idx[n] = neighbours.idx[n];
+    }
     for(unsigned int n = 0 ; n < (1 << kDim) ; n++) {
-      dataIdx[n] = GetDataIndex(dimensions, idx, n);
+      dataIdx[n] = GetDataIndex(dimensions, neighbours.idx, n);
     }
   }
 
@@ -369,6 +475,51 @@ class LookupTable {
   void GetNeighboursIndxHost(const real x[kDim], int idx[kDim], int dataIdx[1 << kDim]) const {
     GetNeighboursIndx(x, dimensionsHost, offsetHost, xinHost, idx, dataIdx);
   }
+
+  // Getter on device, which stores the neighbours it used in "neighbours". Giving that same
+  // structure to GetNeighbours or GetNeighboursIndx below then avoids searching the table twice.
+  KOKKOS_INLINE_FUNCTION
+  real Get(const real x[kDim], LookupTableNeighbours<kDim> &neighbours) const {
+    return(Get(x, dimensionsDev, offsetDev, xinDev, dataDev, neighbours));
+  }
+
+  // Getter on Host, which stores the neighbours it used in "neighbours"
+  KOKKOS_INLINE_FUNCTION
+  real GetHost(const real x[kDim], LookupTableNeighbours<kDim> &neighbours) const {
+    return(Get(x, dimensionsHost, offsetHost, xinHost, dataHost, neighbours));
+  }
+
+  // Getter for the neighbours used by the interpolation, on device, reusing the search stored in
+  // "neighbours" when it was performed for these very same coordinates
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighbours(const real x[kDim], LookupTableNeighbours<kDim> &neighbours,
+                     real xN[2*kDim], real dataN[1 << kDim]) const {
+    GetNeighbours(x, dimensionsDev, offsetDev, xinDev, dataDev, neighbours, xN, dataN);
+  }
+
+  // Getter for the neighbours used by the interpolation, on Host, reusing the search stored in
+  // "neighbours" when it was performed for these very same coordinates
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighboursHost(const real x[kDim], LookupTableNeighbours<kDim> &neighbours,
+                         real xN[2*kDim], real dataN[1 << kDim]) const {
+    GetNeighbours(x, dimensionsHost, offsetHost, xinHost, dataHost, neighbours, xN, dataN);
+  }
+
+  // Getter for the indices of the neighbours, on device, reusing the search stored in
+  // "neighbours" when it was performed for these very same coordinates
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighboursIndx(const real x[kDim], LookupTableNeighbours<kDim> &neighbours,
+                         int idx[kDim], int dataIdx[1 << kDim]) const {
+    GetNeighboursIndx(x, dimensionsDev, offsetDev, xinDev, neighbours, idx, dataIdx);
+  }
+
+  // Getter for the indices of the neighbours, on Host, reusing the search stored in
+  // "neighbours" when it was performed for these very same coordinates
+  KOKKOS_INLINE_FUNCTION
+  void GetNeighboursIndxHost(const real x[kDim], LookupTableNeighbours<kDim> &neighbours,
+                             int idx[kDim], int dataIdx[1 << kDim]) const {
+    GetNeighboursIndx(x, dimensionsHost, offsetHost, xinHost, neighbours, idx, dataIdx);
+  }
 };
 
 template <int kDim, class TFunc, class TInvFunc>