cuda_stream_manager.h 581 Bytes
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>

class CudaStreamManager {
Rick Ho's avatar
Rick Ho committed
8
9
10
11
12
public:
    int device;
    cublasHandle_t* handles;
    cudaStream_t* streams;

Rick Ho's avatar
Rick Ho committed
13
public:
14
15
    CudaStreamManager(int device_): device(device_) {
		this->setup(device);
16
17
    }

18
19
20
21
	void setup(int);
	void sync(int=0);
	void destroy();

Rick Ho's avatar
Rick Ho committed
22
	cudaStream_t stream(size_t=0);
23
	cublasHandle_t handle(size_t=0);
24

Rick Ho's avatar
Rick Ho committed
25
    ~CudaStreamManager() {
26
		this->destroy();
Rick Ho's avatar
Rick Ho committed
27
28
29
    }
}; 

30
CudaStreamManager* getCudaStreamManager(const int device);
Rick Ho's avatar
Rick Ho committed
31
32

#endif  // CUDA_STREAM_MANAGER