wehub-resource-sync
116 lines
3.8 KiB
C++
116 lines
3.8 KiB
C++
/* Copyright (c) 2021 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 <cstdint>
|
|
#include <functional>
|
|
#include <memory>
|
|
|
|
#include "paddle/phi/common/place.h"
|
|
|
|
namespace phi {
|
|
|
|
/// \brief Fancy pointer with deleter. The use of this data type
|
|
/// is to be compatible with allocators from different frameworks
|
|
/// without significant performance loss. This class does not
|
|
/// support being inherited.
|
|
class Allocation {
|
|
public:
|
|
using Place = phi::Place;
|
|
using DeleterFnPtr = void (*)(Allocation*);
|
|
|
|
Allocation() = default;
|
|
|
|
// Don't own resources, only provide access.
|
|
Allocation(void* data, size_t size, const Place& place)
|
|
: ptr_(data), size_(size), place_(place) {}
|
|
|
|
// Own resources.
|
|
Allocation(void* data, size_t size, DeleterFnPtr deleter, const Place& place)
|
|
: ptr_(data), size_(size), deleter_(deleter), place_(place) {}
|
|
|
|
Allocation(Allocation&& other) noexcept { swap(*this, other); }
|
|
Allocation& operator=(Allocation&& other) noexcept {
|
|
// Exchange them explicitly to avoid moving is equivalent
|
|
// to copying.
|
|
swap(*this, other);
|
|
return *this;
|
|
}
|
|
|
|
virtual ~Allocation() {
|
|
if (deleter_) {
|
|
deleter_(this);
|
|
}
|
|
}
|
|
|
|
// Returns the holding pointer.
|
|
// NOTE(liujinnan): Make this method as a virtual method. Because we add
|
|
// memory compact func for VirtualMemoryAutoGrowthBestFitAllocator. it will
|
|
// change ptr_ for defragmentation.
|
|
virtual void* ptr() const noexcept { return ptr_; }
|
|
void set_ptr(void* new_ptr) noexcept { ptr_ = new_ptr; }
|
|
|
|
uint64_t id() const noexcept { return id_; }
|
|
void set_id(uint64_t id) noexcept { id_ = id; }
|
|
|
|
// Returns the size of this memory buffer, i.e., ptr() + size() - 1 is the
|
|
// last valid element.
|
|
//
|
|
// NOTE: Some allocator might alloc more memory than request. The size
|
|
// could larger than its request. For example,
|
|
// the AlignedAllocator will always allocate memory as size + kAlignment.
|
|
// The raw pointer might not aligned, so an offset might be added to raw
|
|
// the pointer. The size of this allocation will be
|
|
// `size + kAlignment - offset`.
|
|
virtual size_t size() const noexcept { return size_; }
|
|
|
|
void* operator->() const noexcept { return ptr_; }
|
|
operator bool() const noexcept { return ptr_; }
|
|
virtual const Place& place() const noexcept { return place_; }
|
|
DeleterFnPtr deleter() const noexcept { return deleter_; }
|
|
|
|
protected:
|
|
friend void swap(Allocation& a, Allocation& b) noexcept;
|
|
void* ptr_{nullptr};
|
|
size_t size_{};
|
|
DeleterFnPtr deleter_{nullptr};
|
|
// TODO(Shixiaowei02): Enum needs to be used instead to reduce
|
|
// the construction overhead by more than 50%.
|
|
uint64_t id_{0}; // for debug
|
|
Place place_;
|
|
};
|
|
|
|
inline void swap(Allocation& a, Allocation& b) noexcept {
|
|
::std::swap(a.ptr_, b.ptr_);
|
|
::std::swap(a.deleter_, b.deleter_);
|
|
::std::swap(a.place_, b.place_);
|
|
::std::swap(a.size_, b.size_);
|
|
}
|
|
|
|
class Allocator {
|
|
public:
|
|
using DeleterType = std::function<void(Allocation*)>;
|
|
using AllocationPtr = std::unique_ptr<Allocation, DeleterType>;
|
|
|
|
virtual ~Allocator() = default;
|
|
virtual AllocationPtr Allocate(size_t bytes_size) = 0;
|
|
virtual void Free(phi::Allocation* allocation) {}
|
|
virtual void PreAlloc() {}
|
|
|
|
virtual bool IsAllocThreadSafe() const { return false; }
|
|
};
|
|
|
|
} // namespace phi
|