pybind.cc 18.9 KB
Newer Older
1 2 3 4 5 6
/* 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

7
http://www.apache.org/licenses/LICENSE-2.0
8 9 10 11 12 13 14

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. */

Q
qijun 已提交
15 16
#include "paddle/pybind/protobuf.h"

Q
QI JUN 已提交
17
#include <mutex>  // for call_once
18
#include <unordered_map>
Q
Qiao Longfei 已提交
19
#include "paddle/framework/backward.h"
F
fengjiayi 已提交
20
#include "paddle/framework/executor.h"
Q
qijun 已提交
21
#include "paddle/framework/feed_fetch_method.h"
22
#include "paddle/framework/framework.pb.h"
D
dzhwinter 已提交
23
#include "paddle/framework/init.h"
Y
Yu Yang 已提交
24
#include "paddle/framework/lod_rank_table.h"
D
dangqingqing 已提交
25
#include "paddle/framework/lod_tensor.h"
Y
Yu Yang 已提交
26
#include "paddle/framework/lod_tensor_array.h"
27
#include "paddle/framework/prune.h"
Q
qijun 已提交
28
#include "paddle/framework/selected_rows.h"
Z
zchen0211 已提交
29
#include "paddle/operators/cond_op.h"
Y
Yan Chunwei 已提交
30
#include "paddle/operators/net_op.h"
Q
qijun 已提交
31
#include "paddle/platform/enforce.h"
Q
qijun 已提交
32
#include "paddle/platform/place.h"
33
#include "paddle/platform/profiler.h"
Q
qiaolongfei 已提交
34
#include "paddle/pybind/const_value.h"
Y
Yu Yang 已提交
35
#include "paddle/pybind/exception.h"
Q
qijun 已提交
36
#include "paddle/pybind/pybind.h"
37
#include "paddle/pybind/tensor_py.h"
38
#include "paddle/string/to_string.h"
39

D
Dong Zhihong 已提交
40 41
#ifdef PADDLE_WITH_CUDA
#include "paddle/operators/nccl/nccl_gpu_common.h"
D
dangqingqing 已提交
42
#include "paddle/platform/cuda_profiler.h"
D
Dong Zhihong 已提交
43
#include "paddle/platform/gpu_info.h"
D
Dong Zhihong 已提交
44 45
#endif

Q
Qiao Longfei 已提交
46 47 48
// disable auto conversion to list in Python
PYBIND11_MAKE_OPAQUE(paddle::framework::LoDTensorArray);

