wehub-resource-sync
123 lines
3.3 KiB
C++
123 lines
3.3 KiB
C++
/* Copyright (c) 2022 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. */
|
|
|
|
#pragma once
|
|
|
|
#include <mutex>
|
|
|
|
#include "paddle/common/macros.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/utils/flat_hash_map.h"
|
|
#if defined(PADDLE_WITH_XPU)
|
|
#include "paddle/phi/core/xpu_cuda_stream.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
class DeviceContext;
|
|
class CPUContext;
|
|
class GPUContext;
|
|
class CustomContext;
|
|
class Allocator;
|
|
class CUDAStream;
|
|
} // namespace phi
|
|
|
|
namespace paddle {
|
|
namespace experimental {
|
|
|
|
template <AllocationType T>
|
|
struct DefaultDeviceContextType;
|
|
|
|
template <>
|
|
struct DefaultDeviceContextType<AllocationType::CPU> {
|
|
using TYPE = phi::CPUContext;
|
|
};
|
|
|
|
template <>
|
|
struct DefaultDeviceContextType<AllocationType::GPU> {
|
|
using TYPE = phi::GPUContext;
|
|
};
|
|
|
|
template <>
|
|
struct DefaultDeviceContextType<AllocationType::CUSTOM> {
|
|
using TYPE = phi::CustomContext;
|
|
};
|
|
|
|
/**
|
|
* The DeviceContextPool here is just a mirror of the DeviceContextPool in
|
|
* fluid, and does not manage the life cycle of the DeviceContext.
|
|
* It is mainly used for external custom operator calls and high-performance
|
|
* C++ APIs.
|
|
*
|
|
* Since DeviceContextPool in fluid is a global singleton, it always exists
|
|
* in program running, so DeviceContextPool here can always access the correct
|
|
* DeviceContext pointer.
|
|
*
|
|
* In order not to depend on the fluid's DeviceContextPool,
|
|
* the DeviceContextPool here needs to be initialized in the fluid, and cannot
|
|
* be initialized by itself.
|
|
*
|
|
* Note: DeviceContextPool is an experimental API and may be removed in the
|
|
* future. From 2.3, we recommend directly using the C++ API to combine new
|
|
* operators.
|
|
*/
|
|
class PADDLE_API DeviceContextPool {
|
|
public:
|
|
static DeviceContextPool& Instance();
|
|
|
|
const phi::DeviceContext* Get(const Place& place);
|
|
|
|
phi::DeviceContext* GetMutable(const Place& place);
|
|
|
|
void SyncDeviceContext(const Place& place);
|
|
|
|
template <AllocationType T>
|
|
const typename DefaultDeviceContextType<T>::TYPE* Get(const Place& place) {
|
|
return reinterpret_cast<const typename DefaultDeviceContextType<T>::TYPE*>(
|
|
Get(place));
|
|
}
|
|
|
|
private:
|
|
DeviceContextPool() = default;
|
|
|
|
paddle::flat_hash_map<Place, const phi::DeviceContext*, Place::Hash>
|
|
context_map_;
|
|
std::mutex mutex_;
|
|
|
|
DISABLE_COPY_AND_ASSIGN(DeviceContextPool);
|
|
};
|
|
|
|
} // namespace experimental
|
|
} // namespace paddle
|
|
|
|
namespace paddle {
|
|
|
|
/**
|
|
* Get the Allocator for the passed place.
|
|
*/
|
|
PADDLE_API phi::Allocator* GetAllocator(const phi::Place& place);
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
/**
|
|
* Get the current CUDA stream for the passed CUDA device.
|
|
*/
|
|
PADDLE_API phi::CUDAStream* GetCurrentCUDAStream(const phi::Place& place);
|
|
#elif defined(PADDLE_WITH_XPU)
|
|
/**
|
|
* Get the current XPUCUDA stream for the passed XPU device.
|
|
*/
|
|
PADDLE_API cudaStream_t GetCurrentCUDAStream(const phi::Place& place);
|
|
#endif
|
|
|
|
} // namespace paddle
|