place.h 2.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
#include <iostream>
Y
Yi Wang 已提交
18 19
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/variant.h"
Y
Yi Wang 已提交
20

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

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

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

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

Y
Yu Yang 已提交
38
  inline int GetDeviceId() const { return device; }
Y
Yi Wang 已提交
39
  // needed for variant equality comparison
D
dzhwinter 已提交
40 41 42 43
  inline bool operator==(const CUDAPlace &o) const {
    return device == o.device;
  }
  inline bool operator!=(const CUDAPlace &o) const { return !(*this == o); }
Y
Yi Wang 已提交
44 45 46 47

  int device;
};

D
dzhwinter 已提交
48
struct IsCUDAPlace : public boost::static_visitor<bool> {
49
  bool operator()(const CPUPlace &) const { return false; }
D
dzhwinter 已提交
50
  bool operator()(const CUDAPlace &gpu) const { return true; }
T
tensor-tang 已提交
51 52
};

Q
QI JUN 已提交
53
typedef boost::variant<CUDAPlace, CPUPlace> Place;
Y
Yi Wang 已提交
54

Y
Yang Yu 已提交
55 56
using PlaceList = std::vector<Place>;

Y
Yi Wang 已提交
57 58
void set_place(const Place &);
const Place &get_place();
Y
Yi Wang 已提交
59

D
dzhwinter 已提交
60
const CUDAPlace default_gpu();
61
const CPUPlace default_cpu();
Y
Yi Wang 已提交
62

Y
Yi Wang 已提交
63 64 65
bool is_gpu_place(const Place &);
bool is_cpu_place(const Place &);
bool places_are_same_class(const Place &, const Place &);
66
bool is_same_place(const Place &, const Place &);
Y
Yi Wang 已提交
67

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

Y
Yang Yu 已提交
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
template <typename Visitor>
struct PlaceVisitorWrapper
    : public boost::static_visitor<typename Visitor::result_type> {
  const Visitor &visitor_;
  explicit PlaceVisitorWrapper(const Visitor &visitor) : visitor_(visitor) {}

  typename Visitor::result_type operator()(const CPUPlace &cpu) const {
    return visitor_(cpu);
  }

  typename Visitor::result_type operator()(const CUDAPlace &cuda) const {
#ifdef PADDLE_WITH_CUDA
    return visitor_(cuda);
#else
    PADDLE_THROW("Paddle is not compiled with CUDA. Cannot visit cuda device");
    return typename Visitor::result_type();
#endif
  }
};

template <typename Visitor>
typename Visitor::result_type VisitPlace(const Place &place,
                                         const Visitor &visitor) {
  return boost::apply_visitor(PlaceVisitorWrapper<Visitor>(visitor), place);
}

Y
Yi Wang 已提交
96 97
}  // namespace platform
}  // namespace paddle