allocators.hpp

Go to the documentation of this file.

// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA)
// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS)
// Copyright (c) 2018-2024 Simons Foundation
//
// 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.txt
//
// 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.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./address_space.hpp"
#include "./malloc.hpp"
#include "./memset.hpp"
#include "../macros.hpp"
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <vector>
#include <utility>
 
#ifndef NDEBUG
#include <iostream>
#endif
 
#if defined(__has_feature)
#if __has_feature(address_sanitizer)
#include <sanitizer/asan_interface.h>
#define NDA_USE_ASAN
#endif
#endif
 
namespace nda::mem {
 
  struct blk_t {
    char *ptr = nullptr;
 
    size_t s = 0;
  };
 
  template <AddressSpace AdrSp = Host>
  class mallocator {
    public:
    mallocator() = default;
 
    mallocator(mallocator const &) = delete;
 
    mallocator(mallocator &&) = default;
 
    mallocator &operator=(mallocator const &) = delete;
 
    mallocator &operator=(mallocator &&) = default;
 
    static constexpr auto address_space = AdrSp;
 
    static blk_t allocate(size_t s) noexcept { return {(char *)malloc<AdrSp>(s), s}; }
 
    static blk_t allocate_zero(size_t s) noexcept {
      if constexpr (AdrSp == mem::Host) {
        return {(char *)std::calloc(s, 1 /* byte */), s}; // NOLINT (C-style cast is fine here)
      } else {
        char *ptr = (char *)malloc<AdrSp>(s);
        memset<AdrSp>(ptr, 0, s);
        return {ptr, s};
      }
    }
 
    static void deallocate(blk_t b) noexcept { free<AdrSp>((void *)b.ptr); }
  };
 
  template <int ChunkSize>
  class bucket {
    // Unique pointer to handle memory allocation for the bucket.
    std::unique_ptr<char[]> _start = std::make_unique<char[]>(TotalChunkSize); // NOLINT (C-style array is fine here)
 
    // Pointer to the start of the bucket.
    char *p = _start.get();
 
    // Bitmask to keep track of which chunks are free.
    uint64_t flags = uint64_t(-1);
 
    public:
    static constexpr int TotalChunkSize = 64 * ChunkSize;
 
    static constexpr auto address_space = Host;
 
#ifdef NDA_USE_ASAN
    bucket() { __asan_poison_memory_region(p, TotalChunkSize); }
    ~bucket() { __asan_unpoison_memory_region(p, TotalChunkSize); }
#else
    bucket() = default;
#endif
    bucket(bucket const &) = delete;
 
    bucket(bucket &&) = default;
 
    bucket &operator=(bucket const &) = delete;
 
    bucket &operator=(bucket &&) = default;
 
    blk_t allocate(size_t s) noexcept {
      // check the size and if there is a free chunk, otherwise abort
      if (s > ChunkSize) std::abort();
      if (flags == 0) std::abort();
 
      // find the first free chunk
      int pos = __builtin_ctzll(flags);
 
      // update the bitmask and return the memory block
      flags &= ~(1ull << pos);
      blk_t b{p + static_cast<ptrdiff_t>(pos * ChunkSize), s};
#ifdef NDA_USE_ASAN
      __asan_unpoison_memory_region(b.ptr, ChunkSize);
#endif
      return b;
    }
 
    blk_t allocate_zero(size_t s) noexcept {
      auto blk = allocate(s);
      std::memset(blk.ptr, 0, s);
      return blk;
    }
 
    void deallocate(blk_t b) noexcept {
#ifdef NDA_USE_ASAN
      __asan_poison_memory_region(b.ptr, ChunkSize);
#endif
      int pos = (b.ptr - p) / ChunkSize;
      flags |= (1ull << pos);
    }
 
    [[nodiscard]] bool is_full() const noexcept { return flags == 0; }
 
    [[nodiscard]] bool empty() const noexcept { return flags == uint64_t(-1); }
 
    [[nodiscard]] const char *data() const noexcept { return p; }
 
    [[nodiscard]] auto mask() const noexcept { return flags; }
 
    [[nodiscard]] bool owns(blk_t b) const noexcept { return b.ptr >= p and b.ptr < p + TotalChunkSize; }
  };
 
  template <int ChunkSize>
  class multi_bucket {
    // Alias for the bucket allocator type.
    using b_t = bucket<ChunkSize>;
 
    // Vector of nda::mem::bucket allocators (ordered in memory).
    std::vector<b_t> bu_vec;
 
    // Iterator to the current bucket in use.
    typename std::vector<b_t>::iterator bu;
 
    // Find a bucket with a free chunk, otherwise create a new one and insert it at the correct position.
    [[gnu::noinline]] void find_non_full_bucket() {
      bu = std::find_if(bu_vec.begin(), bu_vec.end(), [](auto const &b) { return !b.is_full(); });
      if (bu != bu_vec.end()) return;
      b_t b;
      auto pos = std::upper_bound(bu_vec.begin(), bu_vec.end(), b, [](auto const &b1, auto const &b2) { return b1.data() < b2.data(); });
      bu       = bu_vec.insert(pos, std::move(b));
    }
 
    public:
    static constexpr auto address_space = Host;
 
    multi_bucket() : bu_vec(1), bu(bu_vec.begin()) {}
 
    multi_bucket(multi_bucket const &) = delete;
 
    multi_bucket(multi_bucket &&) = default;
 
    multi_bucket &operator=(multi_bucket const &) = delete;
 
    multi_bucket &operator=(multi_bucket &&) = default;
 
    blk_t allocate(size_t s) noexcept {
      if ((bu == bu_vec.end()) or (bu->is_full())) find_non_full_bucket();
      return bu->allocate(s);
    }
 
