[OpenMP] Introduce stream pool to make sure the correctness of device synchronization

[tianshilei1992]

Summary

In previous patch, in order to optimize performance, we only synchronize once
for each target region. The syncrhonization is via stream synchronization.
However, in the extreme situation, the performce might be bad. Consider the
following case: There is a task that requires transferring huge amount of data
(call many times of data transferring function). It is scheduled to the first
stream. And then we have 255 very light tasks scheduled to the remaining 255
streams (by default we have 256 streams). They can be finished before we do
synchronization at the end of the first task. Next, we get another very huge
task. It will be scheduled again to the first stream. Now the first task
finishes its kernel launch and call stream synchronization. Right now, the
stream already contains two kernels, and the synchronization will wait until the
two kernels finish instead of just the first one for the first task.

In this patch, we introduce stream pool. After each synchronization, the stream
will be returned back to the pool to make sure that for each synchronization,
only expected operations are waited.

[jdoerfert]

I like this a lot. I have one high-level question and a request for documentation:
Shouldn't we hide the stream stuff behind the async_info object instead?
I mean, we could avoid accidental misuse by restricting the getStream/returnStream to the constructor/destructor of async_info for example.
We should also unify the naming.

Maybe, put the entire "stream" (=queue) management in a class. All the initialization happens there, all the tear down as well, and when you make it friends with async_info we can restrict access too.

WDYT?

openmp/libomptarget/plugins/cuda/src/rtl.cpp

367

This is clever. So clever in fact that we need a comment above StreamPool (or similar) describing how the different values interact and how the StreamPool evolves. Basically how it works.

[tianshilei1992]

In D77412#1960259, @jdoerfert wrote:

I like this a lot. I have one high-level question and a request for documentation:
Shouldn't we hide the stream stuff behind the async_info object instead?
I mean, we could avoid accidental misuse by restricting the getStream/returnStream to the constructor/destructor of async_info for example.
We should also unify the naming.

Maybe, put the entire "stream" (=queue) management in a class. All the initialization happens there, all the tear down as well, and when you make it friends with async_info we can restrict access too.

WDYT?

Agree. Will create a new class StreamManager.

[tianshilei1992]

Wrap all stream-related operations into StreamManager.

[tianshilei1992]

Rebase my patch

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

250

No need to put your name here, git blame would know anyway.

We should move all the CUDA stuff into the stream manager and the stream manager out of the RTLDeviceInfoTy but both can be done later. For now just replace the friends with a getter.

[tianshilei1992]

Removed friend functions. Later I'll prepare a new patch to improve encapsulation.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

102

Later we can think abut making this a vector of structs instead as we tend to access them per device anyway. We can then align the structs such that different processes interacting with different devices don't really interfere with each other.

117

Outline this into a helper. Also the error checking should be a helper.

So something like:

/// ....
bool checkResult(CUResult err, const char *ErrMsg) {
  if (err == CUDA_SUCCESS)
    return true;
   DP(ErrMsg);
   CUDA_ERR_STRING(err);
   return false;
}

/// ...
bool switchToDevice(DeviceId) {
  CUresult err = cuCtxSetCurrent((*ContextsPtr)[DeviceId]);
  return checkResult(err, "Error when creating CUDA stream to resize stream pool\n");
}

and then use it eveywhere. Should cut down the size and duplication.

136

If there is no reason ever to provide a nullptr as CtxPtr make it a reference everywhere instead.

149

Don't you need to create the sreams, e.g., call resizeStreamPool instead?

545

I was expecting this to be part of the StreamManager. If it can only happen at this point, maybe we want a StreamManager::initializeDevice method or similar. We can then also avoid exposing the stream pool to the outside.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

102

That would be ideal.

136

My bad. At the very beginning I thought by some chances it will be nullptr but it turns out that will never happen. :-)

149

Here I follow the initial design which is later initialization. In the constructor, it only allocate memory but does not initialize.

