Skip to content

Commit

Permalink
Fixed bugs in nucc_allocator
Browse files Browse the repository at this point in the history
  • Loading branch information
@abkein
abkein committed Dec 8, 2024
1 parent d2324be commit 80572bf
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 26 deletions.
12 changes: 6 additions & 6 deletions src/compute_cluster_size_ext.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,15 @@ ComputeClusterSizeExt::ComputeClusterSizeExt(LAMMPS* lmp, int narg, char** arg)
size_cutoff = utils::inumeric(FLERR, arg[4], true, lmp);
if (size_cutoff < 1) { error->all(FLERR, "size_cutoff for compute cluster/size must be greater than 0"); }

keeper1 = new MemoryKeeper<MapMember_t<int, int>>(memory);
keeper1 = new MemoryKeeper(memory);
cluster_map_allocator = new MapAlloc_t<int, int>(keeper1);
cluster_map = new Map_t<int, int>(*cluster_map_allocator);

keeper2 = new MemoryKeeper<MapMember_t<int, Vec_t<int>>>(memory);
keeper2 = new MemoryKeeper(memory);
alloc_map_vec1 = new MapAlloc_t<int, Vec_t<int>>(keeper2);
cIDs_by_size = new Map_t<int, Vec_t<int>>(*alloc_map_vec1);

keeper3 = new MemoryKeeper<MapMember_t<int, Vec_t<int>>>(memory);
keeper3 = new MemoryKeeper(memory);
alloc_map_vec2 = new MapAlloc_t<int, Vec_t<int>>(keeper3);
cIDs_by_size_all = new Map_t<int, Vec_t<int>>(*alloc_map_vec2);

Expand Down Expand Up @@ -112,9 +112,9 @@ void ComputeClusterSizeExt::init()
vector = dist.data();

nloc = static_cast<int>(atom->nlocal * LMP_NUCC_ALLOC_COEFF);
keeper1->pool_size(nloc);
keeper2->pool_size(nloc);
keeper3->pool_size(nloc);
keeper1->pool_size<MapMember_t<int, int>>(nloc);
keeper2->pool_size<MapMember_t<int, Vec_t<int>>>(nloc);
keeper3->pool_size<MapMember_t<int, Vec_t<int>>>(nloc);

