1 /* Copyright 2018 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 #ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ 17 #define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ 18 19 #include <unordered_map> 20 21 #include "tensorflow/core/framework/allocator.h" 22 #include "tensorflow/core/platform/mutex.h" 23 24 #if GOOGLE_CUDA && GOOGLE_TENSORRT 25 #include "third_party/tensorrt/NvInfer.h" 26 #endif // GOOGLE_CUDA && GOOGLE_TENSORRT 27 28 namespace tensorflow { 29 namespace tensorrt { 30 // std::align is not supported, so this function mimic its behavior. 31 void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space); 32 } // namespace tensorrt 33 } // namespace tensorflow 34 35 #if GOOGLE_CUDA && GOOGLE_TENSORRT 36 37 namespace tensorflow { 38 namespace tensorrt { 39 40 class TRTBaseAllocator : public nvinfer1::IGpuAllocator { 41 // Base allocator class so we can have a virtual destructor; 42 public: 43 // python wrapper seems to be not happy with an pure virtual destructor; 44 virtual ~TRTBaseAllocator() = default; 45 }; 46 47 class TRTDeviceAllocator : public TRTBaseAllocator { 48 // Allocator implementation wrapping TF device allocators. 49 public: 50 TRTDeviceAllocator(Allocator* allocator); 51 52 // TODO(aaroey): base class doesn't have a virtual destructor, work with 53 // Nvidia to fix it. ~TRTDeviceAllocator()54 virtual ~TRTDeviceAllocator() { 55 VLOG(1) << "Destroying allocator attached to " << allocator_->Name(); 56 } 57 void* allocate(uint64_t size, uint64_t alignment, 58 uint32_t flags) noexcept override; 59 void free(void* memory) noexcept override; 60 61 private: 62 mutex mu_; 63 Allocator* allocator_; 64 65 // supporting alignment from allocation request requires a map to free; 66 std::unordered_map<void*, void*> mem_map_ TF_GUARDED_BY(mu_); 67 }; 68 69 } // namespace tensorrt 70 } // namespace tensorflow 71 72 #endif // GOOGLE_CUDA && GOOGLE_TENSORRT 73 #endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ 74