You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by lx...@apache.org on 2017/07/07 15:58:36 UTC
[23/50] [abbrv] incubator-mxnet-test git commit: fix pinned mem
USE_CUDA=1 on host with no device (#6864)
fix pinned mem USE_CUDA=1 on host with no device (#6864)
* fix pinned mem USE_CUDA=1 on host with no device
* fix
* fix
* fix
Project: http://git-wip-us.apache.org/repos/asf/incubator-mxnet-test/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-mxnet-test/commit/a94d0e24
Tree: http://git-wip-us.apache.org/repos/asf/incubator-mxnet-test/tree/a94d0e24
Diff: http://git-wip-us.apache.org/repos/asf/incubator-mxnet-test/diff/a94d0e24
Branch: refs/heads/master
Commit: a94d0e2448f767f7deed01168e61b096a0b3598e
Parents: a4fe4b4
Author: Sheng Zha <sz...@users.noreply.github.com>
Authored: Thu Jun 29 22:06:13 2017 -0700
Committer: Eric Junyuan Xie <pi...@users.noreply.github.com>
Committed: Thu Jun 29 22:06:13 2017 -0700
----------------------------------------------------------------------
src/storage/storage.cc | 27 +++++++++++++++++++--------
1 file changed, 19 insertions(+), 8 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-mxnet-test/blob/a94d0e24/src/storage/storage.cc
----------------------------------------------------------------------
diff --git a/src/storage/storage.cc b/src/storage/storage.cc
index 997f033..a7cdbf6 100644
--- a/src/storage/storage.cc
+++ b/src/storage/storage.cc
@@ -27,21 +27,25 @@ class StorageImpl : public Storage {
private:
static constexpr size_t kMaxNumberOfDevices = Context::kMaxDevType + 1;
static constexpr size_t kMaxNumberOfDeviceIDs = Context::kMaxDevID + 1;
+#if MXNET_USE_CUDA
+ static int num_gpu_device;
+#endif // MXNET_USE_CUDA
static void ActivateDevice(Context ctx) {
switch (ctx.dev_type) {
case Context::kCPU: break;
case Context::kGPU:
case Context::kCPUPinned: {
- int gpu_num = 0;
#if MXNET_USE_CUDA
- CUDA_CALL(cudaGetDeviceCount(&gpu_num));
-#endif // MXNET_USE_CUDA
- if (gpu_num > 0) {
-#if MXNET_USE_CUDA
- CUDA_CALL(cudaSetDevice(ctx.dev_id));
-#endif // MXNET_USE_CUDA
+ num_gpu_device = 0;
+ cudaError_t e = cudaGetDeviceCount(&num_gpu_device);
+ if (e != cudaSuccess) {
+ num_gpu_device = 0;
}
+ if (num_gpu_device > 0) {
+ CUDA_CALL(cudaSetDevice(ctx.dev_id));
+ }
+#endif // MXNET_USE_CUDA
break;
}
default:
@@ -52,6 +56,9 @@ class StorageImpl : public Storage {
std::array<common::LazyAllocArray<storage::StorageManager>,
kMaxNumberOfDevices> storage_managers_;
}; // struct Storage::Impl
+#if MXNET_USE_CUDA
+int StorageImpl::num_gpu_device = 0;
+#endif // MXNET_USE_CUDA
Storage::Handle StorageImpl::Alloc(size_t size, Context ctx) {
// space already recycled, ignore request
@@ -69,7 +76,11 @@ Storage::Handle StorageImpl::Alloc(size_t size, Context ctx) {
}
case Context::kCPUPinned: {
#if MXNET_USE_CUDA
- ptr = new storage::NaiveStorageManager<storage::PinnedMemoryStorage>();
+ if (num_gpu_device > 0) {
+ ptr = new storage::NaiveStorageManager<storage::PinnedMemoryStorage>();
+ } else {
+ ptr = new storage::NaiveStorageManager<storage::CPUDeviceStorage>();
+ }
#else
ptr = new storage::NaiveStorageManager<storage::CPUDeviceStorage>();
#endif // MXNET_USE_CUDA