small_buffer_allocator.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <dispenso/detail/math.h>
#include <dispenso/platform.h>
 
namespace dispenso {
 
constexpr size_t kMaxSmallBufferSize = 256;
 
namespace detail {
 
DISPENSO_DLL_ACCESS char* allocSmallBufferImpl(size_t ordinal);
DISPENSO_DLL_ACCESS void deallocSmallBufferImpl(size_t ordinal, void* buf);
 
DISPENSO_DLL_ACCESS size_t approxBytesAllocatedSmallBufferImpl(size_t ordinal);
 
// This has the effect of selecting actual block sizes starting with 4 bytes.  Smaller requests
// (e.g. 1 byte, 2 bytes) will still utilize 4-byte blocks.  Choice of 4 bytes as the smallest
// mainly aligns to sizeof(ptr) on 32-bit platforms, where we'd expect most common use cases to be
// no smaller than one pointer.  Retaining 4-byte buckets on 64-bit platforms doesn't cost much
// (tiny startup/teardown cost, and trivial amount of memory) when not using 4-byte or smaller
// allocations, and makes the code simpler.
constexpr size_t getOrdinal(size_t blockSize) {
  return std::max<ssize_t>(0, log2const(blockSize) - 2);
}
 
template <size_t kBlockSize>
inline std::enable_if_t<(kBlockSize <= kMaxSmallBufferSize), char*> allocSmallOrLarge() {
#if defined(DISPENSO_NO_SMALL_BUFFER_ALLOCATOR)
  return reinterpret_cast<char*>(alignedMalloc(kBlockSize, kBlockSize));
#else
  return allocSmallBufferImpl(getOrdinal(kBlockSize));
#endif // DISPENSO_NO_SMALL_BUFFER_ALLOCATOR
}
 
template <size_t kBlockSize>
inline std::enable_if_t<(kBlockSize > kMaxSmallBufferSize), char*> allocSmallOrLarge() {
  return reinterpret_cast<char*>(alignedMalloc(kBlockSize, kBlockSize));
}
 
template <size_t kBlockSize>
inline std::enable_if_t<(kBlockSize <= kMaxSmallBufferSize), void> deallocSmallOrLarge(void* buf) {
#if defined(DISPENSO_NO_SMALL_BUFFER_ALLOCATOR)
  alignedFree(buf);
#else
  deallocSmallBufferImpl(getOrdinal(kBlockSize), buf);
#endif // DISPENSO_NO_SMALL_BUFFER_ALLOCATOR
}
 
template <size_t kBlockSize>
inline std::enable_if_t<(kBlockSize > kMaxSmallBufferSize), void> deallocSmallOrLarge(void* buf) {
  alignedFree(buf);
}
 
} // namespace detail
 
template <size_t kBlockSize>
inline char* allocSmallBuffer() {
  return detail::allocSmallOrLarge<kBlockSize>();
}
template <size_t kBlockSize>
inline void deallocSmallBuffer(void* buf) {
  detail::deallocSmallOrLarge<kBlockSize>(buf);
}
 
template <size_t kBlockSize>
size_t approxBytesAllocatedSmallBuffer() {
  return detail::approxBytesAllocatedSmallBufferImpl(detail::getOrdinal(kBlockSize));
}
 
} // namespace dispenso