feat(kernels): use binary search to find compartment

Author: Benncs

apps/libs/simulation/public/simulation/kernels/move_kernel.hpp

@@ -27,6 +27,39 @@ namespace Simulation::KernelInline
   constexpr bool disable_move = false;
   constexpr bool enable_move = true;
 
+  // KOKKOS_INLINE_FUNCTION std::size_t __find_next_compartment(
+  //     const ConstNeighborsView<ComputeSpace>& neighbors,
+  //     const CumulativeProbabilityView<ComputeSpace>& cumulative_probability,
+  //     const std::size_t i_compartment,
+  //     const double random_number)
+  // {
+  //   const int max_neighbor = static_cast<int>(neighbors.extent(1));
+
+  //   std::size_t next =
+  //       neighbors(i_compartment, 0); // Default to the first neighbor
+
+  //   // Iterate through the neighbors to find the appropriate next compartment
+  //   for (int k_neighbor = 0; k_neighbor < max_neighbor - 1; ++k_neighbor)
+  //   {
+
+  //     // Get the cumulative probability range for the current neighbor
+  //     const auto pi = cumulative_probability(i_compartment, k_neighbor);
+  //     const auto pn = cumulative_probability(i_compartment, k_neighbor + 1);
+
+  //     // Use of a Condition mask to avoid branching.
+  //     next = (random_number <= pn && pi <= random_number)
+  //                ? neighbors(i_compartment, k_neighbor + 1)
+  //                : next;
+  //   }
+
+  //   return next; // Return the index of the chosen next compartment
+  // }
+
+  /** @brief probably overkill binary search to find next compartment
+
+  Compared with first impl it might not change anything
+  Binary seach  is O(log(n)) vs first linear is (n)
+  */
   KOKKOS_INLINE_FUNCTION std::size_t __find_next_compartment(
       const ConstNeighborsView<ComputeSpace>& neighbors,
       const CumulativeProbabilityView<ComputeSpace>& cumulative_probability,
@@ -35,24 +68,19 @@ namespace Simulation::KernelInline
   {
     const int max_neighbor = static_cast<int>(neighbors.extent(1));
 
-    std::size_t next =
-        neighbors(i_compartment, 0); // Default to the first neighbor
+    int left = 0;
+    int right = max_neighbor - 1;
 
-    // Iterate through the neighbors to find the appropriate next compartment
-    for (int k_neighbor = 0; k_neighbor < max_neighbor - 1; ++k_neighbor)
+    while (left < right)
     {
-
-      // Get the cumulative probability range for the current neighbor
-      const auto pi = cumulative_probability(i_compartment, k_neighbor);
-      const auto pn = cumulative_probability(i_compartment, k_neighbor + 1);
-
-      // Use of a Condition mask to avoid branching.
-      next = (random_number <= pn && pi <= random_number)
-                 ? neighbors(i_compartment, k_neighbor + 1)
-                 : next;
+      int mid = (left + right) >> 1; // NOLINT
+      const auto pm = cumulative_probability(i_compartment, mid);
+      const int mask = static_cast<int>(random_number > pm);
+      left = mask * (mid + 1) + (1 - mask) * left;
+      right = mask * right + (1 - mask) * mid;
     }
 
-    return next; // Return the index of the chosen next compartment
+    return neighbors(i_compartment, left);
   }
 
   struct TagMove
@@ -206,10 +234,10 @@ namespace Simulation::KernelInline
       KOKKOS_ASSERT(i_current_compartment < move.liquid_volume.extent(0));
 
       const bool mask_next =
-          probability_leaving<int>(rng1,
-                                   move.liquid_volume(i_current_compartment),
-                                   move.diag_transition(i_current_compartment),
-                                   d_t);
+          probability_leaving<void>(rng1,
+                                    move.liquid_volume(i_current_compartment),
+                                    move.diag_transition(i_current_compartment),
+                                    d_t);
 
       positions(idx) =
           (mask_next) ? __find_next_compartment(move.neighbors,
@@ -229,52 +257,62 @@ namespace Simulation::KernelInline
       }
     }
 
