From 1e4c8faa097c5423604e343bf20745f57c81f2a9 Mon Sep 17 00:00:00 2001
From: KunoiSayami <i@leanhe.dev>
Date: Wed, 26 Jan 2022 02:29:32 +0800
Subject: test(skiplist): Use atomic arena

Signed-off-by: KunoiSayami <i@leanhe.dev>
---
 util/arena.cu  | 30 +++++++++++++++++++++---------
 util/arena.cuh |  6 +++---
 2 files changed, 24 insertions(+), 12 deletions(-)

(limited to 'util')

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;
diff --git a/util/arena.cuh b/util/arena.cuh
index f9ccb1e..b70dcb9 100644
--- a/util/arena.cuh
+++ b/util/arena.cuh
@@ -42,7 +42,7 @@ class Arena {
 
   struct ArenaNode {
     char * block;
-    ArenaNode * next;
+    cuda::atomic<ArenaNode*> next;
   };
 
   // Allocation state
@@ -53,8 +53,8 @@ class Arena {
   //thrust::host_vector<char *> blocks_;
   //std::vector<char*> blocks_;
 
-  ArenaNode * head_;
-  ArenaNode * blocks_;
+  cuda::atomic<ArenaNode *> head_;
+  cuda::atomic<ArenaNode *> blocks_;
 
 
   // Total memory usage of the arena.
-- 
cgit v1.3.1