cluster_map->reserve(nloc);
cIDs_by_size->reserve(size_cutoff);
Expand Down
8 changes: 4 additions & 4 deletions src/compute_cluster_size_ext.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,22 +101,22 @@ class ComputeClusterSizeExt : public Compute {
inline constexpr const NUCC::Map_t<int, int> *get_cluster_map() const noexcept(true) { return cluster_map; }
inline constexpr const NUCC::Map_t<int, NUCC::Vec_t<int>> *get_cIDs_by_size() const noexcept(true) { return cIDs_by_size; }
inline constexpr const NUCC::Map_t<int, NUCC::Vec_t<int>> *get_cIDs_by_size_all() const noexcept(true) { return cIDs_by_size_all; }
inline constexpr const NUCC::cspan<const NUCC::cluster_data> &get_clusters() const noexcept(true) { return clusters; }
inline constexpr const NUCC::cspan<const NUCC::cluster_data> get_clusters() const noexcept(true) { return clusters; }

private:
int size_cutoff; // number of elements reserved in dist

NUCC::MemoryKeeper<NUCC::MapMember_t<int, int>> *keeper1;
NUCC::MemoryKeeper *keeper1;
NUCC::MapAlloc_t<int, int> *cluster_map_allocator;
NUCC::Map_t<int, int> *cluster_map;
// std::unordered_map<int, int> cluster_map; // clid -> idx

NUCC::MemoryKeeper<NUCC::MapMember_t<int, NUCC::Vec_t<int>>> *keeper2;
NUCC::MemoryKeeper *keeper2;
NUCC::MapAlloc_t<int, NUCC::Vec_t<int>> *alloc_map_vec1;
NUCC::Map_t<int, NUCC::Vec_t<int>> *cIDs_by_size;
// std::unordered_map<int, std::vector<int>> cIDs_by_size; // size -> vector(idx)

NUCC::MemoryKeeper<NUCC::MapMember_t<int, NUCC::Vec_t<int>>> *keeper3;
NUCC::MemoryKeeper *keeper3;
NUCC::MapAlloc_t<int, NUCC::Vec_t<int>> *alloc_map_vec2;
NUCC::Map_t<int, NUCC::Vec_t<int>> *cIDs_by_size_all;
// std::unordered_map<int, std::vector<int>> cIDs_by_size_all;
Expand Down
41 changes: 25 additions & 16 deletions src/nucc_allocator.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@

namespace NUCC {

template <typename T>
class MemoryKeeper {
public:
MemoryKeeper() = delete;
Expand All @@ -20,15 +19,16 @@ class MemoryKeeper {
MemoryKeeper& operator=(const MemoryKeeper&) = delete;
MemoryKeeper& operator=(MemoryKeeper&&) = delete;

constexpr MemoryKeeper(Memory* memory) noexcept : memory_(memory) {}
constexpr MemoryKeeper(LAMMPS_NS::Memory* memory) noexcept : memory_(memory) {}
~MemoryKeeper() noexcept(noexcept(clear())) { clear(); }

template <typename T>
void store(T*& ptr, const size_t size) noexcept(noexcept(infos.emplace_back(ptr, size)))
{
infos.emplace_back(ptr, size);
}

void clear() noexcept(noexcept(std::declval<Memory>().destroy<void>(std::declval<void*&>())))
void clear() noexcept(noexcept(std::declval<LAMMPS_NS::Memory>().destroy<void>(std::declval<void*&>())))
{
for (auto& pool : infos) { memory_->destroy(pool.ptr); }
}
Expand All @@ -39,31 +39,40 @@ class MemoryKeeper {
return _pool_size;
}

template <typename T>
inline constexpr void pool_size(std::size_t n) noexcept { _pool_size = n * sizeof(T); }

inline constexpr void pool_size(std::size_t n) noexcept { _pool_size = n; }

template <typename T>
T* allocate(const std::size_t n)
{
T* ptr = nullptr;
if (n > _pool_size) {
std::size_t pool_size_T = _pool_size / sizeof(T) + 1;
if (n > pool_size_T) {
// If requested size is larger than pool, allocate separately
memory_->create<T>(ptr, n, "CustomAllocator_Large");
infos.emplace_back(ptr, n);
return ptr;
}
if (current == nullptr || left < n) {
std::size_t nbytes = n * sizeof(T);
T* _current = reinterpret_cast<T*>(current);
if ((current == nullptr) || (left < nbytes)) {
// Pool is full or not initialized, request a new pool
memory_->create<T>(current, _pool_size, "CustomAllocator_Pool");
infos.emplace_back(current, _pool_size);
memory_->create(_current, pool_size_T, "CustomAllocator_Pool");
infos.emplace_back(_current, pool_size_T);
left = _pool_size;
}
ptr = current;
left -= n;
current += n;
ptr = _current;
left -= nbytes;
_current += n;
current = reinterpret_cast<void *>(_current);
return ptr;
}

private:
struct PoolInfo {
template <typename T>
constexpr PoolInfo(T*& ptr, const size_t size) noexcept : ptr(reinterpret_cast<void*>(ptr)), size(size * sizeof(T))
{
}
Expand All @@ -74,9 +83,9 @@ class MemoryKeeper {
size_t size = 0;
};

T* current = nullptr;
void* current = nullptr;
std::size_t left = 0;
Memory* const memory_ = nullptr;
LAMMPS_NS::Memory* const memory_ = nullptr;
std::size_t _pool_size = 0;
std::vector<PoolInfo> infos;
};
Expand All @@ -95,16 +104,16 @@ class CustomAllocator {

CustomAllocator() = delete;

constexpr CustomAllocator(MemoryKeeper<T>* const keeper) noexcept : keeper_(keeper) {}
constexpr CustomAllocator(MemoryKeeper* const keeper) noexcept : keeper_(keeper) {}

template <typename U>
constexpr CustomAllocator(const CustomAllocator<U>& other) noexcept : keeper_(other.keeper_)
{
}

inline T* allocate(const std::size_t n) const { return keeper_->allocate(n); }
inline T* allocate(std::size_t n) const { return keeper_->allocate<T>(n); }

inline constexpr void deallocate(T /**p*/, const std::size_t /*n*/) const noexcept
inline constexpr void deallocate(T* /*p*/, const std::size_t /*n*/) const noexcept
{
// Deallocation can be handled when the allocator is destroyed
// For pool allocator, individual deallocations are often no-ops
Expand All @@ -124,7 +133,7 @@ class CustomAllocator {
}

private:
MemoryKeeper<T>* const keeper_;
MemoryKeeper* const keeper_;
};

} // namespace NUCC
Expand Down

0 comments on commit 80572bf

Please sign in to comment.