2 files changed, 34 insertions, 12 deletions

diff --git a/util/arena.cu b/util/arena.cu
index 5075318..db5827d 100644
--- a/util/arena.cu
+++ b/util/arena.cu
@@ -9,12 +9,16 @@ namespace leveldb {
 static const int kBlockSize = 4096;
 
 Arena::Arena()
-    : alloc_ptr_(nullptr), alloc_bytes_remaining_(0), memory_usage_(0) {}
+    : alloc_ptr_(nullptr), alloc_bytes_remaining_(0), memory_usage_(0),
+      head_(nullptr), blocks_(nullptr) {}
 
 Arena::~Arena() {
-  for (size_t i = 0; i < blocks_.size(); i++) {
-    //cudaFree(blocks_[i]);
-    delete [] blocks_[i];
+  ArenaNode * current = this->head_;
+  while (current != nullptr) {
+    ArenaNode * next = current->next;
+    cudaFree(current->block);
+    cudaFree(current);
+    current = next;
   }
 }
 
@@ -56,12 +60,21 @@ __device__ char* Arena::AllocateAligned(size_t bytes) {
   return result;
 }
 
-char* Arena::AllocateNewBlock(size_t block_bytes) {
-  char* result = new char[block_bytes];
-  //cudaMallocManaged((void **)&result, sizeof(char) * block_bytes);
-  blocks_.push_back(result);
+__device__ char* Arena::AllocateNewBlock(size_t block_bytes) {
+  char* result = nullptr;
+  cudaMalloc((void **)&result, sizeof(char) * block_bytes);
+  if (this->blocks_ == nullptr) {
+    cudaMalloc((void**)&this->blocks_, sizeof(ArenaNode));
+    // First alloc
+    this->head_ = this->blocks_;
+  } else {
+    cudaMalloc((void**)&this->blocks_->next, sizeof(ArenaNode));
+    this->blocks_ = this->blocks_->next;
+  }
+  this->blocks_->block = result;
+  this->blocks_->next = nullptr;
   memory_usage_.fetch_add(block_bytes + sizeof(char*),
-                          std::memory_order_relaxed);
+                          cuda::memory_order_relaxed);
   return result;
 }
 
diff --git a/util/arena.cuh b/util/arena.cuh
index 618b426..afbc575 100644
--- a/util/arena.cuh
+++ b/util/arena.cuh
@@ -25,7 +25,7 @@ class Arena {
   ~Arena();
 
   // Return a pointer to a newly allocated memory block of "bytes" bytes.
-  char* Allocate(size_t bytes);
+  __device__ char* Allocate(size_t bytes);
 
   // Allocate memory with the normal alignment guarantees provided by malloc.
   __device__ char* AllocateAligned(size_t bytes);
@@ -40,13 +40,22 @@ class Arena {
   __device__ char* AllocateFallback(size_t bytes);
   __device__ char* AllocateNewBlock(size_t block_bytes);
 
+  struct ArenaNode {
+    char * block;
+    ArenaNode * next;
+  };
+
   // Allocation state
   char* alloc_ptr_;
   size_t alloc_bytes_remaining_;
 
   // Array of new[] allocated memory blocks
   //thrust::host_vector<char *> blocks_;
-  std::vector<char*> blocks_;
+  //std::vector<char*> blocks_;
+
+  ArenaNode * head_;
+  ArenaNode * blocks_;
+
 
   // Total memory usage of the arena.
   //
@@ -55,7 +64,7 @@ class Arena {
   cuda::atomic<size_t> memory_usage_;
 };
 
-inline char* Arena::Allocate(size_t bytes) {
+__device__ inline char* Arena::Allocate(size_t bytes) {
   // The semantics of what to return are a bit messy if we allow
   // 0-byte allocations, so we disallow them here (we don't need
   // them for our internal use).