Add GPU-compatible upper bound and lower bound algorithms to AMReX_Al…

…gorithm (AMReX-Codes#2958)
Thierry992 · Sep 23, 2022 · c4b7982 · c4b7982
1 parent 3e5cc77
commit c4b7982
Showing 1 changed file with 51 additions and 0 deletions.
diff --git a/Src/Base/AMReX_Algorithm.H b/Src/Base/AMReX_Algorithm.H
@@ -157,6 +157,57 @@ namespace amrex
         return hi;
     }
 
+    template<typename ItType, typename ValType>
+    AMREX_GPU_HOST_DEVICE
+    ItType upper_bound (ItType first, ItType last, const ValType& val)
+    {
+#if AMREX_DEVICE_COMPILE
+        std::ptrdiff_t count = last-first;
+        while(count>0){
+            auto it = first;
+            const auto step = count/2;
+            it += step;
+            if (!(val < *it)){
+                first = ++it;
+                count -= step + 1;
+            }
+            else{
+                count = step;
+            }
+        }
+
+        return first;
+#else
+        return std::upper_bound(first, last, val);
+#endif
+    }
+
+    template<typename ItType, typename ValType>
+    AMREX_GPU_HOST_DEVICE
+    ItType lower_bound (ItType first, ItType last, const ValType& val)
+    {
+#ifdef AMREX_DEVICE_COMPILE
+        std::ptrdiff_t count = last-first;
+        while(count>0)
+        {
+            auto it = first;
+            const auto step = count/2;
+            it += step;
+            if (*it < val){
+                first = ++it;
+                count -= step + 1;
+            }
+            else{
+                count = step;
+            }
+        }
+
+        return first;
+#else
+        return std::lower_bound(first, last, val);
+#endif
+    }
+
 namespace detail {
 
 struct clzll_tag {};