place.cc 1.1 KB
Newer Older
1
#include "paddle/majel/place.h"
Y
Yi Wang 已提交
2 3 4 5 6

namespace majel {

namespace detail {

L
liaogang 已提交
7
class PlacePrinter : public boost::static_visitor<> {
Y
Yi Wang 已提交
8
private:
L
liaogang 已提交
9 10
  std::ostream& os_;

Y
Yi Wang 已提交
11
public:
L
liaogang 已提交
12
  PlacePrinter(std::ostream& os) : os_(os) {}
Y
Yi Wang 已提交
13

L
liaogang 已提交
14
  void operator()(const CpuPlace&) { os_ << "CpuPlace"; }
Y
Yi Wang 已提交
15

L
liaogang 已提交
16
  void operator()(const GpuPlace& p) { os_ << "GpuPlace(" << p.device << ")"; }
Y
Yi Wang 已提交
17 18
};

19
}  // namespace detail
Y
Yi Wang 已提交
20 21 22

static Place the_default_place;

L
liaogang 已提交
23
void set_place(const Place& place) { the_default_place = place; }
Y
Yi Wang 已提交
24

L
liaogang 已提交
25
const Place& get_place() { return the_default_place; }
Y
Yi Wang 已提交
26

L
liaogang 已提交
27
const GpuPlace default_gpu() { return GpuPlace(0); }
Y
Yi Wang 已提交
28

L
liaogang 已提交
29
const CpuPlace default_cpu() { return CpuPlace(); }
Y
Yi Wang 已提交
30 31

bool is_gpu_place(const Place& p) {
L
liaogang 已提交
32
  return boost::apply_visitor(IsGpuPlace(), p);
Y
Yi Wang 已提交
33 34 35
}

bool is_cpu_place(const Place& p) {
L
liaogang 已提交
36
  return !boost::apply_visitor(IsGpuPlace(), p);
Y
Yi Wang 已提交
37 38 39
}

bool places_are_same_class(const Place& p1, const Place& p2) {
L
liaogang 已提交
40
  return is_gpu_place(p1) == is_gpu_place(p2);
Y
Yi Wang 已提交
41 42 43
}

std::ostream& operator<<(std::ostream& os, const majel::Place& p) {
L
liaogang 已提交
44 45 46
  majel::detail::PlacePrinter printer(os);
  boost::apply_visitor(printer, p);
  return os;
Y
Yi Wang 已提交
47 48 49
}

}  // namespace majel