1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #ifdef GOOGLE_CUDA
17 #include "third_party/gpus/cuda/include/cuda.h"
18 #include "tensorflow/stream_executor/cuda/cuda_activation.h"
19 #endif // GOOGLE_CUDA
20
21 #include "tensorflow/core/common_runtime/device/device_id_utils.h"
22 #include "tensorflow/core/common_runtime/gpu/gpu_cudamalloc_allocator.h"
23 #include "tensorflow/core/common_runtime/gpu/gpu_id.h"
24 #include "tensorflow/core/common_runtime/gpu/gpu_init.h"
25 #include "tensorflow/core/platform/stream_executor.h"
26
27 namespace tensorflow {
28
GPUcudaMallocAllocator(PlatformDeviceId platform_device_id)29 GPUcudaMallocAllocator::GPUcudaMallocAllocator(
30 PlatformDeviceId platform_device_id) {
31 stream_exec_ = DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
32 platform_device_id)
33 .ValueOrDie();
34 }
35
AllocateRaw(size_t alignment,size_t num_bytes)36 void* GPUcudaMallocAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
37 #ifdef GOOGLE_CUDA
38 // allocate with cudaMalloc
39 se::cuda::ScopedActivateExecutorContext scoped_activation{stream_exec_};
40 CUdeviceptr rv = 0;
41 CUresult res = cuMemAlloc(&rv, num_bytes);
42 if (res != CUDA_SUCCESS) {
43 const char* error_name;
44 const char* error_string;
45 cuGetErrorName(res, &error_name);
46 cuGetErrorString(res, &error_string);
47 LOG(ERROR) << "cuMemAlloc failed to allocate " << num_bytes
48 << "\n Error name: " << error_name
49 << "\n Error string: " << error_string;
50 return nullptr;
51 }
52 VLOG(10) << "AllocateRaw " << Name() << " " << num_bytes << " "
53 << reinterpret_cast<void*>(rv);
54 return reinterpret_cast<void*>(rv);
55 #else
56 return nullptr;
57 #endif // GOOGLE_CUDA
58 }
DeallocateRaw(void * ptr)59 void GPUcudaMallocAllocator::DeallocateRaw(void* ptr) {
60 #ifdef GOOGLE_CUDA
61 // free with cudaFree
62 CUresult res = cuMemFree(reinterpret_cast<CUdeviceptr>(ptr));
63 if (res == CUDA_ERROR_DEINITIALIZED) {
64 // It happens with multi-GPU that TF free the GPU allocation after
65 // the driver is unloaded. It is safe to ignore this error here.
66 // cuGetErrorName and cuGetErrorString doesn't return any useful
67 // information here.
68 // TODO: Find how to fix the shutdown steps in TF.
69 VLOG(1) << "Ignoring CUDA_ERROR_DEINITIALIZED Error";
70 } else if (res != CUDA_SUCCESS) {
71 const char* error_name;
72 const char* error_string;
73 cuGetErrorName(res, &error_name);
74 cuGetErrorString(res, &error_string);
75 LOG(ERROR) << "cuMemFree failed to free " << ptr
76 << "\n Error name: " << error_name
77 << "\n Error string: " << error_string;
78 }
79 VLOG(10) << Name() << " Freed ptr: " << ptr;
80 #endif // GOOGLE_CUDA
81 }
82
TracksAllocationSizes() const83 bool GPUcudaMallocAllocator::TracksAllocationSizes() const { return false; }
84
85 } // namespace tensorflow
86