-    KOKKOS_INLINE_FUNCTION void handle_exit(std::size_t idx,
-                                            std::size_t& dead_count) const
+    KOKKOS_INLINE_FUNCTION void
+    inner_handle_exit(const std::size_t i_flow,
+                      const std::size_t idx,
+                      std ::size_t& dead_count,
+                      const double liquid_volume,
+                      const std::size_t position) const
     {
-      const auto position = positions(idx);
-      const std::size_t n_flow = move.leaving_flow.size();
-      const auto liquid_volume = move.liquid_volume(position);
-
-      for (std::size_t i = 0LU; i < n_flow; ++i)
-      {
+      // const float random_number = random(idx, i);
 
-        // const float random_number = random(idx, i);
-        auto generator = random_pool.get_state();
-        const float random_number = generator.frand(0., 1.);
-        random_pool.free_state(generator);
+      SAMPLE_RANDOM_VARIABLES(random_pool,
+                              const float random_number =
+                                  _generator_state_.frand(0., 1.);)
 
-        const auto& [index, flow] = move.leaving_flow(i);
+      const auto& [index, flow] = move.leaving_flow(i_flow);
 
-        const bool is_leaving =
-            (position == index) &&
-            probability_leaving(random_number, liquid_volume, flow, d_t);
+      const bool is_leaving =
+          (position == index) &&
+          probability_leaving<void>(random_number, liquid_volume, flow, d_t);
 
-        const int leave_mask = static_cast<int>(is_leaving);
+      const int leave_mask = static_cast<int>(is_leaving);
 
-        // If using probes
-        if constexpr (AutoGenerated::FlagCompileTime::use_probe)
+      // If using probes
+      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
+      {
+        // Execute probe set, but only actually do something if leaving
+        const auto probe_value = ages(idx, 0);
+        if (is_leaving && !probes.set(probe_value))
         {
-          // Execute probe set, but only actually do something if leaving
-          const auto probe_value = ages(idx, 0);
-          if (is_leaving && !probes.set(probe_value))
-          {
-            Kokkos::printf("[Kernel]: PROBES OVERFLOW\r\n");
-          }
+          Kokkos::printf("[Kernel]: PROBES OVERFLOW\r\n");
         }
-        if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
+      }
+      if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
+      {
+        if (is_leaving)
         {
-          if (is_leaving)
-          {
-            events.incr<MC::EventType::Exit>();
-          }
+          events.incr<MC::EventType::Exit>();
         }
+      }
 
-        dead_count += leave_mask;
-        ages(idx, 0) = leave_mask * 0 + (1 - leave_mask) * ages(idx, 0);
-        status(idx) = is_leaving ? MC::Status::Exit : status(idx);
+      dead_count += leave_mask;
+      ages(idx, 0) = leave_mask * 0 + (1 - leave_mask) * ages(idx, 0);
+      // status(idx) = is_leaving ? MC::Status::Exit : status(idx);
+      status(idx) = static_cast<MC::Status>(
+          static_cast<int>(MC::Status::Exit) * leave_mask +
+          (1 - leave_mask) * static_cast<int>(status(idx)));
+    }
 
-        // events.wrap_incr<MC::EventType::Exit>();
+    KOKKOS_INLINE_FUNCTION void handle_exit(std::size_t idx,
+                                            std::size_t& dead_count) const
+    {
+      const auto position = positions(idx);
+      const std::size_t n_flow = move.leaving_flow.size();
+      const auto liquid_volume = move.liquid_volume(position);
+      for (std::size_t i = 0LU; i < n_flow; ++i)
+      {
+        inner_handle_exit(i, idx, dead_count, liquid_volume, position);
       }
     }