Priority Queue

benchmarks/CMakeLists.txt

@@ -50,3 +50,7 @@ ConfigureBench(STATIC_MAP_BENCH "${STATIC_MAP_BENCH_SRC}")
 ###################################################################################################
 set(RBK_BENCH_SRC "${CMAKE_CURRENT_SOURCE_DIR}/reduce_by_key/reduce_by_key.cu")
 ConfigureBench(RBK_BENCH "${RBK_BENCH_SRC}")
+
+#TODO: Port priority_queue benchmark to google benchmark
+add_executable(PRIORITY_QUEUE_BENCH "${CMAKE_CURRENT_SOURCE_DIR}/priority_queue/priority_queue_bench.cu")
+target_link_libraries(PRIORITY_QUEUE_BENCH cuco)

benchmarks/priority_queue/priority_queue_bench.cu

@@ -0,0 +1,187 @@
+#include <iostream>
+#include <vector>
+#include <cstdint>
+#include <cstdlib>
+
+#include <cuda_runtime.h>
+
+#include <cuco/priority_queue.cuh>
+#include <cuco/detail/error.hpp>
+
+using namespace cuco;
+
+template <typename Key, typename Value>
+__global__ void DeviceAPIInsert(
+                typename priority_queue<Key, Value>::device_mutable_view view,
+                Pair<Key, Value> *elements,
+                size_t num_elements) {
+  extern __shared__ int shmem[];
+  thread_block g = this_thread_block(); 
+  for (size_t i = blockIdx.x * view.get_node_size();
+       i < num_elements; i += gridDim.x * view.get_node_size()) {
+    view.push(g, elements + i, min(view.get_node_size(), num_elements - i),
+              shmem);
+  }
+}
+
+template <typename Key, typename Value>
+__global__ void DeviceAPIDelete(
+                typename priority_queue<Key, Value>::device_mutable_view view,
+                Pair<Key, Value> *out,
+                size_t num_elements) {
+
+  extern __shared__ int shmem[];
+  thread_block g = this_thread_block(); 
+  for (size_t i = blockIdx.x * view.get_node_size();
+       i < num_elements; i += gridDim.x * view.get_node_size()) {
+    view.pop(g, out + i, min(view.get_node_size(), num_elements - i), shmem);
+  }
+}
+
+// Use CUDA events to time the code in the lambda function
+template <typename F>
+float TimeCode(F f) {
+  cudaEvent_t t1;
+  CUCO_CUDA_TRY(cudaEventCreate(&t1));
+
+  cudaEvent_t t2;
+  CUCO_CUDA_TRY(cudaEventCreate(&t2));
+
+  CUCO_CUDA_TRY(cudaEventRecord(t1));
+  f();
+  CUCO_CUDA_TRY(cudaEventRecord(t2));
+
+  CUCO_CUDA_TRY(cudaEventSynchronize(t1));
+  CUCO_CUDA_TRY(cudaEventSynchronize(t2));
+
+  float result;
+  CUCO_CUDA_TRY(cudaEventElapsedTime(&result, t1, t2));
+  return result;
+}
+
+// Time the insertion of the num_keys elements at d_elements into pq in ms
+float TimeInsert(priority_queue<uint32_t, uint32_t> &pq,
+                            Pair<uint32_t, uint32_t> *d_elements,
+                            size_t num_keys) {
+  return TimeCode([&]() {
+    pq.push(d_elements, num_keys);
+  });
+}
+
+// Time insert of the num_keys elements with the device API at d_elements
+// into pq in ms
+float TimeInsertDeviceAPI(priority_queue<uint32_t, uint32_t> &pq,
+                            Pair<uint32_t, uint32_t> *d_elements,
+                            size_t num_keys) {
+  return TimeCode([&]() {
+    DeviceAPIInsert<<<64000, 256, pq.get_shmem_size(256)>>>
+                   (pq.get_mutable_device_view(), d_elements, num_keys);
+  });
+}
+
+// Time the deletion of num_keys elements from pq in ms
+float TimeDeleteDeviceAPI(priority_queue<uint32_t, uint32_t> &pq,
+                            Pair<uint32_t, uint32_t> *d_elements,
+                            size_t num_keys) {
+  return TimeCode([&]() {
+    DeviceAPIDelete<<<32000, 512, pq.get_shmem_size(512)>>>
+                   (pq.get_mutable_device_view(), d_elements, num_keys);
+  });
+}
+
+// Time the deletion of num_keys elements from pq in ms
+float TimeDelete(priority_queue<uint32_t, uint32_t> &pq,
+                            Pair<uint32_t, uint32_t> *d_elements,
+                            size_t num_keys) {
+  return TimeCode([&]() {
+    pq.pop(d_elements, num_keys);
+  });
+}
+
+// Follow the first experiment in the paper,
+// inserting 512 million 4-byte keys and then deleting them all
+// Repeat in ascending, descending and random key order
+void InsertThenDelete() {
+
+  std::cout << "==Insert then delete==" << std::endl;
+
+  size_t num_keys = 512e6;
+
+  std::cout << num_keys << " keys" << std::endl;
+
+  std::cout << "Order\t\tInsertion (ms)\t\tDeletion (ms)" << std::endl;
+
+  // Allocate GPU memory to store the keys that will be inserted
+  Pair<uint32_t, uint32_t> *d_elements;
+  size_t num_bytes = num_keys * sizeof(Pair<uint32_t, uint32_t>);
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_elements, num_bytes));
+
+  priority_queue<uint32_t, uint32_t> pq(num_keys);
+
+  // Ascending
+  std::vector<Pair<uint32_t, uint32_t>> ascending(num_keys);
+
+  for (uint32_t i = 0; i < num_keys; i++) {
+    ascending[i] = {i, i};
+  }
+
+  CUCO_CUDA_TRY(cudaMemcpy(d_elements, &ascending[0],
+                      num_bytes, cudaMemcpyHostToDevice));
+
+  auto time_elapsed_insert = TimeInsert(pq, d_elements, num_keys);
+  auto time_elapsed_delete = TimeDelete(pq, d_elements, num_keys);
+
+  std::cout << "Ascend\t\t" << time_elapsed_insert << "\t\t"
+                               << time_elapsed_delete << std::endl;
+
+  // Descending
+  std::vector<Pair<uint32_t, uint32_t>> descending(num_keys);
+
+  for (uint32_t i = 0; i < num_keys; i++) {
+    descending[num_keys - i - 1] = {i, i};
+  }
+
+  CUCO_CUDA_TRY(cudaMemcpy(d_elements, &descending[0],
+                      num_bytes, cudaMemcpyHostToDevice));
+
+  time_elapsed_insert = TimeInsert(pq, d_elements, num_keys);
+  time_elapsed_delete = TimeDelete(pq, d_elements, num_keys);
+
+  std::cout << "Descend\t\t" << time_elapsed_insert << "\t\t"
+                               << time_elapsed_delete << std::endl;
+
+  // Random
+  std::vector<Pair<uint32_t, uint32_t>> random(num_keys);
+
+  for (uint32_t i = 0; i < num_keys; i++) {
+    random[i] = {(uint32_t)rand(), i};
+  }
+
+  CUCO_CUDA_TRY(cudaMemcpy(d_elements, &random[0],
+                      num_bytes, cudaMemcpyHostToDevice));
+
+  time_elapsed_insert = TimeInsert(pq, d_elements, num_keys);
+  time_elapsed_delete = TimeDelete(pq, d_elements, num_keys);
+
+  std::cout << "Random\t\t" << time_elapsed_insert << "\t\t"
+                               << time_elapsed_delete << std::endl;
+
+  CUCO_CUDA_TRY(cudaMemcpy(d_elements, &random[0],
+                      num_bytes, cudaMemcpyHostToDevice));
+
+  time_elapsed_insert = TimeInsertDeviceAPI(pq, d_elements, num_keys);
+  time_elapsed_delete = TimeDeleteDeviceAPI(pq, d_elements, num_keys);
+
+  std::cout << "Random Dev. API\t\t" << time_elapsed_insert << "\t\t"
+                               << time_elapsed_delete << std::endl;
+
+  CUCO_CUDA_TRY(cudaFree(d_elements));
+}
+
+
+int main() {
+
+  InsertThenDelete();
+
+  return 0;
+}

