Implement linked-list/vector hybrid NeighborList

include/openmc/cell.h

@@ -11,6 +11,7 @@
 #include "pugixml.hpp"
 
 #include "openmc/constants.h"
+#include "openmc/neighbor_list.h"
 #include "openmc/position.h"
 
 #ifdef DAGMC
@@ -104,7 +105,7 @@ public:
   bool simple_;  //!< Does the region contain only intersections?
 
   //! \brief Neighboring cells in the same universe.
-  std::vector<int> neighbors;
+  NeighborList neighbors;
 
   Position translation_ {0, 0, 0}; //!< Translation vector for filled universe
 

include/openmc/geometry_aux.h

@@ -88,6 +88,15 @@ distribcell_path(int32_t target_cell, int32_t map, int32_t target_offset);
 
 extern "C" int maximum_levels(int32_t univ);
 
+//==============================================================================
+//! Perform any geometry operations that should be done between generations.
+//!
+//! Currently, this just includes converting cell neighbor linked lists into
+//! vectors, but similar optimizations could be added here in the future.
+//==============================================================================
+
+void geometry_finalize_generation();
+
 //==============================================================================
 //! Deallocates global vectors and maps for cells, universes, and lattices.
 //==============================================================================

include/openmc/neighbor_list.h

@@ -0,0 +1,165 @@
+#ifndef OPENMC_NEIGHBOR_LIST_H
+#define OPENMC_NEIGHBOR_LIST_H
+
+#include <algorithm>
+#include <cstdint>
+#include <forward_list>
+#include <mutex>
+
+#include "openmc/openmp_interface.h"
+
+namespace openmc {
+
+// Forward declare the neighbor list iterator type.
+template <typename T_value, typename T_prefix_iter, typename T_suffix_iter>
+class NeighborListIter;
+
+//==============================================================================
+//! An efficient, threadsafe, dynamic container for listing neighboring cells.
+//
+//! This container is a minor improvement upon a linked list.  The linked list
+//! allows for threadsafe dynamic growth; any number of threads can safely
+//! read from the list without locks or reference counting.  Write access must
+//! be protected with a lock, but write events are fairly rare for neighbor
+//! lists.  This container also provides the option to flush the linked list
+//! into a vector at some convenient time (like between generations where the
+//! neighbor lists are not being read) for contiguous data.  The resulting
+//! performance benefit is small but consistent.
+//==============================================================================
+
+class NeighborList
+{
+public:
+  using value_type = int32_t;
+  using prefix_iter = std::vector<value_type>::iterator;
+  using suffix_iter = std::forward_list<value_type>::iterator;
+  using iterator = NeighborListIter<value_type, prefix_iter, suffix_iter>;
+
+  void push(int new_elem)
+  {
+    // Try to acquire the lock.
+    std::unique_lock<ThreadMutex> lock(mutex_, std::try_to_lock);
+    if (lock) {
+      // Make sure this element isn't already in the suffix.  Don't check the
+      // prefix because we don't expect this function to be called if the "new"
+      // element is already there.
+      if (std::find(suffix_.cbegin(), suffix_.cend(), new_elem)
+          == suffix_.cend()) {
+        suffix_.push_front(new_elem);
+      }
+    }
+  }
+
+  void make_consecutive()
+  {
+    while (!suffix_.empty()) {
+      prefix_.push_back(suffix_.front());
+      suffix_.pop_front();
+    }
+  }
+
+  iterator begin();
+
+  iterator end();
+
+private:
+  std::vector<value_type> prefix_;
+  std::forward_list<value_type> suffix_;
+  ThreadMutex mutex_;
+
+  friend class NeighborListIter<value_type, prefix_iter, suffix_iter>;
+};
+
+//==============================================================================
+
+template <typename T_value, typename T_prefix_iter, typename T_suffix_iter>
+class NeighborListIter
+{
+public:
+  NeighborListIter(NeighborList* nl, T_prefix_iter it)
+  {
+    base_ = nl;
+    if (it != base_->prefix_.end()) {
+      in_prefix_ = true;
+      prefix_iter_ = it;
+    } else {
+      in_prefix_ = false;
+      suffix_iter_ = base_->suffix_.begin();
+    }
+  }
+
+  NeighborListIter(NeighborList* nl, T_suffix_iter it)
+  {
+    in_prefix_ = false;
+    suffix_iter_ = it;
+  }
+
+  T_value operator*()
+  {
+    if (in_prefix_) {
+      return *prefix_iter_;
+    } else {
+      return *suffix_iter_;
+    }
+  }
+
+  bool operator==(const NeighborListIter& other)
+  {
+    if (in_prefix_ != other.in_prefix_) return false;
+
+    if (in_prefix_) {
+      return prefix_iter_ == other.prefix_iter_;
+    } else {
+      return suffix_iter_ == other.suffix_iter_;
+    }
+  }
+
+  bool operator!=(const NeighborListIter& other)
+  {return !(*this == other);}
+
+  NeighborListIter& operator++()
+  {
+    if (in_prefix_) {
+      // We are in the prefix so increment the prefix iterator.
+      ++prefix_iter_;
+
+      // If we've reached the end of the prefix, switch to the suffix iterator.
+      if (prefix_iter_ == base_->prefix_.end()) {
+        in_prefix_ = false;
+        suffix_iter_ = base_->suffix_.begin();
+      }
+
+    } else {
+      // We are in the suffix so increment the suffix iterator.
+      ++suffix_iter_;
+    }
+    return *this;
+  }
+
+private:
+  NeighborList* base_;
+
+  union {
+    T_prefix_iter prefix_iter_;
+    T_suffix_iter suffix_iter_;
+  };
+
+  bool in_prefix_;
+};
+
+//==============================================================================
+
+inline NeighborList::iterator
+NeighborList::begin()
+{
+  return iterator(this, prefix_.begin());
+}
+
+inline NeighborList::iterator
+NeighborList::end()
+{
+  return iterator(this, suffix_.end());
+} 
+
+} // namespace openmc
+#endif // OPENMC_NEIGHBOR_LIST_H

