From 9a5a2487be1ffaf4195b7f955621ffb2d7a65b83 Mon Sep 17 00:00:00 2001
From: Ania Brown <anbrown@nvidia.com>
Date: Mon, 23 Jan 2023 13:46:20 +0100
Subject: [PATCH] Add warmup in the SliceUnion

---
 include/atrip/SliceUnion.hpp | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/include/atrip/SliceUnion.hpp b/include/atrip/SliceUnion.hpp
index abdfcb3..bd04709 100644
--- a/include/atrip/SliceUnion.hpp
+++ b/include/atrip/SliceUnion.hpp
@@ -445,6 +445,34 @@ template <typename F=double>
                      std::inserter(freePointers, freePointers.begin()),
                      [](DataPtr<F> ptr) { return ptr; });
 
+#if defined(HAVE_CUDA)
+      LOG(1,"Atrip") << "warming communication up " << slices.size() << "\n";
+      WITH_CHRONO("cuda:warmup",
+                   int nRanks=Atrip::np, requestCount=0;
+                   int nSends=sliceBuffers.size()*nRanks;
+                   MPI_Request *requests = (MPI_Request*) malloc(nSends*2 * sizeof(MPI_Request));
+                   MPI_Status *statuses = (MPI_Status*) malloc(nSends*2 * sizeof(MPI_Status));
+                   for (int sliceId=0; sliceId<sliceBuffers.size(); sliceId++){
+                     for (int rankId=0; rankId<nRanks; rankId++){
+                       MPI_Isend((void*)SOURCES_DATA(sources[0]),
+                                 sliceSize,
+                                 traits::mpi::datatypeOf<F>(),
+                                 rankId,
+                                 100,
+                                 universe,
+                                 &requests[requestCount++]);
+                       MPI_Irecv((void*)sliceBuffers[sliceId],
+                                 sliceSize,
+                                 traits::mpi::datatypeOf<F>(),
+                                 rankId,
+                                 100,
+                                 universe,
+                                 &requests[requestCount++]);
+                     }
+                   }
+                   MPI_Waitall(nSends*2, requests, statuses);
+                  )
+#endif
 
 
       LOG(1,"Atrip") << "#slices " << slices.size() << "\n";