1 files changed, 21 insertions, 9 deletions
diff --git a/util/arena.cu b/util/arena.cu
index 82dbfe6..6d1b76b 100644
--- a/util/arena.cu
+++ b/util/arena.cu
@@ -62,17 +62,29 @@ __device__ char* Arena::AllocateAligned(size_t bytes) {
 
 __device__ char* Arena::AllocateNewBlock(size_t block_bytes) {
   char* result = nullptr;
+  cuda::atomic<ArenaNode*>* block_alloc = 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;
+  cudaMalloc((void **)&block_alloc, sizeof(cuda::atomic<ArenaNode*>));
+  while (true) {
+      ArenaNode * current_ = this->blocks_.load(cuda::memory_order_acquire);
+      if (current_ == nullptr) {
+      // cudaMalloc((void**)&this->blocks_, sizeof(ArenaNode));
+          ArenaNode * current_end = this->blocks_.load();
+          if (!this->blocks_.compare_exchange_weak(current_end, reinterpret_cast<ArenaNode*>(block_alloc)))
+            continue;
+          if (!this->head_.compare_exchange_weak(current_end, reinterpret_cast<ArenaNode*>(block_alloc)))
+             continue;
+          break;
+      }
+      ArenaNode * except_next = current_->next;
+      if (except_next != nullptr)
+          continue;
+      if (!this->blocks_.compare_exchange_weak(current_, reinterpret_cast<ArenaNode*>(block_alloc)))
+          continue;
+      current_->block = result;
+      current_->next = nullptr;
+      break;
   }
-  this->blocks_->block = result;
-  this->blocks_->next = nullptr;
   memory_usage_.fetch_add(block_bytes + sizeof(char*),
                           cuda::memory_order_relaxed);
   return result;