545

Yes, we could do that. But what I'm thinking is, do we assume cuCtxSetCurrent is called before we call StreamManager::initializeDevice, or we will set the context again in StreamManager::initializeDevice?

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

117

We can do the first one. As for the second function, it must be the member function. Let's do that in next patch which moves every CUDA related things into class.

[jdoerfert]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

117

ok

545

Once could move both to the StreamManager::initializeDevice, set ctx and create streams. One could then either reset the context or just document that the context is set by initializeDevice. This is a one time cost so setting the context twice is not my main concern.

[tianshilei1992]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

545

Yes, in the next patch, all these issues are no long existing. Here let's assume that when call StreamManager::initializeDevice, the right context has been set.

[jdoerfert]

LGTM.

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----------------------------------------------------------------------===//
 
 #include <atomic>
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <vector>
 
 #include "omptargetplugin.h"
 
 #ifndef TARGET_NAME
 #define TARGET_NAME CUDA
 #endif

 };
 
 /// List that contains all the kernels.
 /// FIXME: we may need this to be per device and per library.
 std::list<KernelTy> KernelsList;
 
 /// Class containing all the device information.
 class RTLDeviceInfoTy {
+  class StreamManagerTy {
+    int NumberOfDevices;
+    // Per-device stream mutex
+    std::vector<std::unique_ptr<std::mutex>> StreamMtx;
+    // Per-device stream Id indicates the next available stream in the pool
+    std::vector<int> NextStreamId;
+    // Per-device stream pool
+    std::vector<std::vector<CUstream>> StreamPool;
+    // Pointer to per-device context
+    std::vector<CUcontext> *ContextsPtr;
+
+    // If there is no CUstream left in the pool, we will resize the pool to
+    // allocate more CUstream. This function should be called with device mutex,
+    // and we do not resize to smaller one.
+    void resizeStreamPool(const int DeviceId, const size_t NewSize) {
+      std::vector<CUstream> &Pool = StreamPool[DeviceId];
+      const size_t CurrentSize = Pool.size();
+      assert(NewSize > CurrentSize &&
+             "new size is not larger than current size");
+
+      Pool.resize(NewSize, nullptr);
+
+      CUresult err = cuCtxSetCurrent((*ContextsPtr)[DeviceId]);
+      if (err != CUDA_SUCCESS) {
+        DP("Error when setting current CUDA context\n");
+        CUDA_ERR_STRING(err);
+        // We will return if cannot switch to the right context in case of
+        // creating bunch of streams that are not corresponding to the right
+        // device. The offloading will fail later because selected CUstream is
+        // nullptr.
+        return;
+      }
+
+      for (size_t I = CurrentSize; I < NewSize; ++I) {
+        err = cuStreamCreate(&Pool[I], CU_STREAM_NON_BLOCKING);
+        if (err != CUDA_SUCCESS) {
+          DP("Error when creating CUDA stream to resize stream pool\n");
+          CUDA_ERR_STRING(err);
+        }
+      }
+    }
+
+  public:
+    // Shilei: I don't like to put this friend function here but it seems like
+    // the best way not to refine the whole file.
+    friend int32_t __tgt_rtl_init_device(int32_t device_id);
+
+    StreamManagerTy(const int NumberOfDevices, std::vector<CUcontext> *CtxPtr)
+        : NumberOfDevices(NumberOfDevices), ContextsPtr(CtxPtr) {
+      StreamPool.resize(NumberOfDevices);
+      NextStreamId.resize(NumberOfDevices);
+      StreamMtx.resize(NumberOfDevices);
+
+      // Initially let's create 32 streams for each device
+      int EnvNumInitialStreams = 32;
+      char *envStr = getenv("LIBOMPTARGET_NUM_INITIAL_STREAMS");
+      if (envStr) {
+        EnvNumInitialStreams = std::stoi(envStr);
+      }
+
+      // Initialize the stream pool for each device
+      for (std::vector<CUstream> &S : StreamPool) {
+        S.resize(EnvNumInitialStreams);
+      }
+
+      // Initialize the next stream id
+      std::fill(NextStreamId.begin(), NextStreamId.end(), 0);
+
+      // Initialize stream mutex
+      for (std::unique_ptr<std::mutex> &Ptr : StreamMtx) {
+        Ptr = std::make_unique<std::mutex>();
+      }
+    }
+
+    ~StreamManagerTy() {
+      // Destroy streams
+      for (int I = 0; I < NumberOfDevices; ++I) {
+        CUresult err = cuCtxSetCurrent((*ContextsPtr)[I]);
+        if (err != CUDA_SUCCESS) {
+          DP("Error when setting current CUDA context\n");
+          CUDA_ERR_STRING(err);
+        }
+
+        for (CUstream &S : StreamPool[I]) {
+          if (!S)
+            continue;
+          err = cuStreamDestroy(S);
+          if (err != CUDA_SUCCESS) {
+            DP("Error when destroying CUDA stream\n");
+            CUDA_ERR_STRING(err);
+          }
+        }
+      }
+    }
+
+    // Get a CUstream from pool. Per-device next stream id always points to the
+    // next available CUstream. That means, CUstreams [0, id-1] have been
+    // assigned, and [id,] are still available. If there is no CUstream left, we
+    // will ask more CUstreams from CUDA RT. Each time a CUstream is assigned,
+    // the id will increase one.
+    // xxxxxs+++++++++
+    //      ^
+    //      id
+    // After assignment, the pool becomes the following and s is assigned.
+    // xxxxxs+++++++++
+    //       ^
+    //       id
+    CUstream getStream(const int DeviceId) {
+      assert(DeviceId >= 0 &&
+             static_cast<size_t>(DeviceId) < NextStreamId.size() &&
+             "Unexpected device id");
+
+      const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
+      int &Id = NextStreamId[DeviceId];
+      // No CUstream left in the pool, we need to request from CUDA RT
+      if (Id == StreamPool[DeviceId].size()) {
+        // By default we double the stream pool every time
+        resizeStreamPool(DeviceId, Id * 2);
+      }
+      return StreamPool[DeviceId][Id++];
+    }
+
+    // Return a CUstream back to pool. As mentioned above, per-device next
+    // stream is always points to the next available CUstream, so when we return
+    // a CUstream, we need to first decrease the id, and then copy the CUstream
+    // back.
+    // It is worth noting that, the order of streams return might be different
+    // from that they're assigned, that saying, at some point, there might be
+    // two identical CUstreams.
+    // xxax+a+++++
+    //     ^
+    //     id
+    // However, it doesn't matter, because they're always on the two sides of
+    // id. The left one will in the end be overwritten by another CUstream.
+    // Therefore, after several execution, the order of pool might be different
+    // from its initial state.
+    void returnStream(const int DeviceId, CUstream Stream) {
+      assert(DeviceId >= 0 &&
+             static_cast<size_t>(DeviceId) < NextStreamId.size() &&
+             "Unexpected device id");
+      {
+        const std::lock_guard<std::mutex> Lock(*StreamMtx[DeviceId]);
+        int &Id = NextStreamId[DeviceId];
+        assert(Id > 0 && "Wrong stream ID");
+        StreamPool[DeviceId][--Id] = Stream;
+      }
+    }
+  };
+
   std::vector<std::list<FuncOrGblEntryTy>> FuncGblEntries;
-  std::vector<std::unique_ptr<std::atomic_uint>> NextStreamId;
+  std::shared_ptr<StreamManagerTy> StreamManager;
 
 public:
+  // Shilei: I don't like to put this friend function here but it seems like
+  // the best way not to refine the whole file.
+  friend int32_t __tgt_rtl_init_device(int32_t device_id);
+
   int NumberOfDevices;
   std::vector<CUmodule> Modules;
   std::vector<CUcontext> Contexts;
-  std::vector<std::vector<CUstream>> Streams;
 
   // Device properties
   std::vector<int> ThreadsPerBlock;
   std::vector<int> BlocksPerGrid;
   std::vector<int> WarpSize;
 
   // OpenMP properties
   std::vector<int> NumTeams;
   std::vector<int> NumThreads;
 
   // OpenMP Environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
-  int EnvNumStreams;
 
   // OpenMP Requires Flags
   int64_t RequiresFlags;
 
-  //static int EnvNumThreads;
-  static const int HardTeamLimit = 1<<16; // 64k
+  // static int EnvNumThreads;
+  static const int HardTeamLimit = 1 << 16; // 64k
   static const int HardThreadLimit = 1024;
   static const int DefaultNumTeams = 128;
   static const int DefaultNumThreads = 128;
 
+  CUstream getStream(const int DeviceId) {
+    return StreamManager->getStream(DeviceId);
+  }
+
+  void returnStream(const int DeviceId, __tgt_async_info *AsyncInfoPtr) {
+    assert(AsyncInfoPtr && "AsyncInfoPtr is nullptr");
+    assert(AsyncInfoPtr->Queue && "AsyncInfoPtr->Queue is nullptr");
+
+    StreamManager->returnStream(
+        DeviceId, reinterpret_cast<CUstream>(AsyncInfoPtr->Queue));
+    AsyncInfoPtr->Queue = nullptr;
+  }
+
   // Record entry point associated with device
   void addOffloadEntry(int32_t device_id, __tgt_offload_entry entry) {
     assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
 
     E.Entries.push_back(entry);
   }
void clearOffloadEntriesTable(int32_t device_id) {
     assert(device_id < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncGblEntries[device_id].emplace_back();
     FuncOrGblEntryTy &E = FuncGblEntries[device_id].back();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = 0;
   }
 
-  // Get the next stream on a given device in a round robin manner
-  CUstream &getNextStream(const int DeviceId) {
-    assert(DeviceId >= 0 &&
-           static_cast<size_t>(DeviceId) < NextStreamId.size() &&
-           "Unexpected device id!");
-    const unsigned int Id = NextStreamId[DeviceId]->fetch_add(1);
-    return Streams[DeviceId][Id % EnvNumStreams];
-  }
-
   RTLDeviceInfoTy() {
 #ifdef OMPTARGET_DEBUG
     if (char *envStr = getenv("LIBOMPTARGET_DEBUG")) {
       DebugLevel = std::stoi(envStr);
     }
 #endif // OMPTARGET_DEBUG
 
     DP("Start initializing CUDA\n");
 
     CUresult err = cuInit(0);
     if (err != CUDA_SUCCESS) {
       DP("Error when initializing CUDA\n");
       CUDA_ERR_STRING(err);
       return;
     }
 
     NumberOfDevices = 0;
 
     err = cuDeviceGetCount(&NumberOfDevices);
     if (err != CUDA_SUCCESS) {
       DP("Error when getting CUDA device count\n");
       CUDA_ERR_STRING(err);
       return;
     }
 
     if (NumberOfDevices == 0) {
       DP("There are no devices supporting CUDA.\n");
       return;
     }
 
     FuncGblEntries.resize(NumberOfDevices);
     Contexts.resize(NumberOfDevices);
-    Streams.resize(NumberOfDevices);
-    NextStreamId.resize(NumberOfDevices);
     ThreadsPerBlock.resize(NumberOfDevices);
     BlocksPerGrid.resize(NumberOfDevices);
     WarpSize.resize(NumberOfDevices);
     NumTeams.resize(NumberOfDevices);
     NumThreads.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     char *envStr = getenv("OMP_TEAM_LIMIT");
     if (envStr) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(envStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     } else {
       EnvTeamLimit = -1;
     }
     envStr = getenv("OMP_NUM_TEAMS");
     if (envStr) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(envStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     } else {
       EnvNumTeams = -1;
     }
 
-    // By default let's create 256 streams per device
-    EnvNumStreams = 256;
-    envStr = getenv("LIBOMPTARGET_NUM_STREAMS");
-    if (envStr) {
-      EnvNumStreams = std::stoi(envStr);
-    }
-
-    // Initialize streams for each device
-    for (std::vector<CUstream> &S : Streams) {
-      S.resize(EnvNumStreams);
-    }
-
-    // Initialize the next stream id
-    for (std::unique_ptr<std::atomic_uint> &Ptr : NextStreamId) {
-      Ptr = std::make_unique<std::atomic_uint>(0);
-    }
+    StreamManager =
+        std::make_shared<StreamManagerTy>(NumberOfDevices, &Contexts);
 
     // Default state.
     RequiresFlags = OMP_REQ_UNDEFINED;
   }
 
   ~RTLDeviceInfoTy() {
+    // First destruct stream manager in case of Contexts is destructed before it
+    StreamManager = nullptr;
+
     // Close modules
     for (auto &module : Modules)
       if (module) {
         CUresult err = cuModuleUnload(module);
         if (err != CUDA_SUCCESS) {
           DP("Error when unloading CUDA module\n");
           CUDA_ERR_STRING(err);
         }
       }
 
-    // Destroy streams before contexts
-    for (int I = 0; I < NumberOfDevices; ++I) {
-      CUresult err = cuCtxSetCurrent(Contexts[I]);
-      if (err != CUDA_SUCCESS) {
-        DP("Error when setting current CUDA context\n");
-        CUDA_ERR_STRING(err);
-      }
-
-      for (auto &S : Streams[I])
-        if (S) {
-          err = cuStreamDestroy(S);
-          if (err != CUDA_SUCCESS) {
-            DP("Error when destroying CUDA stream\n");
-            CUDA_ERR_STRING(err);
-          }
-        }
-    }
-
     // Destroy contexts
     for (auto &ctx : Contexts)
       if (ctx) {
         CUresult err = cuCtxDestroy(ctx);
         if (err != CUDA_SUCCESS) {
           DP("Error when destroying CUDA context\n");
           CUDA_ERR_STRING(err);
         }
       }
   }
 };
 
 static RTLDeviceInfoTy DeviceInfo;
 
 namespace {
-CUstream selectStream(int32_t Id, __tgt_async_info *AsyncInfo) {
-  if (!AsyncInfo)
-    return DeviceInfo.getNextStream(Id);
+CUstream getStream(int32_t DeviceId, __tgt_async_info *AsyncInfoPtr) {
+  assert(AsyncInfoPtr && "AsyncInfoPtr is nullptr");
 
-  if (!AsyncInfo->Queue)
-    AsyncInfo->Queue = DeviceInfo.getNextStream(Id);
+  if (!AsyncInfoPtr->Queue)
+    AsyncInfoPtr->Queue = DeviceInfo.getStream(DeviceId);
 
-  return reinterpret_cast<CUstream>(AsyncInfo->Queue);
+  return reinterpret_cast<CUstream>(AsyncInfoPtr->Queue);
 }
 
 int32_t dataRetrieve(int32_t DeviceId, void *HstPtr, void *TgtPtr, int64_t Size,
                      __tgt_async_info *AsyncInfoPtr) {
   assert(AsyncInfoPtr && "AsyncInfoPtr is nullptr");
   // Set the context we are using.
   CUresult err = cuCtxSetCurrent(DeviceInfo.Contexts[DeviceId]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting CUDA context\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
-  CUstream Stream = selectStream(DeviceId, AsyncInfoPtr);
+  CUstream Stream = getStream(DeviceId, AsyncInfoPtr);
 
   err = cuMemcpyDtoHAsync(HstPtr, (CUdeviceptr)TgtPtr, Size, Stream);
   if (err != CUDA_SUCCESS) {
     DP("Error when copying data from device to host. Pointers: host = " DPxMOD
        ", device = " DPxMOD ", size = %" PRId64 "\n",
        DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
   return OFFLOAD_SUCCESS;
 }
 
 int32_t dataSubmit(int32_t DeviceId, void *TgtPtr, void *HstPtr, int64_t Size,
                    __tgt_async_info *AsyncInfoPtr) {
   assert(AsyncInfoPtr && "AsyncInfoPtr is nullptr");
   // Set the context we are using.
   CUresult err = cuCtxSetCurrent(DeviceInfo.Contexts[DeviceId]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting CUDA context\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
 
-  CUstream Stream = selectStream(DeviceId, AsyncInfoPtr);
+  CUstream Stream = getStream(DeviceId, AsyncInfoPtr);
 
   err = cuMemcpyHtoDAsync((CUdeviceptr)TgtPtr, HstPtr, Size, Stream);
   if (err != CUDA_SUCCESS) {
     DP("Error when copying data from host to device. Pointers: host = " DPxMOD
        ", device = " DPxMOD ", size = %" PRId64 "\n",
        DPxPTR(HstPtr), DPxPTR(TgtPtr), Size);
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
int32_t __tgt_rtl_init_device(int32_t device_id) {
   }
 
   err = cuCtxSetCurrent(DeviceInfo.Contexts[device_id]);
   if (err != CUDA_SUCCESS) {
     DP("Error when setting current CUDA context\n");
     CUDA_ERR_STRING(err);
   }
 
-  for (CUstream &Stream : DeviceInfo.Streams[device_id]) {
+  for (CUstream &Stream : DeviceInfo.StreamManager->StreamPool[device_id]) {
     err = cuStreamCreate(&Stream, CU_STREAM_NON_BLOCKING);
     if (err != CUDA_SUCCESS) {
       DP("Error when creating CUDA stream\n");
       CUDA_ERR_STRING(err);
     }
   }
 
   // Query attributes to determine number of threads/block and blocks/grid.
   int maxGridDimX;
   err = cuDeviceGetAttribute(&maxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
                              cuDevice);
   if (err != CUDA_SUCCESS) {
     DP("Error getting max grid dimension, use default\n");
if (team_num <= 0) {
     cudaBlocksPerGrid = team_num;
     DP("Using requested number of teams %d\n", team_num);
   }
 
   // Run on the device.
   DP("Launch kernel with %d blocks and %d threads\n", cudaBlocksPerGrid,
      cudaThreadsPerBlock);
 
-  CUstream Stream = selectStream(device_id, async_info);
+  CUstream Stream = getStream(device_id, async_info);
   err = cuLaunchKernel(KernelInfo->Func, cudaBlocksPerGrid, 1, 1,
                        cudaThreadsPerBlock, 1, 1, 0 /*bytes of shared memory*/,
                        Stream, &args[0], 0);
   if (err != CUDA_SUCCESS) {
     DP("Device kernel launch failed!\n");
     CUDA_ERR_STRING(err);
     return OFFLOAD_FAIL;
   }
int32_t __tgt_rtl_synchronize(int32_t device_id, __tgt_async_info *async_info) {
   CUresult Err = cuStreamSynchronize(Stream);
   if (Err != CUDA_SUCCESS) {
     DP("Error when synchronizing stream. stream = " DPxMOD
        ", async info ptr = " DPxMOD "\n",
        DPxPTR(Stream), DPxPTR(async_info));
     CUDA_ERR_STRING(Err);
     return OFFLOAD_FAIL;
   }
+
+  // Once the stream is synchronized, return it to stream pool and reset
+  // async_info. This is to make sure the synchronization only works for its own
+  // tasks.
+  DeviceInfo.returnStream(device_id, async_info);
+
   return OFFLOAD_SUCCESS;
 }
 
 #ifdef __cplusplus
 }
 #endif