[Rhi] [bug] Fix the Unified Allocator to no longer return first two allocations as dupes

cmake/TaichiTests.cmake

@@ -17,6 +17,7 @@ file(GLOB_RECURSE TAICHI_TESTS_SOURCE
         "tests/cpp/common/*.cpp"
         "tests/cpp/ir/*.cpp"
         "tests/cpp/program/*.cpp"
+        "tests/cpp/rhi/common/*.cpp"
         "tests/cpp/struct/*.cpp"
         "tests/cpp/transforms/*.cpp"
         "tests/cpp/offline_cache/*.cpp")

taichi/rhi/common/unified_allocator.cpp

@@ -6,7 +6,7 @@
 
 namespace taichi::lang {
 
-const std::size_t UnifiedAllocator::default_allocator_size =
+std::size_t UnifiedAllocator::default_allocator_size =
     1 << 30;  // 1 GB per allocator
 
 template <typename T>
@@ -35,7 +35,6 @@ void *UnifiedAllocator::allocate(std::size_t size,
 
   // Note: put mutex on MemoryPool instead of Allocator, since Allocators are
   // transparent to user code
-  std::size_t allocation_size = size;
   if (!chunks_.empty() && !exclusive) {
     // Search for a non-exclusive chunk that has enough space
     for (size_t chunk_id = 0; chunk_id < chunks_.size(); chunk_id++) {
@@ -49,7 +48,7 @@ void *UnifiedAllocator::allocate(std::size_t size,
       auto ret = head + alignment - 1 - (head + alignment - 1) % alignment;
       TI_TRACE("UM [data={}] allocate() request={} remain={}", (intptr_t)data,
                size, (tail - head));
-      head = ret + allocation_size;
+      head = ret + size;
       if (head <= tail) {
         // success
         TI_ASSERT(ret % alignment == 0);
@@ -62,6 +61,7 @@ void *UnifiedAllocator::allocate(std::size_t size,
   // Allocate a new chunk
   MemoryChunk chunk;
 
+  std::size_t allocation_size = size;
   if (!exclusive) {
     // Do not allocate large memory chunks for "exclusive" allocation
     // to increate memory & performance efficiency
@@ -74,7 +74,7 @@ void *UnifiedAllocator::allocate(std::size_t size,
   void *ptr =
       HostMemoryPool::get_instance().allocate_raw_memory(allocation_size);
   chunk.data = ptr;
-  chunk.head = chunk.data;
+  chunk.head = (void *)((std::size_t)chunk.data + size);
   chunk.tail = (void *)((std::size_t)chunk.head + allocation_size);
   chunk.is_exclusive = exclusive;
 

taichi/rhi/common/unified_allocator.h

@@ -22,7 +22,7 @@ class UnifiedAllocator {
   };
 
  private:
-  static const std::size_t default_allocator_size;
+  static std::size_t default_allocator_size;
 
   UnifiedAllocator();
 
@@ -35,6 +35,7 @@ class UnifiedAllocator {
   std::vector<MemoryChunk> chunks_;
 
   friend class HostMemoryPool;
+  friend class HostMemoryPoolTestHelper;
 };
 
 }  // namespace taichi::lang

tests/cpp/rhi/common/host_memory_pool_test.cpp

@@ -0,0 +1,37 @@
+#include "gtest/gtest.h"
+
+#include "taichi/rhi/common/host_memory_pool.h"
+
+namespace taichi::lang {
+
+class HostMemoryPoolTestHelper {
+ public:
+  static void setDefaultAllocatorSize(std::size_t size) {
+    UnifiedAllocator::default_allocator_size = size;
+  }
+  static size_t getDefaultAllocatorSize() {
+    return UnifiedAllocator::default_allocator_size;
+  }
+};
+
+TEST(HostMemoryPool, AllocateMemory) {
+  auto oldAllocatorSize = HostMemoryPoolTestHelper::getDefaultAllocatorSize();
+  HostMemoryPoolTestHelper::setDefaultAllocatorSize(102400);  // 100KB
+
+  HostMemoryPool pool;
+
+  void *ptr1 = pool.allocate(1024, 16);
+  void *ptr2 = pool.allocate(1024, 16);
+  void *ptr3 = pool.allocate(1024, 16);
+
+  EXPECT_NE(ptr1, ptr2);
+  EXPECT_NE(ptr1, ptr3);
+  EXPECT_NE(ptr2, ptr3);
+
+  EXPECT_EQ((std::size_t)ptr2, (std::size_t)ptr1 + 1024);
+  EXPECT_EQ((std::size_t)ptr3, (std::size_t)ptr2 + 1024);
+
+  HostMemoryPoolTestHelper::setDefaultAllocatorSize(oldAllocatorSize);
+}
+
+}  // namespace taichi::lang