wehub-resource-sync
124 lines
3.8 KiB
C++
124 lines
3.8 KiB
C++
/* Copyright (c) 2023 PaddlePaddle 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 "paddle/phi/backends/context_pool.h"
|
|
|
|
#include "glog/logging.h"
|
|
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
|
|
COMMON_DECLARE_bool(use_default_stream);
|
|
|
|
namespace phi {
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
|
|
defined(PADDLE_WITH_CUSTOM_DEVICE)
|
|
#if defined(PADDLE_WITH_CUSTOM_DEVICE)
|
|
bool allow_tf32_cublas = true;
|
|
#else
|
|
bool allow_tf32_cublas = false;
|
|
#endif
|
|
void SetAllowTF32Cublas(bool active) { allow_tf32_cublas = active; }
|
|
bool AllowTF32Cublas() { return allow_tf32_cublas; }
|
|
bool allow_tf32_cudnn = true;
|
|
void SetAllowTF32Cudnn(bool active) { allow_tf32_cudnn = active; }
|
|
bool AllowTF32Cudnn() { return allow_tf32_cudnn; }
|
|
#endif // PADDLE_WITH_CUDA
|
|
|
|
static DeviceContextPool* pool = nullptr;
|
|
|
|
TEST_API DeviceContextPool& DeviceContextPool::Instance() {
|
|
PADDLE_ENFORCE_NOT_NULL(pool,
|
|
common::errors::PreconditionNotMet(
|
|
"Need to Create DeviceContextPool firstly!"));
|
|
return *pool;
|
|
}
|
|
|
|
/*! \brief Create should only called by Init function */
|
|
TEST_API DeviceContextPool& DeviceContextPool::Init(
|
|
const std::vector<Place>& places) {
|
|
if (pool == nullptr) {
|
|
pool = new DeviceContextPool(places);
|
|
}
|
|
return *pool;
|
|
}
|
|
|
|
TEST_API bool DeviceContextPool::IsInitialized() { return pool != nullptr; }
|
|
|
|
TEST_API void DeviceContextPool::SetPool(DeviceContextPool* dev_pool) {
|
|
pool = dev_pool;
|
|
}
|
|
|
|
thread_local const std::map<Place,
|
|
std::shared_future<std::unique_ptr<DeviceContext>>>*
|
|
DeviceContextPool::external_device_contexts_ = nullptr;
|
|
|
|
TEST_API DeviceContext* DeviceContextPool::Get(const Place& place) {
|
|
VLOG(6) << "DeviceContextPool Get: " << place;
|
|
const std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
|
|
ptr = nullptr;
|
|
if (external_device_contexts_ && external_device_contexts_->count(place)) {
|
|
ptr = external_device_contexts_;
|
|
} else {
|
|
ptr = &device_contexts_;
|
|
}
|
|
|
|
auto it = ptr->find(place);
|
|
if (it == ptr->end()) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Place %s is not supported. Please check that your paddle compiles "
|
|
"with WITH_GPU, WITH_XPU or WITH_IPU option "
|
|
"or check "
|
|
"that your train process set the correct device id if you use "
|
|
"Executor.",
|
|
place));
|
|
}
|
|
return it->second.get().get();
|
|
}
|
|
|
|
TEST_API size_t DeviceContextPool::Size() const {
|
|
if (external_device_contexts_) {
|
|
return external_device_contexts_->size();
|
|
}
|
|
return device_contexts_.size();
|
|
}
|
|
|
|
TEST_API const
|
|
std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>&
|
|
DeviceContextPool::device_contexts() const {
|
|
if (external_device_contexts_) {
|
|
return *external_device_contexts_;
|
|
}
|
|
return device_contexts_;
|
|
}
|
|
|
|
TEST_API void DeviceContextPool::SetDeviceContexts(
|
|
const std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
|
|
dev_ctxs) {
|
|
external_device_contexts_ = dev_ctxs;
|
|
}
|
|
|
|
DeviceContextPool::DeviceContextPool(const std::vector<Place>& places) {
|
|
memory_utils::EmplaceDeviceContexts(
|
|
&device_contexts_,
|
|
places,
|
|
/*disable_setting_default_stream_for_allocator=*/false,
|
|
/*stream_priority=*/0,
|
|
FLAGS_use_default_stream);
|
|
}
|
|
|
|
} // namespace phi
|