diff options
Diffstat (limited to 'util')
| -rw-r--r-- | util/arena.cu | 30 |
1 files changed, 17 insertions, 13 deletions
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<ArenaNode*>* 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<ArenaNode*>)); + 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<ArenaNode*>(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<ArenaNode*>(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<ArenaNode*>(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; |