diff --git a/src/pinned_memory_manager.cc b/src/pinned_memory_manager.cc index 1152fb559..6c45feb5b 100644 --- a/src/pinned_memory_manager.cc +++ b/src/pinned_memory_manager.cc @@ -36,8 +36,7 @@ #include #endif // TRITON_ENABLE_GPU -namespace triton { -namespace core { +namespace triton { namespace core { namespace { @@ -264,8 +263,8 @@ PinnedMemoryManager::Create(const Options& options) instance_.reset(new PinnedMemoryManager()); if (options.host_policy_map_.empty()) { void* buffer = nullptr; - if (options.pinned_memory_pool_byte_size_ > 0) { #ifdef TRITON_ENABLE_GPU + if (options.pinned_memory_pool_byte_size_ > 0) { auto err = cudaHostAlloc( &buffer, options.pinned_memory_pool_byte_size_, cudaHostAllocPortable); @@ -325,8 +324,8 @@ PinnedMemoryManager::Create(const Options& options) continue; } void* buffer = nullptr; - if (options.pinned_memory_pool_byte_size_ > 0) { #ifdef TRITON_ENABLE_GPU + if (options.pinned_memory_pool_byte_size_ > 0) { auto err = cudaHostAlloc( &buffer, options.pinned_memory_pool_byte_size_, cudaHostAllocPortable); @@ -357,84 +356,89 @@ PinnedMemoryManager::Create(const Options& options) "Failed to add Pinned Memory buffer with host policy: " + std::string(ex.what())); } - } else { - LOG_INFO << "Pinned memory pool disabled"; } - // If no pinned memory is allocated, add an empty entry where all - // allocation will be on normal system memory - if (instance_->pinned_memory_buffers_.empty()) { - try { - instance_->AddPinnedMemoryBuffer( - std::shared_ptr(new PinnedMemory( - nullptr, options.pinned_memory_pool_byte_size_)), - 0); - } - catch (const std::exception& ex) { - return Status( - Status::Code::INTERNAL, - "Failed to add empty Pinned Memory entry: " + - std::string(ex.what())); - } + } + else + { + LOG_INFO << "Pinned memory pool disabled"; + } + // If no pinned memory is allocated, add an empty entry where all allocation + // will be on normal system memory + if (instance_->pinned_memory_buffers_.empty()) { + try { + instance_->AddPinnedMemoryBuffer( + std::shared_ptr(new PinnedMemory( + nullptr, options.pinned_memory_pool_byte_size_)), + 0); + } + catch (const std::exception& ex) { + return Status( + Status::Code::INTERNAL, + "Failed to add empty Pinned Memory entry: " + + std::string(ex.what())); } } - pinned_memory_byte_size_ = options.pinned_memory_pool_byte_size_; - return Status::Success; } + pinned_memory_byte_size_ = options.pinned_memory_pool_byte_size_; + return Status::Success; +} - Status PinnedMemoryManager::Alloc( - void** ptr, uint64_t size, TRITONSERVER_MemoryType* allocated_type, - bool allow_nonpinned_fallback) - { - if (instance_ == nullptr) { - return Status( - Status::Code::UNAVAILABLE, - "PinnedMemoryManager has not been created"); - } +Status +PinnedMemoryManager::Alloc( + void** ptr, uint64_t size, TRITONSERVER_MemoryType* allocated_type, + bool allow_nonpinned_fallback) +{ + if (instance_ == nullptr) { + return Status( + Status::Code::UNAVAILABLE, "PinnedMemoryManager has not been created"); + } - auto pinned_memory_buffer = - instance_->pinned_memory_buffers_.begin()->second.get(); - if (instance_->pinned_memory_buffers_.size() > 1) { - unsigned long node_mask; - if (GetNumaMemoryPolicyNodeMask(&node_mask).IsOk()) { - auto it = instance_->pinned_memory_buffers_.find(node_mask); - if (it != instance_->pinned_memory_buffers_.end()) { - pinned_memory_buffer = it->second.get(); - } + auto pinned_memory_buffer = + instance_->pinned_memory_buffers_.begin()->second.get(); + if (instance_->pinned_memory_buffers_.size() > 1) { + unsigned long node_mask; + if (GetNumaMemoryPolicyNodeMask(&node_mask).IsOk()) { + auto it = instance_->pinned_memory_buffers_.find(node_mask); + if (it != instance_->pinned_memory_buffers_.end()) { + pinned_memory_buffer = it->second.get(); } } - - return instance_->AllocInternal( - ptr, size, allocated_type, allow_nonpinned_fallback, - pinned_memory_buffer); } - Status PinnedMemoryManager::Free(void* ptr) - { - if (instance_ == nullptr) { - return Status( - Status::Code::UNAVAILABLE, - "PinnedMemoryManager has not been created"); - } + return instance_->AllocInternal( + ptr, size, allocated_type, allow_nonpinned_fallback, + pinned_memory_buffer); +} - return instance_->FreeInternal(ptr); +Status +PinnedMemoryManager::Free(void* ptr) +{ + if (instance_ == nullptr) { + return Status( + Status::Code::UNAVAILABLE, "PinnedMemoryManager has not been created"); } - uint64_t PinnedMemoryManager::GetTotalPinnedMemoryByteSize() - { - return pinned_memory_byte_size_; - } + return instance_->FreeInternal(ptr); +} - uint64_t PinnedMemoryManager::GetUsedPinnedMemoryByteSize() - { - std::lock_guard lk(allocated_buffer_mtx_); - uint64_t used_pinned_memory_size = 0; - if (!allocated_pinned_memory_buffers_.empty()) { - for (const auto& it : allocated_pinned_memory_buffers_) { - used_pinned_memory_size += it->GetUsedPinnedMemorySizeInternal(); - } - } +uint64_t +PinnedMemoryManager::GetTotalPinnedMemoryByteSize() +{ + return pinned_memory_byte_size_; +} - return used_pinned_memory_size; +uint64_t +PinnedMemoryManager::GetUsedPinnedMemoryByteSize() +{ + std::lock_guard lk(allocated_buffer_mtx_); + uint64_t used_pinned_memory_size = 0; + if (!allocated_pinned_memory_buffers_.empty()) { + for (const auto& it : allocated_pinned_memory_buffers_) { + used_pinned_memory_size += it->GetUsedPinnedMemorySizeInternal(); + } } + + return used_pinned_memory_size; } -} // namespace triton::core + +}} // namespace triton::core