place.cu 1.1 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 51 52 53 54 55 56 57 58 59 60 61
#include <majel/place.h>

namespace majel {

namespace detail {

class PlacePrinter
    : public boost::static_visitor<> {
private:
    std::ostream& os_;
public:
    PlacePrinter(std::ostream& os) : os_(os) {}

    void operator()(const CpuPlace&) {
        os_ << "CpuPlace";
    }

    void operator()(const GpuPlace& p) {
        os_ << "GpuPlace(" << p.device << ")";
    }
};

}  // namespace majel

static Place the_default_place;

void set_place(const Place& place) {
    the_default_place = place;
}

const Place& get_place() {
    return the_default_place;
}

const GpuPlace default_gpu() {
    return GpuPlace(0);
}

const CpuPlace default_cpu() {
    return CpuPlace();
}

bool is_gpu_place(const Place& p) {
    return boost::apply_visitor(IsGpuPlace(), p);
}

bool is_cpu_place(const Place& p) {
    return !boost::apply_visitor(IsGpuPlace(), p);
}

bool places_are_same_class(const Place& p1, const Place& p2) {
    return is_gpu_place(p1) == is_gpu_place(p2);
}

std::ostream& operator<<(std::ostream& os, const majel::Place& p) {
    majel::detail::PlacePrinter printer(os);
    boost::apply_visitor(printer, p);
    return os;
}

}  // namespace majel