From 1e4c8faa097c5423604e343bf20745f57c81f2a9 Mon Sep 17 00:00:00 2001 From: KunoiSayami Date: Wed, 26 Jan 2022 02:29:32 +0800 Subject: test(skiplist): Use atomic arena Signed-off-by: KunoiSayami --- util/arena.cu | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) (limited to 'util/arena.cu') 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* 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)); + 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(block_alloc))) + continue; + if (!this->head_.compare_exchange_weak(current_end, reinterpret_cast(block_alloc))) + continue; + break; + } + ArenaNode * except_next = current_->next; + if (except_next != nullptr) + continue; + if (!this->blocks_.compare_exchange_weak(current_, reinterpret_cast(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; -- cgit v1.3.1