    blk_t allocate_zero(size_t s) noexcept {
      auto blk = allocate(s);
      std::memset(blk.ptr, 0, s);
      return blk;
    }
 
    void deallocate(blk_t b) noexcept {
      // try the current bucket first
      if (bu != bu_vec.end() and bu->owns(b)) {
        bu->deallocate(b);
        return;
      }
 
      // otherwise, find the owning bucket in the vector and deallocate
      bu = std::lower_bound(bu_vec.begin(), bu_vec.end(), b.ptr, [](auto const &b1, auto p) { return b1.data() <= p; });
      --bu;
      EXPECTS_WITH_MESSAGE((bu != bu_vec.end()), "Error in nda::mem::multi_bucket::deallocate: Owning bucket not found");
      EXPECTS_WITH_MESSAGE((bu->owns(b)), "Error in nda::mem::multi_bucket::deallocate: Owning bucket not found");
      bu->deallocate(b);
 
      // remove bucket the current bucket if it is empty and not the only one
      if (!bu->empty()) return;
      if (bu_vec.size() <= 1) return;
      bu_vec.erase(bu);
      bu = bu_vec.end();
    }
 
    [[nodiscard]] bool empty() const noexcept { return bu_vec.size() == 1 && bu_vec[0].empty(); }
 
    [[nodiscard]] auto const &buckets() const noexcept { return bu_vec; }
 
    [[nodiscard]] bool owns(blk_t b) const noexcept {
      bool res = false;
      for (const auto &mb : bu_vec) { res = res || mb.owns(b); }
      return res;
    }
  };
 
  template <size_t Threshold, Allocator A, Allocator B>
  class segregator {
    // Allocator for small memory blocks.
    A small;
 
    // Allocator for big memory blocks.
    B big;
 
    public:
    static_assert(A::address_space == B::address_space);
 
    static constexpr auto address_space = A::address_space;
 
    segregator() = default;
 
    segregator(segregator const &) = delete;
 
    segregator(segregator &&) = default;
 
    segregator &operator=(segregator const &) = delete;
 
    segregator &operator=(segregator &&) = default;
 
    blk_t allocate(size_t s) noexcept { return s <= Threshold ? small.allocate(s) : big.allocate(s); }
 
    blk_t allocate_zero(size_t s) noexcept { return s <= Threshold ? small.allocate_zero(s) : big.allocate_zero(s); }
 
    void deallocate(blk_t b) noexcept { return b.s <= Threshold ? small.deallocate(b) : big.deallocate(b); }
 
    [[nodiscard]] bool owns(blk_t b) const noexcept { return small.owns(b) or big.owns(b); }
  };
 
  template <Allocator A>
  class leak_check : A {
    // Total memory used by the allocator.
    long memory_used = 0;
 
    public:
    static constexpr auto address_space = A::address_space;
 
    leak_check() = default;
 
    leak_check(leak_check const &) = delete;
 
    leak_check(leak_check &&) = default;
 
    leak_check &operator=(leak_check const &) = delete;
 
    leak_check &operator=(leak_check &&) = default;
 
    ~leak_check() {
      if (!empty()) {
#ifndef NDEBUG
        std::cerr << "Memory leak in allocator: " << memory_used << " bytes leaked\n";
        std::abort();
#endif
      }
    }
 
    blk_t allocate(size_t s) {
      blk_t b = A::allocate(s);
      memory_used += b.s;
      return b;
    }
 
    blk_t allocate_zero(size_t s) {
      blk_t b = A::allocate_zero(s);
      memory_used += b.s;
      return b;
    }
 
    void deallocate(blk_t b) noexcept {
      memory_used -= b.s;
      if (memory_used < 0) {
#ifndef NDEBUG
        std::cerr << "Memory used by allocator < 0: Memory block to be deleted: b.s = " << b.s << ", b.ptr = " << (void *)b.ptr << "\n";
        std::abort();
#endif
      }
      A::deallocate(b);
    }
 
    [[nodiscard]] bool empty() const { return (memory_used == 0); }
 
    [[nodiscard]] bool owns(blk_t b) const noexcept { return A::owns(b); }
 
    [[nodiscard]] long get_memory_used() const noexcept { return memory_used; }
  };
 
  template <Allocator A>
  class stats : A {
    // Histogram of the allocation sizes.
    std::vector<uint64_t> hist = std::vector<uint64_t>(65, 0);
 
    public:
    static constexpr auto address_space = A::address_space;
 
    stats() = default;
 
    stats(stats const &) = delete;
 
    stats(stats &&) = default;
 
    stats &operator=(stats const &) = delete;
 
    stats &operator=(stats &&) = default;
 
    ~stats() {
#ifndef NDEBUG
      print_histogram(std::cerr);
#endif
    }
 
    blk_t allocate(uint64_t s) {
      // __builtin_clzl returns the number of leading zeros
      ++hist[__builtin_clzl(s)];
      return A::allocate(s);
    }
 
    blk_t allocate_zero(uint64_t s) {
      // __builtin_clzl returns the number of leading zeros
      ++hist[__builtin_clzl(s)];
      return A::allocate_zero(s);
    }
 
    void deallocate(blk_t b) noexcept { A::deallocate(b); }
 
    [[nodiscard]] bool owns(blk_t b) const noexcept { return A::owns(b); }
 
    [[nodiscard]] auto const &histogram() const noexcept { return hist; }
 
    void print_histogram(std::ostream &os) const {
      os << "Allocation size histogram :\n";
      os << "[0, 2^0): " << hist.back() << "\n";
      for (int i = 0; i < 64; ++i) { os << "[2^" << i << ", 2^" << i + 1 << "): " << hist[63 - i] << "\n"; }
    }
  };
 
} // namespace nda::mem