device.h

#ifndef SINGA_CORE_DEVICE_H_
#define SINGA_CORE_DEVICE_H_

#include <type_traits>
#include <vector>
#include <string>
#include <functional>
#include <memory>

#include "singa/singa_config.h"
#include "singa/core/common.h"
#include "singa/core/memory.h"
#include "singa/core/scheduler.h"
#include "singa/proto/core.pb.h"

#ifdef USE_CUDA
#include "singa/utils/cuda_utils.h"
#endif // USE_CUDA

#ifdef USE_OPENCL
#include "singa/utils/opencl_utils.h"
#endif // USE_OPENCL

#ifdef USE_MKLDNN
#include "singa/utils/mkldnn_utils.h"
#endif  // USE_MKLDNN

using std::vector;
using std::string;
using std::function;
using std::shared_ptr;

namespace singa {

class Device {
  public:
  // Device() = default;
  virtual ~Device() {}
  Device(int id, int num_executors);

  virtual void SetRandSeed(unsigned seed) = 0;

  Block* NewBlock(int size);

  void FreeBlock(Block* block);

  virtual size_t GetAllocatedMem() {
    return 0u;
  }

  virtual void CopyDataToFrom(Block* dst, Block* src, size_t nBytes,
                      CopyDirection direction, int dst_offset, int src_offset);

  void CopyDataFromHostPtr(Block* dst, const void* src, size_t nBytes,
                           size_t dst_offset = 0);
  void Exec(function<void(Context*)>&& fn, const vector<Block*> read_blocks,
                    const vector<Block*> write_blocks,
                    bool use_rand_generator = false);

  // Wait for one event.
  // void WaitFor();

  void Sync();

  LangType lang() const {
    return lang_;
  }

  virtual std::shared_ptr<Device> host() const { return host_;}

  Context* context(int k) {
    return &ctx_;
  }

  int id() const { return id_; }

 private:
  Device() {};

 protected:
  virtual void DoExec(function<void(Context*)>&& fn, int executor) = 0;

  virtual void CopyToFrom(void* dst, const void* src, size_t nBytes,
                          CopyDirection direction, Context* ctx) = 0;

  virtual void* Malloc(int size) = 0;

  virtual void Free(void* ptr) = 0;

 protected:
  int id_ = 0;
  int num_executors_ = 0;
  unsigned seed_ = 0;
  // Scheduler* scheduler_ = nullptr;
  // VirtualMemory* vm_ = nullptr;
  LangType lang_;
  // SafeQueue<Operation> op_queue_;
  // SafeQueue<Operation> op_log_;
  std::shared_ptr<Device> host_;
  // TODO(wangwei) define multiple contexts, one per executor
  Context ctx_;
};

extern std::shared_ptr<Device> defaultDevice;

class CppCPU : public Device {
 public:
  ~CppCPU();
  CppCPU();

  std::shared_ptr<Device> host() const override { return defaultDevice;}
  void SetRandSeed(unsigned seed) override;

 protected:
  void DoExec(function<void(Context*)>&& fn, int executor) override;

  void CopyToFrom(void* dst, const void* src, size_t nBytes,
                  CopyDirection direction, Context* ctx) override;

  void* Malloc(int size) override;

  void Free(void* ptr) override;
};


// Implement Device using OpenCL libs.
// class OpenclDevice : public Device { };

#ifdef USE_CUDA
// Represent a Nvidia GPU which runs cuda code.
class CudaGPU : public Device {
 public:
  ~CudaGPU();
  CudaGPU(int id = 0);
  CudaGPU(int id, std::shared_ptr<DeviceMemPool> pool);

  void SetRandSeed(unsigned seed) override;
  size_t GetAllocatedMem() override;

 protected:
  void DoExec(function<void(Context*)>&& fn, int executor) override;

  void CopyToFrom(void* dst, const void* src, size_t nBytes,
                  CopyDirection direction, Context* ctx) override;

  void* Malloc(int size) override;

  void Free(void* ptr) override;

 private:
  void Setup();

 private:
    shared_ptr<DeviceMemPool> pool_;
};


#endif  // USE_CUDA

#ifdef USE_OPENCL

// Implement Device using OpenCL libs.
class OpenclDevice : public singa::Device {
public:

  // TODO: Constructor arguments to consider:
  // Path to kernel sources?
  // Select only certain device types?
  OpenclDevice(int id = 0, int num_executors = 1);
  ~OpenclDevice();

// Overridden, inherited methods
  void SetRandSeed(unsigned seed) override;

  virtual void CopyDataToFrom(Block* dst, Block* src, size_t nBytes,
                      CopyDirection direction, int dst_offset = 0,
                      int src_offset = 0) override;

protected:
  viennacl::ocl::device this_device;

  viennacl::ocl::context vcl_ctx;

  void BuildPrograms();

// Overridden, inherited methods.

  void DoExec(function<void(Context*)>&& fn, int executor) override;

  void CopyToFrom(void* dst, const void* src, size_t nBytes,
                  CopyDirection direction, Context* ctx = nullptr) override;

  void* Malloc(int size) override;

  void Free(void* ptr) override;

private:

  static const std::string cl_src_path;
};
#endif  // USE_OPENCL

class Platform {
public:

  static std::shared_ptr<Device> GetDefaultDevice() {
    return defaultDevice;
  }

#ifdef USE_CUDA
  static int GetNumGPUs();

  static const std::vector<int> GetGPUIDs();

  static const std::pair<size_t, size_t> GetGPUMemSize(const int device);

  static const std::vector<std::pair<size_t, size_t>> GetGPUMemSize();

  static const std::string DeviceQuery(int id, bool verbose = false);

  static const std::vector<std::shared_ptr<Device>>
  CreateCudaGPUs(const size_t num_devices, size_t init_size = 0);

  static const std::vector<std::shared_ptr<Device>>
  CreateCudaGPUsOn(const std::vector<int> &devices, size_t init_size = 0);
  
  static bool CheckDevice(const int device_id);
#endif // USE_CUDA

#ifdef USE_OPENCL

  const int GetNumOpenclPlatforms();
  
  const int GetNumOpenclDevices();
  
  static const std::shared_ptr<Device> GetDefaultOpenclDevice();

//  static const std::vector<std::shared_ptr<Device>>
//  CreateOpenclDevices(const size_t num_devices);

//  const std::vector<std::shared_ptr<Device>>
//  CreateOpenclDevices(const vector<int> &id);
#endif // USE_OPENCL

};


}  // namespace singa

#endif  // SINGA_CORE_DEVICE_H_