Remove single-dimension limitation from GPU halo exchange

[alexxy/gromacs.git] / src / gromacs / domdec / gpuhaloexchange_impl.cuh

diff --git a/src/gromacs/domdec/gpuhaloexchange_impl.cuh b/src/gromacs/domdec/gpuhaloexchange_impl.cuh
index 77aaf3f9bfe1e5f5713cc90bd565d060f745aa0f..761938a0133c3c5f5645d17e2121d818356ca9a7 100644 (file)
--- a/src/gromacs/domdec/gpuhaloexchange_impl.cuh
+++ b/src/gromacs/domdec/gpuhaloexchange_impl.cuh
@@ -72,6 +72,7 @@ public:
     /*! \brief Creates GPU Halo Exchange object.
      *
      * \param [inout] dd                       domdec structure
+     * \param [in]    dimIndex                 the dimension index for this instance
      * \param [in]    mpi_comm_mysim           communicator used for simulation
      * \param [in]    deviceContext            GPU device context
      * \param [in]    localStream              local NB CUDA stream
@@ -80,6 +81,7 @@ public:
      * \param [in]    wcycle                   The wallclock counter
      */
     Impl(gmx_domdec_t*        dd,
+         int                  dimIndex,
          MPI_Comm             mpi_comm_mysim,
          const DeviceContext& deviceContext,
          const DeviceStream&  localStream,
@@ -198,12 +200,16 @@ private:
     float3* d_f_ = nullptr;
     //! An event recorded once the exchanged forces are ready on the GPU
     GpuEventSynchronizer fReadyOnDevice_;
+    //! The dimension index corresponding to this halo exchange instance
+    int dimIndex_ = 0;
     //! The pulse corresponding to this halo exchange instance
     int pulse_ = 0;
     //! Number of zones. Always 1 for 1-D case.
     const int nzone_ = 1;
     //! The wallclock counter
     gmx_wallcycle* wcycle_ = nullptr;
+    //! The atom offset for receive (x) or send (f) for dimension index and pulse corresponding to this halo exchange instance
+    int atomOffset_ = 0;
 };
 
 } // namespace gmx