place.h 1.3 KB
Newer Older
Y
Yi Wang 已提交
1 2 3 4
#pragma once
#include <boost/variant.hpp>
#include <iostream>

Y
Yi Wang 已提交
5 6
namespace paddle {
namespace platform {
Y
Yi Wang 已提交
7 8

struct CpuPlace {
Y
Yi Wang 已提交
9 10 11
  // WORKAROUND: for some reason, omitting this constructor
  // causes errors with boost 1.59 and OSX
  CpuPlace() {}
Y
Yi Wang 已提交
12

Y
Yi Wang 已提交
13 14 15
  // needed for variant equality comparison
  inline bool operator==(const CpuPlace &) const { return true; }
  inline bool operator!=(const CpuPlace &) const { return false; }
Y
Yi Wang 已提交
16 17 18
};

struct GpuPlace {
Y
Yi Wang 已提交
19
  GpuPlace() : GpuPlace(0) {}
Y
Yi Wang 已提交
20 21 22
  GpuPlace(int d) : device(d) {}

  // needed for variant equality comparison
Y
Yi Wang 已提交
23 24
  inline bool operator==(const GpuPlace &o) const { return device == o.device; }
  inline bool operator!=(const GpuPlace &o) const { return !(*this == o); }
Y
Yi Wang 已提交
25 26 27 28

  int device;
};

Y
Yi Wang 已提交
29 30 31
struct IsGpuPlace : public boost::static_visitor<bool> {
  bool operator()(const CpuPlace &) const { return false; }
  bool operator()(const GpuPlace &gpu) const { return true; }
Y
Yi Wang 已提交
32 33 34 35
};

typedef boost::variant<GpuPlace, CpuPlace> Place;

Y
Yi Wang 已提交
36 37
void set_place(const Place &);
const Place &get_place();
Y
Yi Wang 已提交
38 39 40 41

const GpuPlace default_gpu();
const CpuPlace default_cpu();

Y
Yi Wang 已提交
42 43 44
bool is_gpu_place(const Place &);
bool is_cpu_place(const Place &);
bool places_are_same_class(const Place &, const Place &);
Y
Yi Wang 已提交
45

Y
Yi Wang 已提交
46
std::ostream &operator<<(std::ostream &, const Place &);
Y
Yi Wang 已提交
47

Y
Yi Wang 已提交
48 49
} // namespace platform
} // namespace paddle