Commit 7af4a1c1 authored by Rick Ho's avatar Rick Ho
Browse files

limit max streams support

parent bf60a846
#include <cassert> #include <cuda_runtime.h>
#include "cuda_stream_manager.h" #include "cuda_stream_manager.h"
...@@ -8,6 +8,11 @@ CudaStreamManager* getCudaStreamManager(const size_t num_expert) { ...@@ -8,6 +8,11 @@ CudaStreamManager* getCudaStreamManager(const size_t num_expert) {
if (!smgr) { if (!smgr) {
smgr = new CudaStreamManager(num_expert); smgr = new CudaStreamManager(num_expert);
} }
assert(smgr->num_expert == num_expert);
return smgr; return smgr;
} }
void CudaStreamManager::sync() {
for (size_t i=0; i<MAX_STREAMS; ++i) {
cudaStreamSynchronize(streams[i]);
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment