diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index b0a3f06a8871b1dc8c6c9d7231dfe2c9764ade3f..e7fa450832d9624a3b279a061c7ef83cb5fbec63 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -11,11 +11,16 @@ 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. */ +#include +#include +#include // NOLINT // for call_once +#include +#include +#include +#include #include "paddle/fluid/pybind/protobuf.h" -#include // for call_once -#include #include "paddle/fluid/framework/backward.h" #include "paddle/fluid/framework/channel.h" #include "paddle/fluid/framework/executor.h" @@ -32,7 +37,6 @@ limitations under the License. */ #include "paddle/fluid/operators/cond_op.h" #include "paddle/fluid/operators/net_op.h" #include "paddle/fluid/platform/enforce.h" -#include "paddle/fluid/platform/gpu_info.h" #include "paddle/fluid/platform/place.h" #include "paddle/fluid/platform/profiler.h" #include "paddle/fluid/pybind/const_value.h" @@ -100,6 +104,14 @@ PYBIND11_PLUGIN(core) { [](Tensor &self, paddle::platform::CUDAPlace &place) { self.mutable_data(place); }) + .def("alloc_int", + [](Tensor &self, paddle::platform::CUDAPinnedPlace &place) { + self.mutable_data(place); + }) + .def("alloc_float", + [](Tensor &self, paddle::platform::CUDAPinnedPlace &place) { + self.mutable_data(place); + }) .def("set", PyCPUTensorSetFromArray) .def("set", PyCPUTensorSetFromArray) .def("set", PyCPUTensorSetFromArray) @@ -317,7 +329,17 @@ All parameter, weight, gradient are variables in Paddle. #else return new paddle::platform::CUDADeviceContext(place); #endif - }); + }) + .def_static("create", + [](paddle::platform::CUDAPinnedPlace& place) + -> paddle::platform::DeviceContext* { +#ifndef PADDLE_WITH_CUDA + PADDLE_THROW( + "CUDAPinnedPlace is not supported in CPU device."); +#else + return new paddle::platform::CUDAPinnedDeviceContext(place); +#endif + });; // clang-format on #ifdef PADDLE_WITH_CUDA py::class_(m, "Communicator").def(py::init<>()); @@ -330,6 +352,10 @@ All parameter, weight, gradient are variables in Paddle. .def(py::init<>()) .def("__str__", string::to_string); + py::class_(m, "CUDAPinnedPlace") + .def(py::init<>()) + .def("__str__", string::to_string); + py::class_(m, "Place") .def(py::init<>()) .def("set_place", @@ -339,7 +365,11 @@ All parameter, weight, gradient are variables in Paddle. .def("set_place", [](platform::Place &self, const platform::CUDAPlace &gpu_place) { self = gpu_place; - }); + }) + .def("set_place", [](platform::Place &self, + const platform::CUDAPinnedPlace &gpu_place) { + self = gpu_place; + }); py::class_(m, "Operator") .def_static("create", @@ -363,6 +393,11 @@ All parameter, weight, gradient are variables in Paddle. .def("run", [](OperatorBase &self, const Scope &scope, const platform::CUDAPlace &place) { self.Run(scope, place); }) + .def("run", + [](OperatorBase &self, const Scope &scope, + const platform::CUDAPinnedPlace &place) { + self.Run(scope, place); + }) .def("type", [](const OperatorBase &op) -> std::string { return op.Type(); }) .def("outputs",