1 files changed, 29 insertions, 14 deletions

diff --git a/main.cu b/main.cu
index 898df07..5539991 100644
--- a/main.cu
+++ b/main.cu
@@ -111,6 +111,7 @@ constexpr size_t FACTOR = 1;
 
 // should change this to dynamic next time
 constexpr size_t KEY_INDEX_SIZE = 32;
+constexpr size_t SAMPLE_SIZE = 1024;
 
 constexpr int block_size = STEP_SIZE;
 
@@ -338,9 +339,8 @@ __device__ Node *GetNewNode(LL key) {
 __device__ LockFreeSkipList *l; // The lock-free skip list
 
 __device__ LL KeyIndex[KEY_INDEX_SIZE];
-#ifdef MEASURE_TIME
-__device__ int *SpendTime;
-#endif
+
+__device__ key_type SampleStorage[SAMPLE_SIZE];
 
 // Kernel for initializing device memory
 
@@ -657,22 +657,37 @@ void initialize(const std::vector<double> &input_population,
   std::sort(sample.begin(), sample.end());
 }
 
-double sample_cdf(double x) {
-  auto it = lower_bound(sample.begin(), sample.end(), x);
-  if (it == sample.end()) {
+__device__ key_type *cudaBinarySearch(key_type *start, key_type *end,
+                                      key_type val) {
+  auto begin = start;
+  key_type *last_known_point = nullptr;
+  while (begin < end) {
+    auto mid = (end - begin) / 2;
+    auto mid_val = *(start + mid);
+    if (val == mid_val) {
+      return start + mid;
+    } else if (val > mid_val) {
+      begin = begin + mid + 1;
+    } else {
+      end = end - mid - 1;
+    }
+    last_known_point = begin;
+  }
+  return last_known_point;
+}
+
+__device__ long double sample_cdf(long double x) {
+  auto it = cudaBinarySearch(SampleStorage, SampleStorage + SAMPLE_SIZE, x);
+  if (it == SampleStorage + SAMPLE_SIZE) {
     return 1;
   }
-  if (it == sample.begin()) {
+  if (it == SampleStorage) {
     return 0;
   }
   auto it_prev = it - 1;
-  return (double(it_prev - sample.begin()) +
-          double(x - *it_prev) / (*it - *it_prev)) /
-         double(sample.size() - 1);
-}
-
-int infer_offset(double x) {
-  return int(sample_cdf(x) * double(population.size()));
+  return (double(it_prev - SampleStorage) +
+          (x - (long double)*it_prev) / (long double)(*it - *it_prev)) /
+         double(SAMPLE_SIZE - 1);
 }
 
 int main(int argc, char **argv) {