include/cuco/detail/pq_pair.cuh

@@ -0,0 +1,21 @@
+#pragma once
+
+namespace cuco {
+
+template <typename Key, typename Value>
+struct Pair {
+  Key key;
+  Value value;
+};
+
+/*
+* Check if two Pairs have the same key and value
+* @param a The first pair
+* @param b The second pair
+*/
+template <typename Key, typename Value>
+bool operator==(const Pair<Key, Value> &a, const Pair<Key, Value> &b) {
+  return a.key == b.key && a.value == b.value;
+}
+
+}

include/cuco/detail/priority_queue.inl

@@ -0,0 +1,158 @@
+#pragma once
+#include <cmath>
+
+#include <cuco/detail/priority_queue_kernels.cuh>
+#include <cuco/detail/error.hpp>
+
+namespace cuco {
+
+template <typename Key, typename Value, bool Max>
+priority_queue<Key, Value, Max>::priority_queue(size_t initial_capacity,
+                                                size_t node_size) {
+
+  node_size_ = node_size;
+
+  // Round up to the nearest multiple of node size
+  int nodes = ((initial_capacity + node_size_ - 1) / node_size_);
+
+  node_capacity_ = nodes;
+  lowest_level_start_ = 1 << (int)log2(nodes);
+
+  // Allocate device variables
+
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_size_, sizeof(int)));
+
+  CUCO_CUDA_TRY(cudaMemset(d_size_, 0, sizeof(int)));
+
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_p_buffer_size_, sizeof(size_t)));
+
+  CUCO_CUDA_TRY(cudaMemset(d_p_buffer_size_, 0, sizeof(size_t)));
+
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_heap_,
+                          sizeof(Pair<Key, Value>)
+                          * (node_capacity_ * node_size_ + node_size_)));
+
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_locks_,
+             sizeof(int) * (node_capacity_ + 1)));
+
+  CUCO_CUDA_TRY(cudaMemset(d_locks_, 0,
+                          sizeof(int) * (node_capacity_ + 1)));
+
+  CUCO_CUDA_TRY(cudaMalloc((void**)&d_pop_tracker_, sizeof(int)));
+
+}
+
+template <typename Key, typename Value, bool Max>
+priority_queue<Key, Value, Max>::~priority_queue() {
+  CUCO_ASSERT_CUDA_SUCCESS(cudaFree(d_size_));
+  CUCO_ASSERT_CUDA_SUCCESS(cudaFree(d_p_buffer_size_));
+  CUCO_ASSERT_CUDA_SUCCESS(cudaFree(d_heap_));
+  CUCO_ASSERT_CUDA_SUCCESS(cudaFree(d_locks_));
+  CUCO_ASSERT_CUDA_SUCCESS(cudaFree(d_pop_tracker_));
+}
+
+
+template <typename Key, typename Value, bool Max>
+void priority_queue<Key, Value, Max>::push(Pair<Key, Value> *elements,
+                                           size_t num_elements,
+                                           int block_size,
+                                           int grid_size,
+                                           bool warp_level,
+                                           cudaStream_t stream) {
+
+  const int kBlockSize = block_size;
+  const int kNumBlocks = grid_size;
+
+  if (!warp_level) {
+    PushKernel<Max><<<kNumBlocks, kBlockSize,
+                 get_shmem_size(kBlockSize), stream>>>
+              (elements, num_elements, d_heap_, d_size_,
+               node_size_, d_locks_, d_p_buffer_size_, lowest_level_start_);
+  } else {
+    PushKernelWarp<Max><<<kNumBlocks, kBlockSize,
+                 get_shmem_size(32) * kBlockSize / 32, stream>>>
+              (elements, num_elements, d_heap_, d_size_,
+               node_size_, d_locks_, d_p_buffer_size_,
+               lowest_level_start_, get_shmem_size(32));
+  }
+
+  CUCO_CUDA_TRY(cudaGetLastError());
+}
+
+template <typename Key, typename Value, bool Max>
+void priority_queue<Key, Value, Max>::pop(Pair<Key, Value> *out,
+                                          size_t num_elements,
+                                          int block_size,
+                                          int grid_size,
+                                          bool warp_level,
+                                          cudaStream_t stream) {
+
+  const int kBlockSize = block_size;
+  const int kNumBlocks = grid_size;
+
+  cudaMemset(d_pop_tracker_, 0, sizeof(int));
+  if (!warp_level) {
+    PopKernel<Max><<<kNumBlocks, kBlockSize,
+                 get_shmem_size(kBlockSize), stream>>>
+             (out, num_elements, d_heap_, d_size_,
+              node_size_, d_locks_, d_p_buffer_size_,
+              d_pop_tracker_, lowest_level_start_, node_capacity_);
+  } else {
+    PopKernelWarp<Max><<<kNumBlocks, kBlockSize,
+                 get_shmem_size(32) * kBlockSize / 32, stream>>>
+             (out, num_elements, d_heap_, d_size_,
+              node_size_, d_locks_, d_p_buffer_size_,
+              d_pop_tracker_, lowest_level_start_,
+              node_capacity_, get_shmem_size(32));
+
+  }
+
+  CUCO_CUDA_TRY(cudaGetLastError());
+}
+
+template <typename Key, typename Value, bool Max>
+template <typename CG>
+__device__ void priority_queue<Key, Value, Max>::device_mutable_view::push(
+                                                  CG const& g,
+                                                  Pair<Key, Value> *elements,
+                                                  size_t num_elements,
+                                                  void *temp_storage) {
+
+  SharedMemoryLayout<Key, Value> shmem =
+       GetSharedMemoryLayout<Key, Value>((int*)temp_storage,
+                                         g.size(), node_size_);
+  if (num_elements == node_size_) {
+    PushSingleNode<Max>(g, elements, d_heap_, d_size_, node_size_,
+                   d_locks_, lowest_level_start_, shmem);
+  } else if (num_elements < node_size_) {
+    PushPartialNode<Max>(g, elements, num_elements, d_heap_,
+                         d_size_, node_size_, d_locks_,
+                         d_p_buffer_size_, lowest_level_start_, shmem);
+  }
+}
+
+template <typename Key, typename Value, bool Max>
+template <typename CG>
+__device__ void priority_queue<Key, Value, Max>::device_mutable_view::pop(
+                                                      CG const& g,
+                                                      Pair<Key, Value> *out,
+                                                      size_t num_elements,
+                                                      void *temp_storage) {
+  int pop_tracker = 0;
+
+  SharedMemoryLayout<Key, Value> shmem =
+       GetSharedMemoryLayout<Key, Value>((int*)temp_storage,
+                                         g.size(), node_size_);
+
+  if (num_elements == node_size_) {
+    PopSingleNode<Max>(g, out, d_heap_, d_size_, node_size_, d_locks_,
+                  d_p_buffer_size_, &pop_tracker, lowest_level_start_,
+                  node_capacity_, shmem);
+  } else {
+    PopPartialNode<Max>(g, out, num_elements, d_heap_, d_size_, node_size_,
+                   d_locks_, d_p_buffer_size_, lowest_level_start_,
+                   node_capacity_, shmem);
+  }
+}
+
+}