Make SlabAllocator use MemoryBlockAllocator

SlabAllocator now uses MemoryBlockAllocator as its backing block
allocator instead of malloc/free directly. Each SlabAllocator
instance owns a MemoryBlockAllocator sized to its slab requirements.

SlabAllocator also calls TrimMemory() on its MemoryBlockAllocator
during DeleteEmptySlabs().

Bug: 542009502
Change-Id: Ied6fa15d0169bf7dd6d3b534cc0aec0fc4797a57
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/329575
Commit-Queue: Quyen Le <lehoangquyen@google.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn/common/MemoryBlockAllocator.cpp b/src/dawn/common/MemoryBlockAllocator.cpp
index 467580a..c1a2ccd 100644
--- a/src/dawn/common/MemoryBlockAllocator.cpp
+++ b/src/dawn/common/MemoryBlockAllocator.cpp
@@ -90,6 +90,15 @@
     return DAWN_UNSAFE_BUFFERS(HeapArray<std::byte>::Uninit(mBlockSize));
 }
 
+void MemoryBlockAllocator::Return(HeapArray<std::byte>&& block) {
+    if (block.size() != mBlockSize) {
+        return;
+    }
+    mFreeList.Use([&](auto freeList) {
+        freeList->Append(FreeBlock::FromHeapArray(mTickSerial, std::move(block)));
+    });
+}
+
 void MemoryBlockAllocator::Return(std::vector<HeapArray<std::byte>>&& blocks) {
     mFreeList.Use([&](auto freeList) {
         for (HeapArray<std::byte>& block : blocks) {
diff --git a/src/dawn/common/MemoryBlockAllocator.h b/src/dawn/common/MemoryBlockAllocator.h
index 95bc532..6137c0d 100644
--- a/src/dawn/common/MemoryBlockAllocator.h
+++ b/src/dawn/common/MemoryBlockAllocator.h
@@ -67,6 +67,7 @@
     ~MemoryBlockAllocator();
 
     HeapArray<std::byte> Allocate(size_t minimumSize);
+    void Return(HeapArray<std::byte>&& block);
     void Return(std::vector<HeapArray<std::byte>>&& blocks);
     void Tick();
     void TrimMemory();
diff --git a/src/dawn/common/SlabAllocator.cpp b/src/dawn/common/SlabAllocator.cpp
index ba24dc0..b433b4f 100644
--- a/src/dawn/common/SlabAllocator.cpp
+++ b/src/dawn/common/SlabAllocator.cpp
@@ -33,6 +33,7 @@
 
 #include "src/dawn/common/AlignedAlloc.h"
 #include "src/dawn/common/Math.h"
+#include "src/dawn/common/MemoryBlockAllocator.h"
 #include "src/utils/assert.h"
 #include "src/utils/compiler.h"
 
@@ -46,8 +47,8 @@
 // Slab
 
 SlabAllocatorImpl::Slab::Slab() = default;
-SlabAllocatorImpl::Slab::Slab(char allocation[], IndexLinkNode* head)
-    : allocation(allocation), freeList(head) {}
+SlabAllocatorImpl::Slab::Slab(HeapArray<std::byte> allocation, IndexLinkNode* head)
+    : allocation(std::move(allocation)), freeList(head) {}
 
 SlabAllocatorImpl::Slab::Slab(Slab&& rhs) = default;
 
@@ -57,14 +58,18 @@
 SlabAllocatorImpl::SentinelSlab::SentinelSlab(SentinelSlab&& rhs) = default;
 
 SlabAllocatorImpl::SentinelSlab::~SentinelSlab() {
+    DAWN_CHECK(next == nullptr);
+}
+
+void SlabAllocatorImpl::SentinelSlab::Destroy(MemoryBlockAllocator* allocator) {
     // Delete the full linked list.
     while (next) {
         Slab* slab = next;
         slab->Splice();
         DAWN_ASSERT(slab->blocksInUse == 0);
-        char* allocation = slab->allocation;
+        HeapArray<std::byte> allocation = std::move(slab->allocation);
         slab->~Slab();  // Placement delete.
-        AlignedFree(allocation);
+        allocator->Return(std::move(allocation));
     }
 }
 
@@ -79,8 +84,12 @@
       mBlockStride(Align(mIndexLinkNodeOffset + u32_sizeof<IndexLinkNode>, objectAlignment)),
       mBlocksPerSlab(blocksPerSlab),
       mTotalAllocationSize(static_cast<size_t>(mSlabBlocksOffset) +
-                           static_cast<size_t>(mBlocksPerSlab) * mBlockStride) {
+                           static_cast<size_t>(mBlocksPerSlab) * mBlockStride),
+      mMemoryBlockAllocator(std::make_unique<MemoryBlockAllocator>(mTotalAllocationSize)) {
     DAWN_ASSERT(blocksPerSlab > 0);
+    // TODO(542009502): Currently only support standard alignment. So that standard malloc used by
+    // MemoryBlockAllocator can return a pointer satisfies this alignment.
+    DAWN_CHECK(mAllocationAlignment <= alignof(std::max_align_t));
     DAWN_ASSERT(IsPowerOfTwo(mAllocationAlignment));
 }
 
@@ -91,11 +100,16 @@
       mBlockStride(rhs.mBlockStride),
       mBlocksPerSlab(rhs.mBlocksPerSlab),
       mTotalAllocationSize(rhs.mTotalAllocationSize),
+      mMemoryBlockAllocator(std::move(rhs.mMemoryBlockAllocator)),
       mAvailableSlabs(std::move(rhs.mAvailableSlabs)),
       mFullSlabs(std::move(rhs.mFullSlabs)),
       mRecycledSlabs(std::move(rhs.mRecycledSlabs)) {}
 
-SlabAllocatorImpl::~SlabAllocatorImpl() = default;
+SlabAllocatorImpl::~SlabAllocatorImpl() {
+    mAvailableSlabs.Destroy(mMemoryBlockAllocator.get());
+    mFullSlabs.Destroy(mMemoryBlockAllocator.get());
+    mRecycledSlabs.Destroy(mMemoryBlockAllocator.get());
+}
 
 SlabAllocatorImpl::IndexLinkNode* SlabAllocatorImpl::OffsetFrom(
     IndexLinkNode* node,
@@ -175,16 +189,19 @@
 }
 
 void SlabAllocatorImpl::DeleteEmptySlabs() {
-    auto DeleteEmptyFromList = [](const SentinelSlab& sentinel) {
+    // TODO(542009502): this should be removed. SlabAllocator should eagerly return the memory block
+    // to the memory block allocator then memory block allocator will implicitly trim the memory
+    // when needed.
+    auto DeleteEmptyFromList = [&](const SentinelSlab& sentinel) {
         for (Slab* current = sentinel.next; current != nullptr;) {
             if (current->blocksInUse == 0) {
                 Slab* next = current->next;
 
                 // Remove from list and then delete to avoid dangling pointers.
                 current->Splice();
-                char* allocation = current->allocation;
+                HeapArray<std::byte> allocation = std::move(current->allocation);
                 current->~Slab();
-                AlignedFree(allocation);
+                mMemoryBlockAllocator->Return(std::move(allocation));
 
                 current = next;
             } else {
@@ -194,6 +211,7 @@
     };
     DeleteEmptyFromList(mRecycledSlabs);
     DeleteEmptyFromList(mAvailableSlabs);
+    mMemoryBlockAllocator->TrimMemory();
 }
 
 uint32_t SlabAllocatorImpl::CountAllocatedSlabsForTesting() const {
@@ -264,9 +282,11 @@
         return;
     }
 
-    char* alignedPtr = static_cast<char*>(AlignedAlloc(mTotalAllocationSize, mAllocationAlignment));
+    HeapArray<std::byte> allocation = mMemoryBlockAllocator->Allocate(mTotalAllocationSize);
+    std::byte* alignedPtr = allocation.data();
+    DAWN_CHECK(IsPtrAligned(alignedPtr, mAllocationAlignment));
 
-    char* dataStart = DAWN_UNSAFE_TODO(alignedPtr + mSlabBlocksOffset);
+    void* dataStart = allocation.subspan(mSlabBlocksOffset).data();
 
     IndexLinkNode* node = NodeFromObject(dataStart);
     for (Index i = 0; i < mBlocksPerSlab; ++i) {
@@ -277,7 +297,7 @@
         OffsetFrom(node, checked_cast<std::make_signed_t<Index>>(mBlocksPerSlab - 1));
     lastNode->nextIndex = kInvalidIndex;
 
-    mAvailableSlabs.Prepend(new (alignedPtr) Slab(alignedPtr, node));
+    mAvailableSlabs.Prepend(new (alignedPtr) Slab(std::move(allocation), node));
 }
 
 }  // namespace dawn
diff --git a/src/dawn/common/SlabAllocator.h b/src/dawn/common/SlabAllocator.h
index 5a12088..71be4f9 100644
--- a/src/dawn/common/SlabAllocator.h
+++ b/src/dawn/common/SlabAllocator.h
@@ -31,16 +31,20 @@
 #include <algorithm>
 #include <cstdint>
 #include <limits>
+#include <memory>
 #include <type_traits>
 #include <utility>
 
 #include "partition_alloc/pointers/raw_ptr.h"
 #include "src/dawn/common/Numeric.h"
 #include "src/dawn/common/PlacementAllocated.h"
+#include "src/utils/heap_array.h"
 #include "src/utils/numeric.h"
 
 namespace dawn {
 
+class MemoryBlockAllocator;
+
 // The SlabAllocator allocates objects out of one or more fixed-size contiguous "slabs" of memory.
 // This makes it very quick to allocate and deallocate fixed-size objects because the allocator only
 // needs to index an offset into pre-allocated memory. It is similar to a pool-allocator that
@@ -107,13 +111,13 @@
         // | ---------- allocation --------- |
         // | pad | Slab | data ------------> |
         Slab();
-        Slab(char allocation[], IndexLinkNode* head);
+        Slab(HeapArray<std::byte> allocation, IndexLinkNode* head);
         Slab(Slab&& rhs);
 
         // Extract the Slab from the linked list.
         void Splice();
 
-        raw_ptr<char> allocation = nullptr;
+        HeapArray<std::byte> allocation;
         raw_ptr<IndexLinkNode> freeList = nullptr;
 
         raw_ptr<Slab> prev = nullptr;
@@ -174,6 +178,9 @@
     const Index mBlocksPerSlab;  // The total number of blocks in a slab.
 
     const size_t mTotalAllocationSize;
+    // TODO(crbug.com/398193014): Reuse Device's MemoryBlockAllocator across multiple
+    // SlabAllocators.
+    std::unique_ptr<MemoryBlockAllocator> mMemoryBlockAllocator;
 
     struct SentinelSlab : Slab {
         SentinelSlab();
@@ -181,6 +188,7 @@
 
         SentinelSlab(SentinelSlab&& rhs);
 
+        void Destroy(MemoryBlockAllocator* allocator);
         void Prepend(Slab* slab);
     };
 
diff --git a/src/dawn/tests/unittests/SlabAllocatorTests.cpp b/src/dawn/tests/unittests/SlabAllocatorTests.cpp
index edaf2f7..f869f3c 100644
--- a/src/dawn/tests/unittests/SlabAllocatorTests.cpp
+++ b/src/dawn/tests/unittests/SlabAllocatorTests.cpp
@@ -25,6 +25,7 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#include <cstddef>
 #include <set>
 #include <vector>
 
@@ -41,7 +42,7 @@
     uint32_t value;
 };
 
-struct alignas(256) AlignedFoo : public Foo {
+struct alignas(std::max_align_t) MaxAlignedFoo : public Foo {
     using Foo::Foo;
 };
 
@@ -82,11 +83,11 @@
         }
     }
 
-    // Check large alignment
+    // Check max alignment
     {
-        SlabAllocator<AlignedFoo> allocator(9 * sizeof(AlignedFoo));
+        SlabAllocator<MaxAlignedFoo> allocator(9 * sizeof(MaxAlignedFoo));
 
-        std::vector<AlignedFoo*> objects;
+        std::vector<MaxAlignedFoo*> objects;
         for (uint32_t i = 0; i < 21; ++i) {
             auto* ptr = allocator.Allocate(i);
             EXPECT_TRUE(std::find(objects.begin(), objects.end(), ptr) == objects.end());
@@ -98,11 +99,11 @@
             EXPECT_EQ(objects[i]->value, i);
 
             // Check that the alignment is correct.
-            EXPECT_TRUE(IsPtrAligned(objects[i], 256));
+            EXPECT_TRUE(IsPtrAligned(objects[i], alignof(std::max_align_t)));
         }
 
         // Deallocate all of the objects.
-        for (AlignedFoo* object : objects) {
+        for (MaxAlignedFoo* object : objects) {
             allocator.Deallocate(object);
         }
     }
@@ -252,13 +253,20 @@
 // larger that the totalObjectBytes would allocate space for no objects but still attempt to fulfill
 // requests.
 TEST(SlabAllocatorTests, TotalObjectBytesTooSmall) {
-    SlabAllocator<AlignedFoo> allocator(sizeof(AlignedFoo) - 1);
+    SlabAllocator<Foo> allocator(sizeof(Foo) - 1);
 
-    AlignedFoo* obj = allocator.Allocate(4u);
+    Foo* obj = allocator.Allocate(4u);
     EXPECT_EQ(obj->value, 4u);
 
     allocator.Deallocate(obj);
 }
 
+TEST(SlabAllocatorDeathTest, AlignmentLargerThanMaxAlignT) {
+    struct alignas(256) OverAlignedFoo : public Foo {
+        using Foo::Foo;
+    };
+    EXPECT_DEATH_IF_SUPPORTED(SlabAllocator<OverAlignedFoo>(100), "");
+}
+
 }  // anonymous namespace
 }  // namespace dawn