wehub-resource-sync
126 lines
4.0 KiB
C++
126 lines
4.0 KiB
C++
/* Copyright (c) 2016 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 <memory>
|
|
|
|
#include "paddle/phi/backends/device_manager.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/device_context.h"
|
|
#include "paddle/phi/core/memory/allocation/allocator.h"
|
|
#include "paddle/phi/core/stream.h"
|
|
|
|
#ifdef PADDLE_WITH_XPU
|
|
#include "xpu/runtime.h"
|
|
#include "xpu/runtime_ex.h"
|
|
#endif
|
|
|
|
namespace paddle {
|
|
namespace memory {
|
|
|
|
using allocation::AllocationPtr;
|
|
using allocation::Allocator;
|
|
using phi::Allocation;
|
|
|
|
PADDLE_API extern std::shared_ptr<Allocation> AllocShared(const Place& place,
|
|
size_t size);
|
|
|
|
PADDLE_API extern AllocationPtr Alloc(const Place& place, size_t size);
|
|
|
|
PADDLE_API extern uint64_t Release(const Place& place);
|
|
|
|
// Compact memory of free blocks held by the VmmAllocator.
|
|
PADDLE_API extern size_t Compact(const GPUPlace& place);
|
|
|
|
PADDLE_API extern std::shared_ptr<Allocation> AllocShared(
|
|
const Place& place, size_t size, const phi::Stream& stream);
|
|
|
|
PADDLE_API extern AllocationPtr Alloc(const Place& place,
|
|
size_t size,
|
|
const phi::Stream& stream);
|
|
|
|
PADDLE_API extern bool InSameStream(
|
|
const std::shared_ptr<Allocation>& allocation, const phi::Stream& stream);
|
|
|
|
PADDLE_API extern void* GetBasePtr(
|
|
const std::shared_ptr<Allocation>& allocation);
|
|
|
|
#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && \
|
|
!defined(PADDLE_WITH_CUSTOM_DEVICE)
|
|
PADDLE_API extern uint64_t Release(const GPUPlace& place, gpuStream_t stream);
|
|
|
|
PADDLE_API bool RecordStream(std::shared_ptr<Allocation> allocation,
|
|
gpuStream_t stream);
|
|
|
|
void EraseStream(std::shared_ptr<Allocation> allocation, gpuStream_t stream);
|
|
|
|
PADDLE_API gpuStream_t GetStream(const std::shared_ptr<Allocation>& allocation);
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_XPU
|
|
bool RecordStream(std::shared_ptr<Allocation> allocation, XPUStream stream);
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
extern uint64_t Release(const CustomPlace& place, phi::stream::stream_t stream);
|
|
|
|
bool RecordStream(std::shared_ptr<Allocation> allocation,
|
|
phi::stream::stream_t stream);
|
|
|
|
void EraseStream(std::shared_ptr<Allocation> allocation,
|
|
phi::stream::stream_t stream);
|
|
|
|
phi::stream::stream_t GetStream(const std::shared_ptr<Allocation>& allocation);
|
|
#endif
|
|
|
|
template <typename StreamType>
|
|
struct ThrustAllocator {
|
|
typedef char value_type;
|
|
ThrustAllocator(Place place, StreamType stream) {
|
|
place_ = place;
|
|
stream_ = stream;
|
|
}
|
|
~ThrustAllocator() {}
|
|
char* allocate(std::ptrdiff_t num_bytes) {
|
|
auto storage = memory::AllocShared(
|
|
place_,
|
|
num_bytes,
|
|
phi::Stream(reinterpret_cast<phi::StreamId>(stream_)));
|
|
char* ptr = reinterpret_cast<char*>(storage->ptr());
|
|
busy_allocation_.emplace(std::make_pair(ptr, storage));
|
|
return ptr;
|
|
}
|
|
void deallocate(char* ptr, size_t) {
|
|
allocation_map_type::iterator iter = busy_allocation_.find(ptr);
|
|
PADDLE_ENFORCE_NE(iter,
|
|
busy_allocation_.end(),
|
|
common::errors::InvalidArgument(
|
|
"Attempting to deallocate a pointer not "
|
|
"found in busy_allocation_ map."));
|
|
|
|
busy_allocation_.erase(iter);
|
|
}
|
|
|
|
private:
|
|
typedef std::unordered_map<char*, std::shared_ptr<phi::Allocation>>
|
|
allocation_map_type;
|
|
allocation_map_type busy_allocation_;
|
|
Place place_;
|
|
StreamType stream_;
|
|
};
|
|
|
|
} // namespace memory
|
|
} // namespace paddle
|