wehub-resource-sync
140 lines
5.0 KiB
C++
140 lines
5.0 KiB
C++
// Copyright (c) 2026 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 <c10/macros/Macros.h>
|
|
|
|
#include <cstddef>
|
|
#include <memory>
|
|
#include <utility>
|
|
|
|
namespace c10 {
|
|
|
|
using DeleterFnPtr = void (*)(void*);
|
|
|
|
namespace detail {
|
|
|
|
// Does not delete anything
|
|
inline void deleteNothing(void* /*unused*/) {}
|
|
|
|
// A detail::UniqueVoidPtr is an owning smart pointer like unique_ptr, but
|
|
// with three major differences:
|
|
//
|
|
// 1) It is specialized to void
|
|
//
|
|
// 2) It is specialized for a function pointer deleter
|
|
// void(void* ctx); i.e., the deleter doesn't take a
|
|
// reference to the data, just to a context pointer
|
|
// (erased as void*). In fact, internally, this pointer
|
|
// is implemented as having an owning reference to
|
|
// context, and a non-owning reference to data; this is why
|
|
// you release_context(), not release() (the conventional
|
|
// API for release() wouldn't give you enough information
|
|
// to properly dispose of the object later.)
|
|
//
|
|
// 3) The deleter is guaranteed to be called when the unique
|
|
// pointer is destructed and the context is non-null; this is different
|
|
// from std::unique_ptr where the deleter is not called if the
|
|
// data pointer is null.
|
|
//
|
|
// Some of the methods have slightly different types than std::unique_ptr
|
|
// to reflect this.
|
|
//
|
|
class UniqueVoidPtr {
|
|
private:
|
|
// Lifetime tied to ctx_
|
|
void* data_;
|
|
std::unique_ptr<void, DeleterFnPtr> ctx_;
|
|
|
|
public:
|
|
UniqueVoidPtr() : data_(nullptr), ctx_(nullptr, &deleteNothing) {}
|
|
explicit UniqueVoidPtr(void* data)
|
|
: data_(data), ctx_(nullptr, &deleteNothing) {}
|
|
UniqueVoidPtr(void* data, void* ctx, DeleterFnPtr ctx_deleter)
|
|
: data_(data), ctx_(ctx, ctx_deleter ? ctx_deleter : &deleteNothing) {}
|
|
void* operator->() const { return data_; }
|
|
void clear() {
|
|
ctx_ = nullptr;
|
|
data_ = nullptr;
|
|
}
|
|
void* get() const { return data_; }
|
|
|
|
bool /* success */ unsafe_reset_data_and_ctx(void* new_data_and_ctx) {
|
|
if (__builtin_expect(
|
|
static_cast<bool>((ctx_.get_deleter() != &deleteNothing)), 0)) {
|
|
return false;
|
|
}
|
|
// seems quicker than calling the no-op deleter when we reset
|
|
(void)ctx_.release();
|
|
ctx_.reset(new_data_and_ctx);
|
|
data_ = new_data_and_ctx;
|
|
return true;
|
|
}
|
|
|
|
void* get_context() const { return ctx_.get(); }
|
|
void* release_context() { return ctx_.release(); }
|
|
std::unique_ptr<void, DeleterFnPtr>&& move_context() {
|
|
return std::move(ctx_);
|
|
}
|
|
[[nodiscard]] bool compare_exchange_deleter(DeleterFnPtr expected_deleter,
|
|
DeleterFnPtr new_deleter) {
|
|
if (get_deleter() != expected_deleter) return false;
|
|
ctx_ = std::unique_ptr<void, DeleterFnPtr>(ctx_.release(), new_deleter);
|
|
return true;
|
|
}
|
|
|
|
template <typename T>
|
|
T* cast_context(DeleterFnPtr expected_deleter) const {
|
|
if (get_deleter() != expected_deleter) return nullptr;
|
|
return static_cast<T*>(get_context());
|
|
}
|
|
operator bool() const { return data_ || ctx_; }
|
|
DeleterFnPtr get_deleter() const { return ctx_.get_deleter(); }
|
|
};
|
|
|
|
// Note [How UniqueVoidPtr is implemented]
|
|
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
// UniqueVoidPtr solves a common problem for allocators of tensor data, which
|
|
// is that the data pointer (e.g., float*) which you are interested in, is not
|
|
// the same as the context pointer (e.g., DLManagedTensor) which you need
|
|
// to actually deallocate the data. Under a conventional deleter design, you
|
|
// have to store extra context in the deleter itself so that you can actually
|
|
// delete the right thing. Implementing this with standard C++ is somewhat
|
|
// error-prone: if you use a std::unique_ptr to manage tensors, the deleter will
|
|
// not be called if the data pointer is nullptr, which can cause a leak if the
|
|
// context pointer is non-null (and the deleter is responsible for freeing both
|
|
// the data pointer and the context pointer).
|
|
//
|
|
// So, in our reimplementation of unique_ptr, which just store the context
|
|
// directly in the unique pointer, and attach the deleter to the context
|
|
// pointer itself. In simple cases, the context pointer is just the pointer
|
|
// itself.
|
|
|
|
inline bool operator==(const UniqueVoidPtr& sp, std::nullptr_t) noexcept {
|
|
return !sp;
|
|
}
|
|
inline bool operator==(std::nullptr_t, const UniqueVoidPtr& sp) noexcept {
|
|
return !sp;
|
|
}
|
|
inline bool operator!=(const UniqueVoidPtr& sp, std::nullptr_t) noexcept {
|
|
return sp;
|
|
}
|
|
inline bool operator!=(std::nullptr_t, const UniqueVoidPtr& sp) noexcept {
|
|
return sp;
|
|
}
|
|
|
|
} // namespace detail
|
|
} // namespace c10
|