From f83373a5486bf7cad48fbd662c7f7f9484e79a65 Mon Sep 17 00:00:00 2001 From: KunoiSayami Date: Thu, 17 Feb 2022 19:25:39 +0800 Subject: fix(arena): Fix arena alloc Signed-off-by: KunoiSayami --- util/arena.cu | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) (limited to 'util') diff --git a/util/arena.cu b/util/arena.cu index 6d1b76b..f954577 100644 --- a/util/arena.cu +++ b/util/arena.cu @@ -61,30 +61,34 @@ __device__ char* Arena::AllocateAligned(size_t bytes) { } __device__ char* Arena::AllocateNewBlock(size_t block_bytes) { + // Allocate new target char* result = nullptr; - cuda::atomic* block_alloc = nullptr; + // Allocate storage target + ArenaNode* node = nullptr; + // Malloc memory cudaMalloc((void **)&result, sizeof(char) * block_bytes); - cudaMalloc((void **)&block_alloc, sizeof(cuda::atomic)); + cudaMalloc((void **)&node, sizeof(ArenaNode)); + + node->next.store(nullptr); + 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))) + if (!this->blocks_.compare_exchange_strong(current_, node)) + continue; + if (!this->head_.compare_exchange_strong(current_, node)) continue; - if (!this->head_.compare_exchange_weak(current_end, reinterpret_cast(block_alloc))) - continue; + node->block = result; break; } - ArenaNode * except_next = current_->next; - if (except_next != nullptr) - continue; - if (!this->blocks_.compare_exchange_weak(current_, reinterpret_cast(block_alloc))) - continue; + ArenaNode * expect_next = nullptr; + if (current_->next.compare_exchange_strong(expect_next, node)) { + continue; + } current_->block = result; - current_->next = nullptr; break; } + memory_usage_.fetch_add(block_bytes + sizeof(char*), cuda::memory_order_relaxed); return result; -- cgit v1.3.1