include/openmc/openmp_interface.h

@@ -0,0 +1,50 @@
+#ifndef OPENMC_OPENMP_INTERFACE_H
+#define OPENMC_OPENMP_INTERFACE_H
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+namespace openmc {
+
+class ThreadMutex
+{
+
+#ifdef _OPENMP
+//==============================================================================
+// Implementation for when OpenMP is actually defined.
+
+public:
+  ThreadMutex()
+  {omp_init_lock(&mutex_);}
+
+  ~ThreadMutex()
+  {omp_destroy_lock(&mutex_);}
+
+  void lock()
+  {omp_set_lock(&mutex_);}
+
+  bool try_lock()
+  {return omp_test_lock(&mutex_);}
+
+  void unlock()
+  {omp_unset_lock(&mutex_);}
+
+private:
+  omp_lock_t mutex_;
+
+#else
+//==============================================================================
+// Empty implementation for when OpenMP is not defined.
+
+public:
+  ThreadMutex() {}
+  ~ThreadMutex() = default;
+  void lock() {}
+  bool try_lock() {return true;}
+  void unlock() {}
+#endif
+};
+
+} // namespace openmc
+#endif // OPENMC_OPENMP_INTERFACE_H

src/geometry.cpp

@@ -63,39 +63,58 @@ check_cell_overlap(Particle* p)
 //==============================================================================
 
 bool
