assembly_schedule.hh

Go to the documentation of this file.

#pragma once
 
// external libraries
#include <deal.II/base/multithread_info.h>
#include <deal.II/distributed/tria_base.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/grid/filtered_iterator.h>
#include <deal.II/grid/tria.h>
 
// DiFfRG
#include <DiFfRG/common/config_tree.hh>
#include <DiFfRG/common/threads.hh>
#include <DiFfRG/common/types.hh>
 
// std
#include <algorithm>
#include <cmath>
#include <optional>
 
namespace DiFfRG
{
  struct AssemblySchedule {
    uint queue_length = 1;
    uint chunk_size = 1;
 
    friend bool operator==(const AssemblySchedule &, const AssemblySchedule &) = default;
  };
 
  namespace assembly_cost
  {
    inline constexpr double local_fe = 250.;
    inline constexpr double algebraic = 1000.;
    inline constexpr double momentum_integral = 3000.;
  } // namespace assembly_cost
 
  namespace assembly_schedule_defaults
  {
    inline constexpr double chunk_work_ns = 8.0e3;
 
    inline constexpr double worker_work_ns = 64.0e3;
 
    inline constexpr uint chunk_cap = 256;
  } // namespace assembly_schedule_defaults
 
  struct AssemblyScheduleOverrides {
    std::optional<uint> queue_length;
    std::optional<uint> chunk_size;
 
    static AssemblyScheduleOverrides from_config(const ConfigTree &config)
    {
      AssemblyScheduleOverrides overrides;
      if (config.contains("/discretization/mesh_workers"))
        overrides.queue_length = config.get_uint("/discretization/mesh_workers");
      if (config.contains("/discretization/batch_size"))
        overrides.chunk_size = config.get_uint("/discretization/batch_size");
      return overrides;
    }
  };
 
  inline AssemblySchedule make_assembly_schedule(const uint n_local_cells, const uint thread_budget,
                                                 const double cost_ns, const AssemblyScheduleOverrides &overrides = {})
  {
    using namespace assembly_schedule_defaults;
 
    const uint threads = std::max(1u, thread_budget);
    const double cost = std::max(1.0, cost_ns);
 
    uint chunk =
        static_cast<uint>(std::clamp(std::llround(chunk_work_ns / cost), 1ll, static_cast<long long>(chunk_cap)));
 
    // Allocating past the end of the mesh is pure cost: mesh_loop reserves queue_length *
    // chunk_size CopyData objects whether or not there are cells to put in them.
    if (n_local_cells > 0) chunk = std::min(chunk, n_local_cells);
    chunk = std::max(1u, chunk);
 
    const uint n_chunks = (n_local_cells + chunk - 1) / chunk;
    const auto workers_paid_for = std::llround(double(n_local_cells) * cost / worker_work_ns);
    // Below the thread count, never at it: the integrators spawn their own TBB work from inside a
    // cell worker, in the same arena, and a pipeline holding every thread starves that nesting.
    const uint w_max = threads > 1 ? threads - 1 : 1u;
 
    uint queue = static_cast<uint>(std::clamp(workers_paid_for, 1ll, static_cast<long long>(w_max)));
    if (n_chunks > 0) queue = std::min(queue, n_chunks);
 
    if (overrides.queue_length) queue = *overrides.queue_length;
    if (overrides.chunk_size) chunk = *overrides.chunk_size;
 
    // deal.II only guards these with Assert, which is compiled out in Release, and a zero here is
    // undefined behaviour inside tbb::parallel_pipeline rather than a diagnostic.
    return {std::max(1u, queue), std::max(1u, chunk)};
  }
 
  template <int dim> auto locally_owned_cells(const dealii::DoFHandler<dim> &dof_handler)
  {
    return dealii::filter_iterators(dof_handler.active_cell_iterators(),
                                    dealii::IteratorFilters::LocallyOwnedCell());
  }
 
  template <typename Discretization> uint n_locally_owned_cells(const Discretization &discretization)
  {
    const auto &tria = discretization.get_triangulation();
    if constexpr (Discretization::Mesh::is_parallel)
      return static_cast<uint>(tria.n_locally_owned_active_cells());
    else
      return static_cast<uint>(tria.n_active_cells());
  }
} // namespace DiFfRG