diff options
Diffstat (limited to 'util')
| -rw-r--r-- | util/arena.cu | 30 | ||||
| -rw-r--r-- | util/arena.cuh | 6 |
2 files changed, 24 insertions, 12 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; 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. |