feat(skiplist): Convert skiplist to cuda version

Signed-off-by: KunoiSayami <[email protected]>

1 files changed, 28 insertions, 11 deletions

diff --git a/db/skiplist.cuh b/db/skiplist.cuh
index 5310d9c..2520a13 100644
--- a/db/skiplist.cuh
+++ b/db/skiplist.cuh
@@ -48,7 +48,8 @@ class SkipList {
   // Create a new SkipList object that will use "cmp" for comparing keys,
   // and will allocate memory using "*arena".  Objects allocated in the arena
   // must remain allocated for the lifetime of the skiplist object.
-  explicit SkipList(Comparator cmp, Arena* arena);
+  //explicit SkipList(Comparator cmp, Arena* arena);
+  __device__ explicit SkipList(Comparator cmp, Arena* arena);
 
   SkipList(const SkipList&) = delete;
   SkipList& operator=(const SkipList&) = delete;
@@ -111,7 +112,7 @@ class SkipList {
   __device__ bool Equal(const Key& a, const Key& b) const { return (compare_(a, b) == 0); }
 
   // Return true if key is greater than the data stored in "n"
-  bool KeyIsAfterNode(const Key& key, Node* n) const;
+  __device__ bool KeyIsAfterNode(const Key& key, Node* n) const;
 
   // Return the earliest node that comes at or after key.
   // Return nullptr if there is no such node.
@@ -145,38 +146,54 @@ class SkipList {
 // Implementation details follow
 template <typename Key, class Comparator>
 struct SkipList<Key, Comparator>::Node {
-  explicit Node(const Key& k) : key(k) {}
+  //explicit Node(const Key& k) : key(k) {}
 
   Key const key;
 
+  explicit __device__ Node(const Key& k): key(k) {}
+
+  __device__ void* operator new(size_t bytes) {
+    Node * ptr = nullptr;
+    cudaMalloc((void**)&ptr, bytes);
+    return ptr;
+  }
+
+  __device__ void* operator new(size_t bytes, void * const ptr) {
+    return ptr;
+  }
+
+  __device__ void operator delete(void *ptr) {
+      cudaFree(ptr);
+  }
+
   // Accessors/mutators for links.  Wrapped in methods so we can
   // add the appropriate barriers as necessary.
-  Node* Next(int n) {
+  __device__ Node* Next(int n) {
     assert(n >= 0);
     // Use an 'acquire load' so that we observe a fully initialized
     // version of the returned Node.
-    return next_[n].load(std::memory_order_acquire);
+    return next_[n].load(cuda::memory_order_acquire);
   }
   __device__ void SetNext(int n, Node* x) {
     assert(n >= 0);
     // Use a 'release store' so that anybody who reads through this
     // pointer observes a fully initialized version of the inserted node.
-    next_[n].store(x, std::memory_order_release);
+    next_[n].store(x, cuda::memory_order_release);
   }
 
   // No-barrier variants that can be safely used in a few locations.
   __device__ Node* NoBarrier_Next(int n) {
     assert(n >= 0);
-    return next_[n].load(std::memory_order_relaxed);
+    return next_[n].load(cuda::memory_order_relaxed);
   }
   __device__ void NoBarrier_SetNext(int n, Node* x) {
     assert(n >= 0);
-    next_[n].store(x, std::memory_order_relaxed);
+    next_[n].store(x, cuda::memory_order_relaxed);
   }
 
  private:
   // Array of length equal to the node height.  next_[0] is lowest level link.
-  std::atomic<Node*> next_[1];
+  cuda::atomic<Node*> next_[1];
 };
 
 template <typename Key, class Comparator>
@@ -253,7 +270,7 @@ __device__ int SkipList<Key, Comparator>::RandomHeight() {
 }
 
 template <typename Key, class Comparator>
-bool SkipList<Key, Comparator>::KeyIsAfterNode(const Key& key, Node* n) const {
+__device__ bool SkipList<Key, Comparator>::KeyIsAfterNode(const Key& key, Node* n) const {
   // null n is considered infinite
   return (n != nullptr) && (compare_(n->key, key) < 0);
 }
@@ -323,7 +340,7 @@ typename SkipList<Key, Comparator>::Node* SkipList<Key, Comparator>::FindLast()
 }
 
 template <typename Key, class Comparator>
-SkipList<Key, Comparator>::SkipList(Comparator cmp, Arena* arena)
+__device__ SkipList<Key, Comparator>::SkipList(Comparator cmp, Arena* arena)
     : compare_(cmp),
       arena_(arena),
       head_(NewNode(0 /* any key will do */, kMaxHeight)),