place.h 1.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

Y
Yi Wang 已提交
15
#pragma once
16

Y
Yi Wang 已提交
17 18 19
#include <boost/variant.hpp>
#include <iostream>

Y
Yi Wang 已提交
20 21
namespace paddle {
namespace platform {
Y
Yi Wang 已提交
22

23
struct CPUPlace {
Y
Yi Wang 已提交
24 25
  // WORKAROUND: for some reason, omitting this constructor
  // causes errors with boost 1.59 and OSX
26
  CPUPlace() {}
Y
Yi Wang 已提交
27

Y
Yi Wang 已提交
28
  // needed for variant equality comparison
29 30
  inline bool operator==(const CPUPlace &) const { return true; }
  inline bool operator!=(const CPUPlace &) const { return false; }
Y
Yi Wang 已提交
31 32
};

33 34
struct GPUPlace {
  GPUPlace() : GPUPlace(0) {}
Y
Yi Wang 已提交
35
  explicit GPUPlace(int d) : device(d) {}
Y
Yi Wang 已提交
36 37

  // needed for variant equality comparison
38 39
  inline bool operator==(const GPUPlace &o) const { return device == o.device; }
  inline bool operator!=(const GPUPlace &o) const { return !(*this == o); }
Y
Yi Wang 已提交
40 41 42 43

  int device;
};

44 45 46
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 已提交
47 48
};

49
typedef boost::variant<GPUPlace, CPUPlace> Place;
Y
Yi Wang 已提交
50

Y
Yi Wang 已提交
51 52
void set_place(const Place &);
const Place &get_place();
Y
Yi Wang 已提交
53

54 55
const GPUPlace default_gpu();
const CPUPlace default_cpu();
Y
Yi Wang 已提交
56

Y
Yi Wang 已提交
57 58 59
bool is_gpu_place(const Place &);
bool is_cpu_place(const Place &);
bool places_are_same_class(const Place &, const Place &);
Y
Yi Wang 已提交
60

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

Y
Yi Wang 已提交
63 64
}  // namespace platform
}  // namespace paddle