From 51227bd4473c8cad19fe026d397dde11025a4ac6 Mon Sep 17 00:00:00 2001 From: sneaxiy Date: Wed, 23 Jan 2019 11:56:24 +0000 Subject: [PATCH] lazy_allocator test=develop --- paddle/fluid/platform/gpu_info.cc | 19 ++++++++++++++++++- paddle/fluid/pybind/pybind.cc | 20 ++++++++++++++++++-- 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/platform/gpu_info.cc b/paddle/fluid/platform/gpu_info.cc index ca89d91aadb..a402407709a 100644 --- a/paddle/fluid/platform/gpu_info.cc +++ b/paddle/fluid/platform/gpu_info.cc @@ -15,6 +15,8 @@ limitations under the License. */ #include "paddle/fluid/platform/gpu_info.h" #include +#include +#include #include "gflags/gflags.h" #include "paddle/fluid/platform/enforce.h" @@ -58,7 +60,17 @@ DEFINE_string(selected_gpus, "", namespace paddle { namespace platform { -int GetCUDADeviceCount() { +static int GetCUDADeviceCountImpl() { + const auto *cuda_visible_devices = std::getenv("CUDA_VISIBLE_DEVICES"); + if (cuda_visible_devices != nullptr) { + std::string cuda_visible_devices_str(cuda_visible_devices); + if (std::all_of(cuda_visible_devices_str.begin(), + cuda_visible_devices_str.end(), + [](char ch) { return ch == ' '; })) { + return 0; + } + } + int count; PADDLE_ENFORCE( cudaGetDeviceCount(&count), @@ -66,6 +78,11 @@ int GetCUDADeviceCount() { return count; } +int GetCUDADeviceCount() { + static auto dev_cnt = GetCUDADeviceCountImpl(); + return dev_cnt; +} + int GetCUDAComputeCapability(int id) { PADDLE_ENFORCE_LT(id, GetCUDADeviceCount(), "id must less than GPU count"); cudaDeviceProp device_prop; diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index c4704837566..81546b33f5c 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -626,7 +626,18 @@ All parameter, weight, gradient are variables in Paddle. py::class_(m, "Communicator").def(py::init<>()); #endif py::class_(m, "CUDAPlace") - .def(py::init()) + .def("__init__", + [](platform::CUDAPlace &self, int dev_id) { +#ifdef PADDLE_WITH_CUDA + PADDLE_ENFORCE( + dev_id >= 0 && dev_id < platform::GetCUDADeviceCount(), + "Invalid CUDAPlace(%d), must inside [0, %d)", dev_id, + platform::GetCUDADeviceCount()); + new (&self) platform::CUDAPlace(dev_id); +#else + PADDLE_THROW("Cannot use CUDAPlace in CPU only version"); +#endif + }) .def("__str__", string::to_string); py::class_(m, "CPUPlace") @@ -634,7 +645,12 @@ All parameter, weight, gradient are variables in Paddle. .def("__str__", string::to_string); py::class_(m, "CUDAPinnedPlace") - .def(py::init<>()) + .def("__init__", + [](platform::CUDAPinnedPlace &) { +#ifndef PADDLE_WITH_CUDA + PADDLE_THROW("Cannot use CUDAPinnedPlace in CPU only version"); +#endif + }) .def("__str__", string::to_string); py::class_(m, "Place") -- GitLab