place.h 1.3 KB
Newer Older
Y
Yi Wang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
#pragma once
#include <boost/variant.hpp>
#include <iostream>

namespace majel {

struct CpuPlace {
  CpuPlace() {}  // WORKAROUND: for some reason, omitting this constructor
                 // causes errors with boost 1.59 and OSX
  // needed for variant equality comparison
  inline bool operator==(const CpuPlace&) const { return true; }

  inline bool operator!=(const CpuPlace&) const { return false; }
};

struct GpuPlace {
  GpuPlace(int d) : device(d) {}

  // needed for variant equality comparison
  inline bool operator==(const GpuPlace& o) const { return device == o.device; }

  inline bool operator!=(const GpuPlace& o) const { return !(*this == o); }

  GpuPlace() : GpuPlace(0) {}
  int device;
};

class IsGpuPlace : public boost::static_visitor<bool> {
public:
  bool operator()(const CpuPlace&) const { return false; }

  bool operator()(const GpuPlace& gpu) const { return true; }
};

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

void set_place(const Place&);

const Place& get_place();

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

bool is_gpu_place(const Place&);
bool is_cpu_place(const Place&);
bool places_are_same_class(const Place&, const Place&);

std::ostream& operator<<(std::ostream&, const majel::Place&);

}  // namespace majel