-find_cell_inner(Particle* p, const std::vector<int>* search_cells)
+find_cell_inner(Particle* p, NeighborList* neighbor_list)
 {
-  // If a set of cells to search was not specified, search all cells in this
-  // universe.
-  if (!search_cells) {
-    int i_universe = p->coord[p->n_coord-1].universe;
-    search_cells = &model::universes[i_universe]->cells_;
-  }
-
-  // Find which cell of this universe the particle is in.
+  // Find which cell of this universe the particle is in.  Use the neighbor list
+  // to shorten the search if one was provided.
   bool found = false;
   int32_t i_cell;
-  for (int i = 0; i < search_cells->size(); i++) {
-    i_cell = (*search_cells)[i];
+  if (neighbor_list) {
+    for (auto it = neighbor_list->begin(); it != neighbor_list->end(); ++it) {
+      i_cell = *it;
 
-    // Make sure the search cell is in the same universe.
-    int i_universe = p->coord[p->n_coord-1].universe;
-    if (model::cells[i_cell]->universe_ != i_universe) continue;
+      // Make sure the search cell is in the same universe.
+      int i_universe = p->coord[p->n_coord-1].universe;
+      if (model::cells[i_cell]->universe_ != i_universe) continue;
 
-    Position r {p->coord[p->n_coord-1].xyz};
-    Direction u {p->coord[p->n_coord-1].uvw};
-    int32_t surf = p->surface;
-    if (model::cells[i_cell]->contains(r, u, surf)) {
-      p->coord[p->n_coord-1].cell = i_cell;
-
-      if (settings::verbosity >= 10 || simulation::trace) {
-        std::stringstream msg;
-        msg << "    Entering cell " << model::cells[i_cell]->id_;
-        write_message(msg, 1);
+      // Check if this cell contains the particle.
+      Position r {p->coord[p->n_coord-1].xyz};
+      Direction u {p->coord[p->n_coord-1].uvw};
+      auto surf = p->surface;
+      if (model::cells[i_cell]->contains(r, u, surf)) {
+        p->coord[p->n_coord-1].cell = i_cell;
+        found = true;
+        break;
       }
-      found = true;
-      break;
     }
+
+  } else {
+    int i_universe = p->coord[p->n_coord-1].universe;
+    const auto& cells {model::universes[i_universe]->cells_};
+    for (auto it = cells.begin(); it != cells.end(); it++) {
+      i_cell = *it;
+
+      // Make sure the search cell is in the same universe.
+      int i_universe = p->coord[p->n_coord-1].universe;
+      if (model::cells[i_cell]->universe_ != i_universe) continue;
+
+      // Check if this cell contains the particle.
+      Position r {p->coord[p->n_coord-1].xyz};
+      Direction u {p->coord[p->n_coord-1].uvw};
+      auto surf = p->surface;
+      if (model::cells[i_cell]->contains(r, u, surf)) {
+        p->coord[p->n_coord-1].cell = i_cell;
+        found = true;
+        break;
+      }
+    }
+  }
+
+  // Announce the cell that the particle is entering.
+  if (found && (settings::verbosity >= 10 || simulation::trace)) {
+    std::stringstream msg;
+    msg << "    Entering cell " << model::cells[i_cell]->id_;
+    write_message(msg, 1);
   }
 
   if (found) {
@@ -271,15 +290,14 @@ find_cell(Particle* p, bool use_neighbor_lists)
 
     // Search for the particle in that cell's neighbor list.  Return if we
     // found the particle.
-    std::vector<int>* search_cells = &c.neighbors;
-    bool found = find_cell_inner(p, search_cells);
+    bool found = find_cell_inner(p, &c.neighbors);
     if (found) return found;
 
     // The particle could not be found in the neighbor list.  Try searching all
     // cells in this universe, and update the neighbor list if we find a new
     // neighboring cell.
     found = find_cell_inner(p, nullptr);
-    if (found) c.neighbors.push_back(p->coord[coord_lvl].cell);
+    if (found) c.neighbors.push(p->coord[coord_lvl].cell);
     return found;
 
   } else {

src/geometry_aux.cpp

@@ -432,6 +432,14 @@ maximum_levels(int32_t univ)
 
 //==============================================================================
 
+void
+geometry_finalize_generation()
+{
+  for (Cell* c : model::cells) c->neighbors.make_consecutive();
+}
+
+//==============================================================================
+
 void
 free_memory_geometry_c()
 {

src/simulation.cpp

@@ -4,6 +4,7 @@
 #include "openmc/container_util.h"
 #include "openmc/eigenvalue.h"
 #include "openmc/error.h"
+#include "openmc/geometry_aux.h"
 #include "openmc/message_passing.h"
 #include "openmc/output.h"
 #include "openmc/particle.h"
@@ -449,6 +450,9 @@ void finalize_generation()
     // For fixed-source mode, we need to sample the external source
     fill_source_bank_fixedsource();
   }
+
+  // Clean up neighbor lists
+  geometry_finalize_generation();
 }
 
 void initialize_history(Particle* p, int64_t index_source)