stream_manager.h

// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
#include "hip/hip_runtime.h"
#ifndef CUDA_STREAM_MANAGER_H
#define CUDA_STREAM_MANAGER_H

#include "../hip/utils/helper_cuda.h"

#ifdef FMOE_USE_NCCL
#include <rccl/rccl.h>

#define NCCL_SAFE_CALL(__fn__) { \
    auto __res__ = __fn__; \
    if (__res__ != ncclSuccess) { \
        fprintf(stderr, "NCCL Error at %s:%d value %d\n", __FILE__, __LINE__, __res__); \
        exit(-1); \
    } \
}

#endif

class CudaStreamManager {
public:
    int device;
    hipblasHandle_t* handles;
    hipStream_t* streams;
    bool use_default;
#ifdef FMOE_USE_NCCL
    char ncclgood;
    ncclComm_t ncclcomm;
#endif

public:
    CudaStreamManager(int device_): device(device_), use_default(false) {
        this->setup(device);
    }

    void setup(int);
    void sync(int=0);
    void syncTorch();
    void destroy();

    hipStream_t torchStream();
    hipStream_t stream(size_t=0);
    hipblasHandle_t handle(size_t=0);

    ~CudaStreamManager() {
        this->destroy();
    }
}; 

CudaStreamManager* getCudaStreamManager(const int device);

#endif  // CUDA_STREAM_MANAGER