device.h

Go to the documentation of this file.

#ifndef MODULE_DEVICE_H_
#define MODULE_DEVICE_H_
 
#include "types.h"
#include "device_helpers.h"
#include <fstream>
#include <mutex>
 
#ifdef __MPI
#include "mpi.h"
#endif
 
namespace base_device
{
 
namespace information
{
 
std::string get_device_name(std::string device_flag);
 
int get_device_num(std::string device_flag);
 
void output_device_info(std::ostream& output, const std::string& device);
 
bool probe_gpu_availability();
 
std::string get_device_flag(const std::string& device,
                            const std::string& basis_type);
 
#if __MPI
int get_node_rank_with_mpi_shared(const MPI_Comm mpi_comm = MPI_COMM_WORLD);
#endif
 
template <typename Device>
void print_device_info(const Device* dev, std::ofstream& ofs_device)
{
    return;
}
 
template <typename Device>
void record_device_memory(const Device* dev, std::ofstream& ofs_device, std::string str, size_t size)
{
    return;
}
 
#if defined(__CUDA) || defined(__ROCM)
template <>
void print_device_info<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU *ctx, std::ofstream &ofs_device);
 
template <>
void record_device_memory<base_device::DEVICE_GPU>(const base_device::DEVICE_GPU* dev, std::ofstream& ofs_device, std::string str, size_t size);
#endif
 
} // end of namespace information
 
class DeviceContext {
public:
    static DeviceContext& instance();
 
    void init();
 
    int get_device_id() const { return device_id_; }
 
    int get_device_count() const { return device_count_; }
 
    int get_local_rank() const { return local_rank_; }
 
    // Disable copy and assignment
    DeviceContext(const DeviceContext&) = delete;
    DeviceContext& operator=(const DeviceContext&) = delete;
 
private:
    DeviceContext() = default;
    ~DeviceContext() = default;
 
    bool initialized_ = false;
    bool gpu_enabled_ = false;
    int device_id_ = -1;
    int device_count_ = 0;
    int local_rank_ = 0;
    std::mutex init_mutex_;
};
 
} // end of namespace base_device
 
#endif // MODULE_DEVICE_H_