Refactor: separate the roles of Decomposition & FFTLayout

Author: ahojukka5

include/openpfc/core/decomposition.hpp

@@ -143,7 +143,7 @@ using Int3 = pfc::types::Int3;
 struct Decomposition {
 
   const pfc::World &m_global_world; ///< The World object.
-  const std::array<int, 3> &m_grid; ///< The number of parts in each dimension.
+  const std::array<int, 3> m_grid;  ///< The number of parts in each dimension.
   const std::vector<pfc::World> m_subworlds; ///< The sub-worlds for each part.
 
   Decomposition(const World &world, const Int3 grid)
@@ -187,7 +187,13 @@ inline const auto create(const World &world, const Int3 &grid) noexcept {
 };
 
 inline const auto create(const World &world, const int &nparts) noexcept {
-  return create(world, heffte::proc_setup_min_surface(to_indices(world), nparts));
+  auto indices = to_indices(world);
+  auto grid = heffte::proc_setup_min_surface(indices, nparts);
+  return create(world, grid);
+}
+
+inline const auto get_num_domains(const Decomposition &decomposition) noexcept {
+  return get_subworlds(decomposition).size();
 }
 
 } // namespace decomposition

include/openpfc/fft.hpp

@@ -40,20 +40,22 @@ struct FFTLayout {
  * @param decomposition The Decomposition object defining the domain
  * decomposition.
  * @param r2c_direction The direction of real-to-complex symmetry.
- * @param num_domains The number of domains for the FFT layout.
  * @return An FFTLayout object containing the layout information.
  */
-const FFTLayout create(const Decomposition &decomposition, int r2c_direction,
-                       int num_domains);
+const FFTLayout create(const Decomposition &decomposition, int r2c_direction);
 
-inline auto get_real_box(const FFTLayout &layout, int i) {
+inline const auto &get_real_box(const FFTLayout &layout, int i) {
   return layout.m_real_boxes.at(i);
 }
 
-inline auto get_complex_box(const FFTLayout &layout, int i) {
+inline const auto &get_complex_box(const FFTLayout &layout, int i) {
   return layout.m_complex_boxes.at(i);
 }
 
+inline auto get_r2c_direction(const FFTLayout &layout) {
+  return layout.m_r2c_direction;
+}
+
 } // namespace layout
 
 using pfc::types::Int3;
@@ -164,24 +166,36 @@ inline const auto get_outbox(const FFT &fft) noexcept {
   return get_fft_object(fft).outbox();
 }
 
+using heffte::plan_options;
+using layout::FFTLayout;
+
 /**
- * @brief Creates an FFT object based on the given decomposition and MPI
- * communicator.
+ * @brief Creates an FFT object based on the given FFTLayout and rank ID.
+ *
+ * @param fft_layout The FFTLayout object defining the FFT configuration.
+ * @param rank_id The rank ID of the current process in the MPI communicator.
+ * @param options Plan options for configuring the FFT behavior.
+ * @return An FFT object containing the FFT configuration and data.
+ */
+FFT create(const FFTLayout &fft_layout, int rank_id, plan_options options);
+
+/**
+ * @brief Creates an FFT object based on the given decomposition and rank ID.
  *
  * @param decomposition The Decomposition object defining the domain
  * decomposition.
- * @param comm The MPI communicator for parallel computations.
- * @param options Optional plan options for configuring the FFT behavior.
+ * @param rank_id The rank ID of the current process in the MPI communicator.
  * @return An FFT object containing the FFT configuration and data.
  */
-FFT create(const Decomposition &decomposition, MPI_Comm comm,
-           heffte::plan_options options);
+FFT create(const Decomposition &decomposition, int rank_id);
+
 /**
  * @brief Creates an FFT object based on the given decomposition.
  *
  * @param decomposition The Decomposition object defining the domain
  * decomposition.
  * @return An FFT object containing the FFT configuration and data.
+ * @throws std::logic_error, if decomposition size and rank size do not match.
  */
 FFT create(const Decomposition &decomposition);
 

include/openpfc/ui.hpp

@@ -652,10 +652,11 @@ template <class ConcreteModel> class App {
     std::cout << "World: " << world << std::endl;
 
     int num_ranks = m_worker.get_num_ranks();
+    int rank_id = m_worker.get_rank();
     auto decomp = decomposition::create(world, num_ranks);
-    auto plan_options =
-        ui::from_json<heffte::plan_options>(m_settings["plan_options"]);
-    auto fft = fft::create(decomp, m_comm, plan_options);
+    auto options = ui::from_json<heffte::plan_options>(m_settings["plan_options"]);
+    auto fft_layout = fft::layout::create(decomp, 0);
+    auto fft = fft::create(fft_layout, rank_id, options);
     Time time(ui::from_json<Time>(m_settings));
     ConcreteModel model(world);
     model.set_fft(fft);

src/openpfc/fft.cpp

@@ -7,29 +7,53 @@ namespace pfc {
 namespace fft {
 namespace layout {
 
-const FFTLayout create(const Decomposition &decomposition, int r2c_direction,
-                       int num_domains) {
-  if (num_domains <= 0) {
-    throw std::logic_error("Cannot construct domain decomposition: !(nprocs > 0)");
+#include <array>
+#include <iostream>
+#include <sstream>
+
+// Helper function to print std::array
+template <typename T, std::size_t N>
+std::ostream &operator<<(std::ostream &os, const std::array<T, N> &arr) {
+  os << "{";
+  for (std::size_t i = 0; i < N; ++i) {
+    os << arr[i];
+    if (i < N - 1) {
+      os << ", ";
+    }
   }
-  auto world = get_world(decomposition);
-  auto [N1r, N2r, N3r] = get_size(world); // real size
-  auto [N1c, N2c, N3c] = get_size(world); // complex size
+  os << "}";
+  return os;
+}
+
+using heffte::split_world;
+
+const auto get_real_indices(const Decomposition &decomposition) {
+  auto world = get_global_world(decomposition);
+  auto [N1, N2, N3] = get_size(world);
+  return heffte::box3d<int>({0, 0, 0}, {N1 - 1, N2 - 1, N3 - 1});
+}
+
+const auto get_complex_indices(const Decomposition &decomposition,
+                               int r2c_direction) {
+  auto [N1, N2, N3] = get_size(get_global_world(decomposition));
   if (r2c_direction == 0) {
-    N1c = N1c / 2 + 1;
+    return heffte::box3d<int>({0, 0, 0}, {N1 / 2, N2 - 1, N3 - 1});
   } else if (r2c_direction == 1) {
-    N2c = N2c / 2 + 1;
+    return heffte::box3d<int>({0, 0, 0}, {N1 - 1, N2 / 2, N3 - 1});
   } else if (r2c_direction == 2) {
-    N3c = N3c / 2 + 1;
+    return heffte::box3d<int>({0, 0, 0}, {N1 - 1, N2 - 1, N3 / 2});
   } else {
     throw std::logic_error("Invalid r2c_direction: " +
                            std::to_string(r2c_direction));
   }
-  box3di real_indices({0, 0, 0}, {N1r - 1, N2r - 1, N3r - 1});
-  box3di complex_indices({0, 0, 0}, {N1c - 1, N2c - 1, N3c - 1});
-  Int3 grid = heffte::proc_setup_min_surface(real_indices, num_domains);
-  std::vector<box3di> real_boxes = heffte::split_world(real_indices, grid);
-  std::vector<box3di> complex_boxes = heffte::split_world(complex_indices, grid);
+}
+
+const FFTLayout create(const Decomposition &decomposition, int r2c_direction) {
+  auto real_indices = get_real_indices(decomposition);
+  auto complex_indices = get_complex_indices(decomposition, r2c_direction);
+  auto grid = get_grid(decomposition);
+  auto real_boxes = split_world(real_indices, grid);
+  auto complex_boxes = split_world(complex_indices, grid);
   return FFTLayout{decomposition, r2c_direction, real_boxes, complex_boxes};
 }
 
@@ -49,25 +73,42 @@ int get_mpi_size(MPI_Comm comm) {
   return size;
 }
 
-FFT create(const Decomposition &decomposition, MPI_Comm comm,
-           heffte::plan_options options) {
-  int rank = get_mpi_rank(comm);
-  int mpi_num_ranks = get_mpi_size(comm);
-  if (mpi_num_ranks <= 0) {
-    throw std::logic_error("Cannot construct domain decomposition: !(nprocs > 0)");
-  }
-  auto r2c_dir = 0;
-  auto fft_layout = fft::layout::create(decomposition, r2c_dir, mpi_num_ranks);
-  auto inbox = get_real_box(fft_layout, rank);
-  auto outbox = get_complex_box(fft_layout, rank);
-  using fft_r2c = heffte::fft3d_r2c<heffte::backend::fftw>;
+using heffte::plan_options;
+using layout::FFTLayout;
+using fft_r2c = heffte::fft3d_r2c<heffte::backend::fftw>;
+
+FFT create(const FFTLayout &fft_layout, int rank_id, plan_options options) {
+  auto inbox = get_real_box(fft_layout, rank_id);
+  auto outbox = get_complex_box(fft_layout, rank_id);
+  auto r2c_dir = get_r2c_direction(fft_layout);
+  auto comm = get_comm();
   return FFT(fft_r2c(inbox, outbox, r2c_dir, comm, options));
 }
 
+FFT create(const Decomposition &decomposition, int rank_id) {
+  auto options = heffte::default_options<heffte::backend::fftw>();
+  auto r2c_dir = 0;
+  auto fft_layout = layout::create(decomposition, r2c_dir);
+  return create(fft_layout, rank_id, options);
+}
+
 FFT create(const Decomposition &decomposition) {
   auto comm = get_comm();
-  auto options = heffte::default_options<heffte::backend::fftw>();
-  return create(decomposition, comm, options);
+  auto mpi_comm_size = get_mpi_size(comm);
+  auto rank_id = get_mpi_rank(comm);
+  auto decomposition_size = get_num_domains(decomposition);
+  if (mpi_comm_size != decomposition_size) {
+    throw std::logic_error(
+        "Mismatch between MPI communicator size and domain decomposition size: " +
+        std::to_string(mpi_comm_size) + " != " + std::to_string(decomposition_size) +
+        ". This indicates that the number of MPI ranks does not match the number of "
+        "domains in the decomposition. To resolve this issue, you can manually "
+        "specify the rank by calling fft::create(decomposition, rank_id) instead.");
+  }
+  // if mpi communicator size matches decomposition size, we can safely assume
+  // that the intention is to decompose the whole communicator into the
+  // decomposition
+  return create(decomposition, rank_id);
 }
 
 } // namespace fft