1 files changed, 8 insertions, 5 deletions
diff --git a/util/arena.cuh b/util/arena.cuh
index 68fc55d..618b426 100644
--- a/util/arena.cuh
+++ b/util/arena.cuh
@@ -11,6 +11,8 @@
 #include <cstdint>
 #include <vector>
 
+#include <cuda/atomic>
+
 namespace leveldb {
 
 class Arena {
@@ -26,30 +28,31 @@ class Arena {
   char* Allocate(size_t bytes);
 
   // Allocate memory with the normal alignment guarantees provided by malloc.
-  char* AllocateAligned(size_t bytes);
+  __device__ char* AllocateAligned(size_t bytes);
 
   // Returns an estimate of the total memory usage of data allocated
   // by the arena.
   size_t MemoryUsage() const {
-    return memory_usage_.load(std::memory_order_relaxed);
+    return memory_usage_.load(cuda::memory_order_relaxed);
   }
 
  private:
-  char* AllocateFallback(size_t bytes);
-  char* AllocateNewBlock(size_t block_bytes);
+  __device__ char* AllocateFallback(size_t bytes);
+  __device__ char* AllocateNewBlock(size_t block_bytes);
 
   // Allocation state
   char* alloc_ptr_;
   size_t alloc_bytes_remaining_;
 
   // Array of new[] allocated memory blocks
+  //thrust::host_vector<char *> blocks_;
   std::vector<char*> blocks_;
 
   // Total memory usage of the arena.
   //
   // TODO(costan): This member is accessed via atomics, but the others are
   //               accessed without any locking. Is this OK?
-  std::atomic<size_t> memory_usage_;
+  cuda::atomic<size_t> memory_usage_;
 };
 
 inline char* Arena::Allocate(size_t bytes) {