Refactor tracking of GPU short-range work/skipping

[alexxy/gromacs.git] / src / gromacs / nbnxm / nbnxm_gpu.h

diff --git a/src/gromacs/nbnxm/nbnxm_gpu.h b/src/gromacs/nbnxm/nbnxm_gpu.h
index 7e88129f4ee89728bebc38427d0d5617c2f0e863..a1e480045680829d8fd67106d4ce8f8dbdef1324 100644 (file)
--- a/src/gromacs/nbnxm/nbnxm_gpu.h
+++ b/src/gromacs/nbnxm/nbnxm_gpu.h
@@ -55,6 +55,11 @@
 struct nbnxn_atomdata_t;
 enum class GpuTaskCompletion;
 
+namespace gmx
+{
+class GpuBonded;
+}
+
 namespace Nbnxm
 {
 
@@ -69,13 +74,11 @@ class Grid;
  * \param [in]    nb        GPU nonbonded data.
  * \param [in]    nbdata    Host-side atom data structure.
  * \param [in]    aloc      Atom locality flag.
- * \param [in]    haveOtherWork  True if there are other tasks that require the nbnxn coordinate input.
  */
 GPU_FUNC_QUALIFIER
 void gpu_copy_xq_to_gpu(gmx_nbnxn_gpu_t gmx_unused               *nb,
                         const struct nbnxn_atomdata_t gmx_unused *nbdata,
-                        AtomLocality gmx_unused                   aloc,
-                        bool gmx_unused                           haveOtherWork) GPU_FUNC_TERM
+                        AtomLocality gmx_unused                   aloc) GPU_FUNC_TERM
 
 /*! \brief
  * Launch asynchronously the nonbonded force calculations.
@@ -133,17 +136,14 @@ void gpu_launch_kernel_pruneonly(gmx_nbnxn_gpu_t gmx_unused     *nb,
                                  int gmx_unused                  numParts) GPU_FUNC_TERM
 
 /*! \brief
- * Launch asynchronously the download of nonbonded forces from the GPU
+ * Launch asynchronously the download of short-range forces from the GPU
  * (and energies/shift forces if required).
- * When haveOtherWork=true, the copy-back is done even when there was
- * no non-bonded work.
  */
 GPU_FUNC_QUALIFIER
 void gpu_launch_cpyback(gmx_nbnxn_gpu_t  gmx_unused *nb,
                         nbnxn_atomdata_t gmx_unused *nbatom,
                         int              gmx_unused  flags,
-                        AtomLocality     gmx_unused  aloc,
-                        bool             gmx_unused  haveOtherWork) GPU_FUNC_TERM
+                        AtomLocality     gmx_unused  aloc) GPU_FUNC_TERM
 
 /*! \brief Attempts to complete nonbonded GPU task.
  *
@@ -171,7 +171,6 @@ void gpu_launch_cpyback(gmx_nbnxn_gpu_t  gmx_unused *nb,
  * \param[in]  nb     The nonbonded data GPU structure
  * \param[in]  flags  Force flags
  * \param[in]  aloc   Atom locality identifier
- * \param[in]  haveOtherWork  Tells whether there is other work than non-bonded work in the nbnxn stream(s)
  * \param[out] e_lj   Pointer to the LJ energy output to accumulate into
  * \param[out] e_el   Pointer to the electrostatics energy output to accumulate into
  * \param[out] fshift Pointer to the shift force buffer to accumulate into
@@ -182,7 +181,6 @@ GPU_FUNC_QUALIFIER
 bool gpu_try_finish_task(gmx_nbnxn_gpu_t gmx_unused  *nb,
                          int             gmx_unused   flags,
                          AtomLocality    gmx_unused   aloc,
-                         bool            gmx_unused   haveOtherWork,
                          real            gmx_unused  *e_lj,
                          real            gmx_unused  *e_el,
                          rvec            gmx_unused  *fshift,
@@ -198,7 +196,6 @@ bool gpu_try_finish_task(gmx_nbnxn_gpu_t gmx_unused  *nb,
  * \param[in] nb The nonbonded data GPU structure
  * \param[in] flags Force flags
  * \param[in] aloc Atom locality identifier
- * \param[in]  haveOtherWork  Tells whether there is other work than non-bonded work in the nbnxn stream(s)
  * \param[out] e_lj Pointer to the LJ energy output to accumulate into
  * \param[out] e_el Pointer to the electrostatics energy output to accumulate into
  * \param[out] fshift Pointer to the shift force buffer to accumulate into
@@ -207,7 +204,6 @@ GPU_FUNC_QUALIFIER
 void gpu_wait_finish_task(gmx_nbnxn_gpu_t gmx_unused *nb,
                           int             gmx_unused  flags,
                           AtomLocality    gmx_unused  aloc,
-                          bool            gmx_unused  haveOtherWork,
                           real            gmx_unused *e_lj,
                           real            gmx_unused *e_el,
                           rvec            gmx_unused *fshift) GPU_FUNC_TERM
@@ -242,6 +238,37 @@ CUDA_FUNC_QUALIFIER
 void nbnxnInsertNonlocalGpuDependency(const gmx_nbnxn_gpu_t gmx_unused    *nb,
                                       const InteractionLocality gmx_unused interactionLocality) CUDA_FUNC_TERM
 
+/*! \brief Set up internal flags that indicate what type of short-range work there is.
+ *
+ * As nonbondeds and bondeds share input/output buffers and GPU queues,
+ * both are considered when checking for work in the current domain.
+ *
+ * This function is expected to be called every time the work-distribution
+ * can change (i.e. at search/domain decomposition steps).
+ *
+ * \param[inout]  nb         Pointer to the nonbonded GPU data structure
+ * \param[in]     gpuBonded  Pointer to the GPU bonded data structure
+ * \param[in]     iLocality  Interaction locality identifier
+ */
+GPU_FUNC_QUALIFIER
+void setupGpuShortRangeWork(gmx_nbnxn_gpu_t                  gmx_unused *nb,
+                            const gmx::GpuBonded             gmx_unused *gpuBonded,
+                            const Nbnxm::InteractionLocality gmx_unused  iLocality) GPU_FUNC_TERM
+
+/*! \brief Returns true if there is GPU short-range work for the given atom locality.
+ *
+ * Note that as, unlike nonbonded tasks, bonded tasks are not split into local/nonlocal,
+ * and therefore if there are GPU offloaded bonded interactions, this function will return
+ * true for both local and nonlocal atom range.
+ *
+ * \param[inout]  nb        Pointer to the nonbonded GPU data structure
+ * \param[in]     aLocality Atom locality identifier
+ */
+GPU_FUNC_QUALIFIER
+bool haveGpuShortRangeWork(const gmx_nbnxn_gpu_t     gmx_unused *nb,
+                           const Nbnxm::AtomLocality gmx_unused  aLocality) GPU_FUNC_TERM_WITH_RETURN(false)
+
+
 } // namespace Nbnxm
 
 #endif