diff options
author | 2017-09-25 06:04:56 -0700 | |
---|---|---|
committer | 2017-09-25 06:08:20 -0700 | |
commit | cd4f5840a834e5380536b6f04a768408c6eebf3d (patch) | |
tree | 631c57f8f7faf63f3c934ef3bdcbfd09cebafb3f /tensorflow/core/common_runtime | |
parent | c4873e25bdec79f9c96e479a1bcb51ca63c2badf (diff) |
Minimal support for running OpsTest on GPU, using CUDA unified memory.
PiperOrigin-RevId: 169897567
Diffstat (limited to 'tensorflow/core/common_runtime')
-rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_managed_allocator.cc | 42 | ||||
-rw-r--r-- | tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h | 36 |
2 files changed, 78 insertions, 0 deletions
diff --git a/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.cc b/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.cc new file mode 100644 index 0000000000..41f7249d89 --- /dev/null +++ b/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.cc @@ -0,0 +1,42 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +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. +==============================================================================*/ + +#include "tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h" + +#ifdef GOOGLE_CUDA +#include "cuda/include/cuda.h" +#endif // GOOGLE_CUDA + +namespace tensorflow { + +void* GpuManagedAllocator::AllocateRaw(size_t alignment, size_t num_bytes) { +#ifdef GOOGLE_CUDA + CUdeviceptr ptr = 0; + CHECK_EQ(cuMemAllocManaged(&ptr, num_bytes, CU_MEM_ATTACH_GLOBAL), + CUDA_SUCCESS); + CHECK(!(ptr & (alignment - 1))); + return reinterpret_cast<void*>(ptr); +#else + return nullptr; +#endif +} + +void GpuManagedAllocator::DeallocateRaw(void* ptr) { +#ifdef GOOGLE_CUDA + CHECK_EQ(cuMemFree(reinterpret_cast<CUdeviceptr>(ptr)), CUDA_SUCCESS); +#endif +} + +} // namespace tensorflow diff --git a/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h b/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h new file mode 100644 index 0000000000..006b2ca448 --- /dev/null +++ b/tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h @@ -0,0 +1,36 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +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. +==============================================================================*/ + +#ifndef THIRD_PARTY_TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_MANAGED_ALLOCATOR_H_ +#define THIRD_PARTY_TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_MANAGED_ALLOCATOR_H_ + +#include "tensorflow/core/framework/allocator.h" + +namespace tensorflow { + +// An allocator for CUDA unified memory. Memory allocated with this allocator +// can be accessed from both host and device. CUDA transparently migrates dirty +// pages, which can be slow. Therefore, this allocator is intended for +// convenience in functional tests only. +class GpuManagedAllocator : public Allocator { + public: + string Name() override { return "GpuManagedAllocator"; } + void* AllocateRaw(size_t alignment, size_t num_bytes) override; + void DeallocateRaw(void* ptr) override; +}; + +} // namespace tensorflow + +#endif // THIRD_PARTY_TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_MANAGED_ALLOCATOR_H_ |