49
namespace paddle {
50
namespace pybind {
51 52 53
static size_t UniqueIntegerGenerator(const std::string &prefix) {
  static std::unordered_map<std::string, std::atomic<size_t>> generators;
  return generators[prefix].fetch_add(1);
54 55
}

56
bool IsCompiledWithCUDA() {
57
#ifndef PADDLE_WITH_CUDA
Q
qijun 已提交
58 59 60 61 62 63
  return false;
#else
  return true;
#endif
}

64 65
PYBIND11_PLUGIN(core) {
  py::module m("core", "C++ core of PaddlePaddle");
66

67 68 69 70
  // using framework in this function. Since it is inside a function, it will
  // not cause namespace pollution.
  using namespace paddle::framework;  // NOLINT

Y
Yu Yang 已提交
71 72
  BindException(m);

73 74 75
  py::class_<Tensor>(m, "Tensor", py::buffer_protocol())
      .def_buffer(
          [](Tensor &self) -> py::buffer_info { return CastToPyBuffer(self); })
Y
Yu Yang 已提交
76
      .def("get_dims",
77
           [](const Tensor &self) { return vectorize(self.dims()); })
Y
Yu Yang 已提交
78
      .def("set_dims",
Q
qijun 已提交
79
           [](Tensor &self, const std::vector<int64_t> &dim) {
Y
Yu Yang 已提交
80
             self.Resize(make_ddim(dim));
Y
Yu Yang 已提交
81
           })
D
dzhwinter 已提交
82 83 84 85
      .def("set_layout",
           [](Tensor &self, const std::string &layout) {
             self.set_layout(StringToDataLayout(layout));
           })
Y
Yu Yang 已提交
86
      .def("alloc_float",
D
dzhwinter 已提交
87
           [](Tensor &self, paddle::platform::CUDAPlace &place) {
Q
qijun 已提交
88
             self.mutable_data<float>(place);
Y
Yu Yang 已提交
89
           })
Q
qijun 已提交
90
      .def("alloc_float",
Y
Yu Yang 已提交
91
           [](Tensor &self, paddle::platform::CPUPlace &place) {
Q
qijun 已提交
92
             self.mutable_data<float>(place);
Y
Yu Yang 已提交
93 94
           })
      .def("alloc_int",
Y
Yu Yang 已提交
95
           [](Tensor &self, paddle::platform::CPUPlace &place) {
Q
qijun 已提交
96
             self.mutable_data<int>(place);
Y
Yu Yang 已提交
97
           })
Q
qijun 已提交
98
      .def("alloc_int",
D
dzhwinter 已提交
99
           [](Tensor &self, paddle::platform::CUDAPlace &place) {
Q
qijun 已提交
100
             self.mutable_data<int>(place);
Q
qijun 已提交
101
           })
Y
Yu Yang 已提交
102 103
      .def("set", PyCPUTensorSetFromArray<float>)
      .def("set", PyCPUTensorSetFromArray<int>)
104
      .def("set", PyCPUTensorSetFromArray<double>)
105
      .def("set", PyCPUTensorSetFromArray<int64_t>)
Y
Yu Yang 已提交
106
      .def("set", PyCPUTensorSetFromArray<bool>)
107
#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
108 109
      .def("set", PyCUDATensorSetFromArray<float>)
      .def("set", PyCUDATensorSetFromArray<int>)
110
      .def("set", PyCUDATensorSetFromArray<double>)
111
      .def("set", PyCUDATensorSetFromArray<int64_t>)
Y
Yu Yang 已提交
112
      .def("set", PyCUDATensorSetFromArray<bool>)
Q
qijun 已提交
113
#endif
114
      .def("shape", [](Tensor &self) { return vectorize(self.dims()); })
115 116 117 118 119
      .def("set_float_element", TensorSetElement<float>)
      .def("get_float_element", TensorGetElement<float>)
      .def("set_double_element", TensorSetElement<double>)
      .def("get_double_element", TensorGetElement<double>)
      .def("dtype", [](Tensor &self) { return ToDataType(self.type()); });
Y
Yu Yang 已提交
120

121
  py::class_<LoDTensor, Tensor>(m, "LoDTensor")
122 123
      .def_buffer(
          [](Tensor &self) -> py::buffer_info { return CastToPyBuffer(self); })
124 125 126
      .def(
          "__init__",
          [](LoDTensor &instance, const std::vector<std::vector<size_t>> &lod) {
D
dzhwinter 已提交
127 128 129 130
            LoD new_lod;
            new_lod.reserve(lod.size());
            std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
            new (&instance) LoDTensor(new_lod);
131
          })
Y
Yu Yang 已提交
132
      .def("__init__", [](LoDTensor &instance) { new (&instance) LoDTensor(); })
D
dangqingqing 已提交
133
      .def("set_lod",
134
           [](LoDTensor &self, const std::vector<std::vector<size_t>> &lod) {
Y
Yu Yang 已提交
135
             LoD new_lod;
136 137 138
             new_lod.reserve(lod.size());
             std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
             self.set_lod(new_lod);
D
dangqingqing 已提交
139
           })
140
      .def("lod", [](LoDTensor &self) -> std::vector<std::vector<size_t>> {
D
dzhwinter 已提交
141 142 143 144 145
        auto lod = self.lod();
        std::vector<std::vector<size_t>> new_lod;
        new_lod.reserve(lod.size());
        std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
        return new_lod;
D
dangqingqing 已提交
146 147
      });

Q
qijun 已提交
148 149 150 151 152 153 154 155 156 157 158 159 160
  py::class_<SelectedRows>(m, "SelectedRows")
      .def("__init__",
           [](SelectedRows &instance) { new (&instance) SelectedRows(); })
      .def("__init__",
           [](SelectedRows &instance, const std::vector<int64_t> rows,
              const int64_t &height) {
             new (&instance) SelectedRows(rows, height);
           })
      .def("get_tensor",
           [](SelectedRows &self) { return self.mutable_value(); },
           py::return_value_policy::reference)
      .def("set_height", &SelectedRows::set_height)
      .def("height", &SelectedRows::height)
Q
qijun 已提交
161 162 163 164 165 166 167 168 169
      .def("set_rows",
           [](SelectedRows &self, std::vector<int64_t> rows) {
#ifndef PADDLE_WITH_CUDA
             self.set_rows(rows);
#else
        Vector<int64_t> new_rows(rows);
        self.set_rows(new_rows);
#endif
           })
170 171 172 173 174 175 176 177 178 179 180
      .def("rows", [](SelectedRows &self) {
#ifndef PADDLE_WITH_CUDA
        return self.rows();
#else
         auto rows = self.rows();
         std::vector<int64_t> new_rows;
         new_rows.reserve(rows.size());
         std::copy(rows.begin(), rows.end(), std::back_inserter(new_rows));
         return new_rows;
#endif
      });
Q
qijun 已提交
181

182
  py::class_<Variable>(m, "Variable", R"DOC(Variable Class.
183 184 185

All parameter, weight, gradient are variables in Paddle.
)DOC")
186
      .def("is_int", [](const Variable &var) { return var.IsType<int>(); })
187
      .def("set_int",
188 189
           [](Variable &var, int val) -> void { *var.GetMutable<int>() = val; })
      .def("get_int", [](const Variable &var) -> int { return var.Get<int>(); })
190 191 192 193 194 195 196
      .def("is_float", [](const Variable &var) { return var.IsType<float>(); })
      .def("set_float",
           [](Variable &var, float val) -> void {
             *var.GetMutable<float>() = val;
           })
      .def("get_float",
           [](const Variable &var) -> float { return var.Get<float>(); })
Y
Yu Yang 已提交
197
      .def("get_tensor",
198 199
           [](Variable &self) -> LoDTensor * {
             return self.GetMutable<LoDTensor>();
D
dangqingqing 已提交
200 201
           },
           py::return_value_policy::reference)
Y
Yu Yang 已提交
202 203 204
      .def("get_lod_rank_table",
           [](Variable &self) { return self.GetMutable<LoDRankTable>(); },
           py::return_value_policy::reference)
Q
qijun 已提交
205 206 207 208 209
      .def("get_selected_rows",
           [](Variable &self) -> SelectedRows * {
             return self.GetMutable<SelectedRows>();
           },
           py::return_value_policy::reference)
Y
Yu Yang 已提交
210 211 212
      .def("get_lod_tensor_array",
           [](Variable &self) { return self.GetMutable<LoDTensorArray>(); },
           py::return_value_policy::reference)
D
Dong Zhihong 已提交
213 214 215 216 217 218 219
#ifdef PADDLE_WITH_CUDA
      .def("get_communicator",
           [](Variable &self) -> platform::Communicator * {
             return self.GetMutable<platform::Communicator>();
           },
           py::return_value_policy::reference)
#endif
Y
Yan Chunwei 已提交
220
      .def("get_net",
D
dongzhihong 已提交
221 222
           [](Variable &self) -> operators::NetOp * {
             return self.GetMutable<operators::NetOp>();
Y
Yan Chunwei 已提交
223
           },
Y
Yu Yang 已提交
224
           py::return_value_policy::reference);
225

226
  py::class_<Scope>(m, "Scope", "")
D
dongzhihong 已提交
227
      .def("var",
228
           [](Scope &self, const std::string &name) -> Variable * {
D
dongzhihong 已提交
229
             return self.Var(name);
Y
Yu Yang 已提交
230
           },
231
           py::return_value_policy::reference)
232
      .def("find_var", &Scope::FindVar, py::return_value_policy::reference)
Y
Yu Yang 已提交
233
      .def(py::init<>())
234
      .def("new_scope", [](Scope &self) -> Scope * { return &self.NewScope(); },
235
           py::return_value_policy::reference)
Y
Yu Yang 已提交
236
      .def("drop_kids", &Scope::DropKids);
237

Y
Yu Yang 已提交
238 239
  //! @note: Be careful! PyBind will return std::string as an unicode, not
  //! Python str. If you want a str object, you should cast them in Python.
Y
Yu Yang 已提交
240 241
  m.def("get_all_op_protos", []() -> std::vector<py::bytes> {
    std::vector<py::bytes> ret_values;
242 243 244 245 246 247 248 249 250 251
    for (auto &iter : OpInfoMap::Instance().map()) {
      auto &info = iter.second;
      if (info.HasOpProtoAndChecker()) {
        std::string str;
        PADDLE_ENFORCE(
            info.Proto().SerializeToString(&str),
            "Serialize OpProto Error. This could be a bug of Paddle.");
        ret_values.emplace_back(str);
      }
    }
Y
Yu Yang 已提交
252 253
    return ret_values;
  });
254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269
  m.def(
      "get_grad_op_desc", [](const OpDesc &op_desc,
                             const std::unordered_set<std::string> &no_grad_set,
                             const std::vector<BlockDesc *> &grad_sub_block) {
        std::unordered_map<std::string, std::string> grad_to_var;
        std::vector<std::unique_ptr<OpDesc>> grad_op_descs =
            framework::OpInfoMap::Instance()
                .Get(op_desc.Type())
                .GradOpMaker()(op_desc, no_grad_set, &grad_to_var,
                               grad_sub_block);
        std::vector<OpDesc *> grad_op_desc_ptrs(grad_op_descs.size());
        std::transform(grad_op_descs.begin(), grad_op_descs.end(),
                       grad_op_desc_ptrs.begin(),
                       [](std::unique_ptr<OpDesc> &p) { return p.release(); });
        return std::make_pair(grad_op_desc_ptrs, grad_to_var);
      });
Y
Yu Yang 已提交
270
  m.def("prune", [](const ProgramDesc &origin,
271
                    const std::vector<std::array<size_t, 2>> &targets) {
Y
Yu Yang 已提交
272
    ProgramDesc prog_with_targets(origin);
273
    for (const auto &t : targets) {
274
      prog_with_targets.MutableBlock(t[0])->Op(t[1])->MarkAsTarget();
275
    }
276
    proto::ProgramDesc pruned_desc;
277
    Prune(*prog_with_targets.Proto(), &pruned_desc);
Y
Yu Yang 已提交
278
    return new ProgramDesc(pruned_desc);
279
  });
Y
Yu Yang 已提交
280
  m.def("inference_optimize", [](ProgramDesc &origin) {
281
    proto::ProgramDesc pruned_desc;
282
    InferenceOptimize(*(origin.Proto()), &pruned_desc);
Y
Yu Yang 已提交
283
    return new ProgramDesc(pruned_desc);
284
  });
F
fengjiayi 已提交
285 286
  m.def("empty_var_name", []() { return framework::kEmptyVarName; });
  m.def("grad_var_suffix", []() { return framework::kGradVarSuffix; });
287 288 289
  m.def_submodule(
       "var_names",
       "The module will return special predefined variable name in Paddle")
Y
Yi Wang 已提交
290 291
      .def("empty", []() { return kEmptyVarName; })
      .def("temp", []() { return kTempVarName; });
Q
qijun 已提交
292
  // clang-format off
Y
Yu Yang 已提交
293
  py::class_<paddle::platform::DeviceContext>(m, "DeviceContext")
Q
qijun 已提交
294 295
      .def_static("create",
                  [](paddle::platform::CPUPlace& place)
Q
qijun 已提交
296
                      -> paddle::platform::DeviceContext* {
Q
qijun 已提交
297 298 299
                    return new paddle::platform::CPUDeviceContext();
                  })
      .def_static("create",
D
dzhwinter 已提交
300
                  [](paddle::platform::CUDAPlace& place)
Q
qijun 已提交
301
                      -> paddle::platform::DeviceContext* {
302
#ifndef PADDLE_WITH_CUDA
D
dzhwinter 已提交
303
                    PADDLE_THROW("CUDAPlace is not supported in CPU device.");
Q
qijun 已提交
304
#else
Q
qijun 已提交
305
                    return new paddle::platform::CUDADeviceContext(place);
Q
qijun 已提交
306
#endif
Q
qijun 已提交
307
                  });
D
Dong Zhihong 已提交
308
// clang-format on
Q
qijun 已提交
309

D
Dong Zhihong 已提交
310 311 312
#ifdef PADDLE_WITH_CUDA
  py::class_<platform::Communicator>(m, "Communicator").def(py::init<>());
#endif
D
dzhwinter 已提交
313
  py::class_<platform::CUDAPlace>(m, "CUDAPlace")
314
      .def(py::init<int>())
D
dzhwinter 已提交
315
      .def("__str__", string::to_string<const platform::CUDAPlace &>);
Q
qijun 已提交
316

317 318 319
  py::class_<paddle::platform::CPUPlace>(m, "CPUPlace")
      .def(py::init<>())
      .def("__str__", string::to_string<const platform::CPUPlace &>);
Y
Yu Yang 已提交
320

Y
Yu Yang 已提交
321 322 323 324 325 326 327
  py::class_<platform::Place>(m, "Place")
      .def(py::init<>())
      .def("set_place",
           [](platform::Place &self, const platform::CPUPlace &cpu_place) {
             self = cpu_place;
           })
      .def("set_place",
D
dzhwinter 已提交
328
           [](platform::Place &self, const platform::CUDAPlace &gpu_place) {
Y
Yu Yang 已提交
329 330 331
             self = gpu_place;
           });

Y
Yu Yang 已提交
332 333 334
  py::class_<OperatorBase>(m, "Operator")
      .def_static("create",
                  [](py::bytes protobin) {
335
                    proto::OpDesc desc;
Y
Yu Yang 已提交
336 337 338 339 340
                    PADDLE_ENFORCE(desc.ParsePartialFromString(protobin),
                                   "Cannot parse user input to OpDesc");
                    PADDLE_ENFORCE(desc.IsInitialized(),
                                   "User OpDesc is not initialized, reason %s",
                                   desc.InitializationErrorString());
341
                    return OpRegistry::CreateOp(desc);
Y
Yu Yang 已提交
342 343 344 345 346 347
                  })
      .def("backward",
           [](const OperatorBase &forwardOp,
              const std::unordered_set<std::string> &no_grad_vars) {
             return Backward(forwardOp, no_grad_vars).release();
           })
348
      .def("run",
349
           [](OperatorBase &self, const Scope &scope,
D
dzhwinter 已提交
350 351 352
              const platform::CPUPlace &place) { self.Run(scope, place); })
      .def("run",
           [](OperatorBase &self, const Scope &scope,
D
dzhwinter 已提交
353
              const platform::CUDAPlace &place) { self.Run(scope, place); })
Y
Yu Yang 已提交
354 355 356 357 358 359 360
      .def("type",
           [](const OperatorBase &op) -> std::string { return op.Type(); })
      .def("outputs",
           [](const OperatorBase &op)
               -> std::map<std::string, std::vector<std::string>> {
                 return op.Outputs();
               })
Q
qijun 已提交
361 362
      .def("output_vars",
           [](const OperatorBase &op) { return op.OutputVars(true); })
Y
Yu Yang 已提交
363
      .def("inputs", [](const OperatorBase &op) { return op.Inputs(); })
Q
qijun 已提交
364
      .def("input_vars", [](const OperatorBase &op) { return op.InputVars(); })
Y
Yu Yang 已提交
365 366 367 368
      .def("__str__", &OperatorBase::DebugString)
      .def("no_intermediate_outputs",
           [](const OperatorBase &op) { return op.OutputVars(false); })
      .def("support_gpu", &OperatorBase::SupportGPU);
Y
Yu Yang 已提交
369

Y
Yu Yang 已提交
370 371 372 373 374 375 376
  py::class_<operators::NetOp, OperatorBase>(m, "Net")
      .def_static("create",
                  []() -> operators::NetOp * {
                    auto *retv = new operators::NetOp;
                    retv->SetType("plain_net");
                    return retv;
                  })
377 378
      .def("append_op", [](operators::NetOp &self,
                           const OperatorBase &op) { self.AppendOp(op); })
D
dongzhihong 已提交
379 380 381 382
      .def("complete_add_op", &operators::NetOp::CompleteAddOp)
      .def("complete_add_op", [](std::shared_ptr<operators::NetOp> &self) {
        self->CompleteAddOp();
      });
Y
Yan Chunwei 已提交
383

Z
cond op  
zchen0211 已提交
384 385 386 387
  // cond_op
  py::class_<operators::CondOp, OperatorBase>(m, "CondOp")
      .def_static("create",
                  [](py::bytes protobin) -> operators::CondOp * {
388
                    proto::OpDesc desc;
Z
cond op  
zchen0211 已提交
389 390 391 392 393
                    PADDLE_ENFORCE(desc.ParsePartialFromString(protobin),
                                   "Cannot parse user input to OpDesc");
                    PADDLE_ENFORCE(desc.IsInitialized(),
                                   "User OpDesc is not initialized, reason %s",
                                   desc.InitializationErrorString());
394
                    auto cond_op = OpRegistry::CreateOp(desc);
Z
cond op  
zchen0211 已提交
395 396 397 398 399 400 401 402 403 404 405
                    return static_cast<operators::CondOp *>(cond_op.release());
                  })
      .def("set_truenet",
           [](operators::CondOp &self, const operators::NetOp &net) -> void {
             self.set_truenet(net.Clone());
           })
      .def("set_falsenet",
           [](operators::CondOp &self, const operators::NetOp &net) -> void {
             self.set_falsenet(net.Clone());
           });

F
fengjiayi 已提交
406
  py::class_<framework::Executor>(m, "Executor")
D
dzhwinter 已提交
407
      .def(py::init<const platform::Place &>())
408 409 410
      .def("run",
           (void (Executor::*)(const ProgramDesc &, Scope *, int, bool, bool)) &
               Executor::Run);
F
fengjiayi 已提交
411

412
  m.def("unique_integer", UniqueIntegerGenerator);
D
dzhwinter 已提交
413
  m.def("init_gflags", framework::InitGflags);
Y
Yang Yu 已提交
414
  m.def("init_glog", framework::InitGLOG);
D
dzhwinter 已提交
415
  m.def("init_devices", &framework::InitDevices);
416

417
  m.def("is_compiled_with_cuda", IsCompiledWithCUDA);
418

419
  m.def("set_feed_variable", framework::SetFeedVariable);
Q
qijun 已提交
420
  m.def("get_fetch_variable", framework::GetFetchVariable);
Q
qijun 已提交
421

F
fengjiayi 已提交
422 423 424 425
  BindProgramDesc(m);
  BindBlockDesc(m);
  BindVarDsec(m);
  BindOpDesc(m);
Q
qiaolongfei 已提交
426
  BindConstValue(m);
Y
Yu Yang 已提交
427

Y
Yu Yang 已提交
428 429 430 431 432 433 434 435 436
  py::class_<framework::LoDRankTable>(m, "LodRankTable")
      .def("items", [](framework::LoDRankTable &table) {
        std::vector<std::pair<size_t, size_t>> res;
        for (auto &item : table.items()) {
          res.push_back({item.index, item.length});
        }
        return res;
      });

Y
Yu Yang 已提交
437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453
  py::class_<LoDTensorArray>(m, "LoDTensorArray")
      .def("__getitem__",
           [](LoDTensorArray &self, size_t i) { return &self.at(i); },
           py::return_value_policy::reference)
      .def("__len__", [](LoDTensorArray &self) { return self.size(); })
      .def("__setitem__",
           [](LoDTensorArray &self, size_t i, const LoDTensor &t) {
             PADDLE_ENFORCE_LT(i, self.size());
             self[i].ShareDataWith(t);
             self[i].set_lod(t.lod());
           })
      .def("append", [](LoDTensorArray &self, const LoDTensor &t) {
        self.emplace_back();
        self.back().ShareDataWith(t);
        self.back().set_lod(t.lod());
      });

Y
Yu Yang 已提交
454
  m.def("op_support_gpu", OpSupportGPU);
D
Dong Zhihong 已提交
455
#ifdef PADDLE_WITH_CUDA
D
Dong Zhihong 已提交
456
  m.def("get_cuda_device_count", platform::GetCUDADeviceCount);
D
dangqingqing 已提交
457 458 459 460

  m.def("nvprof_init", platform::CudaProfilerInit);
  m.def("nvprof_start", platform::CudaProfilerStart);
  m.def("nvprof_stop", platform::CudaProfilerStop);
D
Dong Zhihong 已提交
461
#endif
Y
Yu Yang 已提交
462

463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480
  py::enum_<platform::ProfilerState>(m, "ProfilerState", py::arithmetic())
      .value("kDisabled", platform::ProfilerState::kDisabled)
      .value("kCPU", platform::ProfilerState::kCPU)
      .value("kCUDA", platform::ProfilerState::kCUDA)
      .export_values();

  py::enum_<platform::EventSortingKey>(m, "EventSortingKey", py::arithmetic())
      .value("kDefault", platform::EventSortingKey::kDefault)
      .value("kCalls", platform::EventSortingKey::kCalls)
      .value("kTotal", platform::EventSortingKey::kTotal)
      .value("kMin", platform::EventSortingKey::kMin)
      .value("kMax", platform::EventSortingKey::kMax)
      .value("kAve", platform::EventSortingKey::kAve)
      .export_values();

  m.def("enable_profiler", platform::EnableProfiler);
  m.def("disable_profiler", platform::DisableProfiler);
  m.def("reset_profiler", platform::ResetProfiler);
481
  return m.ptr();
L
Luo Tao 已提交
482
}
483
}  // namespace pybind
484
}  // namespace paddle