diff --git a/llvm/include/llvm/Support/ThreadPoolAllocator.h b/llvm/include/llvm/Support/ThreadPoolAllocator.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Support/ThreadPoolAllocator.h @@ -0,0 +1,95 @@ +//===- ThreadPoolAllocator.h ------------------------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_SUPPORT_THREADPOOLALLOCATOR_H +#define LLVM_SUPPORT_THREADPOOLALLOCATOR_H + +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Allocator.h" +#include "llvm/Support/Parallel.h" + +namespace llvm { +namespace parallel { + +/// ThreadPoolAllocator allows separating allocations by thread id. +/// It is possible because ThreadPoolExecutor creates threads, keeps them until +/// the destructor of ThreadPoolExecutor is called, and assigns ids to the +/// threads. Thus ThreadExecutorAllocator should be used with only threads +/// created by ThreadPoolExecutor. To work properly, ThreadPoolExecutor should +/// be initialized before ThreadExecutorAllocator is created. +/// The initDefaultThreadsExecutor() should be used for that. +/// TODO: The same approach might be implemented for ThreadPool. + +template class ThreadPoolAllocator { +public: + ThreadPoolAllocator() { Allocators.resize(parallel::getMaxThreadsNum()); } + + /// Allocate \a Size bytes of \a Alignment aligned memory. + void *Allocate(size_t Size, size_t Alignment) { + return Allocators[getThreadIndex()].Allocate(Size, Alignment); + } + + /// Deallocate \a Ptr to \a Size bytes of memory allocated by this + /// allocator. + void Deallocate(const void *Ptr, size_t Size, size_t Alignment) { + return Allocators[getThreadIndex()].Deallocate(Ptr, Size, Alignment); + } + + // Reset state of allocators. + void Reset() { + for (AllocatorTy &Allocator : Allocators) + Allocator.Reset(); + } + + // Return total memory size used by all allocators. + size_t getTotalMemory() const { + size_t TotalMemory = 0; + + for (const AllocatorTy &Allocator : Allocators) + TotalMemory += Allocator.getTotalMemory(); + + return TotalMemory; + } + + // Return allocated size by all allocators. + size_t getBytesAllocated() const { + size_t BytesAllocated = 0; + + for (const AllocatorTy &Allocator : Allocators) + BytesAllocated += Allocator.getBytesAllocated(); + + return BytesAllocated; + } + + // Set red zone for all allocators. + void setRedZoneSize(size_t NewSize) { + for (AllocatorTy &Allocator : Allocators) + Allocator.setRedZoneSize(NewSize); + } + + // Print statistic for each allocator. + void PrintStats() const { + for (size_t Idx = 0; Idx < Allocators.size(); Idx++) { + errs() << "\n Allocator " << Idx << "\n"; + Allocators[Idx].PrintStats(); + } + } + + // Return number of used allocators. + size_t getNumberOfAllocators() const { return Allocators.size(); } + +protected: + SmallVector Allocators; +}; + +using ThreadSafeBumpPtrAllocator = class ThreadPoolAllocator; + +} // end namespace parallel +} // end namespace llvm + +#endif // LLVM_SUPPORT_THREADPOOLALLOCATOR_H diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt --- a/llvm/unittests/Support/CMakeLists.txt +++ b/llvm/unittests/Support/CMakeLists.txt @@ -81,6 +81,7 @@ TaskQueueTest.cpp ThreadPool.cpp Threading.cpp + ThreadSafeBumpPtrAllocatorTest.cpp TimerTest.cpp TimeProfilerTest.cpp ToolOutputFileTest.cpp diff --git a/llvm/unittests/Support/ThreadSafeBumpPtrAllocatorTest.cpp b/llvm/unittests/Support/ThreadSafeBumpPtrAllocatorTest.cpp new file mode 100644 --- /dev/null +++ b/llvm/unittests/Support/ThreadSafeBumpPtrAllocatorTest.cpp @@ -0,0 +1,53 @@ +//===- ThreadSafeBumpPtrAllocatorTest.cpp ---------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "llvm/Support/Parallel.h" +#include "llvm/Support/ThreadPoolAllocator.h" +#include "gtest/gtest.h" +#include + +using namespace llvm; +using namespace parallel; + +namespace { + +TEST(ThreadSafeBumpPtrAllocatorTest, Simple) { + ThreadSafeBumpPtrAllocator Allocator; + + uint64_t *Var = + (uint64_t *)Allocator.Allocate(sizeof(uint64_t), alignof(uint64_t)); + + *Var = 0xFE; + EXPECT_EQ(0xFEul, *Var); + EXPECT_EQ(sizeof(uint64_t), Allocator.getBytesAllocated()); + EXPECT_TRUE(Allocator.getBytesAllocated() <= Allocator.getTotalMemory()); + + ThreadSafeBumpPtrAllocator Allocator2(std::move(Allocator)); + + EXPECT_EQ(sizeof(uint64_t), Allocator2.getBytesAllocated()); + EXPECT_TRUE(Allocator2.getBytesAllocated() <= Allocator2.getTotalMemory()); + + EXPECT_EQ(0xFEul, *Var); +} + +TEST(ThreadSafeBumpPtrAllocatorTest, ParallelAllocation) { + ThreadSafeBumpPtrAllocator Allocator; + + static size_t constexpr NumAllocations = 1000; + + parallelFor(0, NumAllocations, [&](size_t Idx) { + uint64_t *ptr = + (uint64_t *)Allocator.Allocate(sizeof(uint64_t), alignof(uint64_t)); + *ptr = Idx; + }); + + EXPECT_EQ(sizeof(uint64_t) * NumAllocations, Allocator.getBytesAllocated()); + EXPECT_EQ(Allocator.getNumberOfAllocators(), parallel::getMaxThreadsNum()); +} + +} // anonymous namespace