cuda_stream_manager.h 1.17 KB
Newer Older
Rick Ho's avatar
Rick Ho committed
1
2
3
4
5
6
7
#ifndef CUDA_STREAM_MANAGER_H
#define CUDA_STREAM_MANAGER_H

#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <helper_cuda.h> 

8
9
#include <cstdio>

Rick Ho's avatar
Rick Ho committed
10
11

class CudaStreamManager {
Rick Ho's avatar
Rick Ho committed
12
13
14
15
16
17
public:
    size_t num_expert;
    int device;
    cublasHandle_t* handles;
    cudaStream_t* streams;

Rick Ho's avatar
Rick Ho committed
18
public:
Rick Ho's avatar
Rick Ho committed
19
    CudaStreamManager() : num_expert(0), streams(NULL) {
20
21
        int current_device;
        checkCudaErrors(cudaGetDevice(&current_device));
22
23
24
25
26
#ifdef MOE_DEBUG
        printf("constructor at device %d\n", current_device);
#endif
    }

Rick Ho's avatar
Rick Ho committed
27
    void setup(const size_t num_expert, const int device=-1);
Rick Ho's avatar
Rick Ho committed
28
	cudaStream_t stream(size_t=0);
29

Rick Ho's avatar
Rick Ho committed
30
    ~CudaStreamManager() {
31
32
33
#ifdef MOE_DEBUG
        printf("destructor at device %d\n", device);
#endif
Rick Ho's avatar
Rick Ho committed
34
35
        for (size_t i=0; i<num_expert; ++i) {
            checkCudaErrors(cudaStreamDestroy(*(streams+i)));
Rick Ho's avatar
Rick Ho committed
36
37
			checkCudaErrors(cublasDestroy(handles[i]));
		}
Jiezhong Qiu's avatar
Jiezhong Qiu committed
38
        delete[] streams;
Rick Ho's avatar
Rick Ho committed
39
    }
Rick Ho's avatar
Rick Ho committed
40
41

	void sync(int=-1);
Rick Ho's avatar
Rick Ho committed
42
43
}; 

Rick Ho's avatar
Rick Ho committed
44
45
46
47
48
49
#define ENSURE_SMGR(__smgr__, __num_expert__) { \
	if (__smgr__.num_expert == 0) { \
		__smgr__.setup(__num_expert__); \
	} \
}

50
// CudaStreamManager* getCudaStreamManager(const size_t num_expert, const int device);
Rick Ho's avatar
Rick Ho committed
51
52

#endif  // CUDA_STREAM_MANAGER