Issue 1982 - Phase 2, StiffnessAssembler

core/specfem/linear_system/dof_map.hpp

@@ -0,0 +1,165 @@
+#pragma once
+
+#ifdef SPECFEM_ENABLE_TRILINOS
+
+#include "specfem/assembly/assembly.hpp"
+#include "specfem/element.hpp"
+#include "specfem/enums.hpp"
+#include "specfem/setup.hpp"
+#include <Teuchos_Comm.hpp>
+#include <Teuchos_RCP.hpp>
+#include <Tpetra_Core.hpp>
+#include <Tpetra_Map.hpp>
+#include <stdexcept>
+#include <string>
+
+namespace specfem {
+namespace linear_system {
+
+/**
+ * @brief Scalar used for assembled matrices and vectors.
+ *
+ * Aliased to `type_real` so the linear system follows the build's precision:
+ * `float` by default (matching the float-only Tpetra installs on the
+ * clusters), `double` with `SPECFEM_ENABLE_DOUBLE_PRECISION`. Single switch
+ * point if the two ever need to diverge.
+ */
+using scalar_type = type_real;
+
+/// Tpetra map with the default local/global ordinals and node type
+using map_type = Tpetra::Map<>;
+
+/// Global ordinal used for degree-of-freedom ids
+using global_ordinal_type = map_type::global_ordinal_type;
+
+/**
+ * @brief Maps SPECFEM++ (iglob, icomp) degrees of freedom of one medium to
+ * Tpetra global ids and owns the row (owned) and column (overlap) maps.
+ *
+ * The GID layout is component-blocked: `gid = icomp * nglob + iglob`,
+ * deliberately matching SPECFEM++ field storage
+ * (`Kokkos::View<type_real **, Kokkos::LayoutLeft>` of shape
+ * `(nglob, ncomp)`), so a Tpetra vector maps 1:1 onto field memory with no
+ * permutation at solve time. It also matches the element-local ordering
+ * `specfem::linear_system::local_dof_index`. The layout lives ONLY in
+ * @ref gid -- change it there to swap the ordering.
+ *
+ * One DofMap instance describes one medium: `nglob` and `ncomp` are
+ * per-medium quantities (a future multi-medium system holds one DofMap and
+ * one matrix block per medium).
+ *
+ * Serial-only in this milestone: SPECFEM++ numbers global points per rank
+ * with no globally consistent numbering across ranks, so the constructor
+ * throws for communicators with more than one rank. The owned and overlap
+ * maps are kept separate in the API so distributed assembly (owned iglob ->
+ * owned GIDs, shared-interface points -> overlap map, Export(ADD) into the
+ * owned matrix) can be added without changing callers; at one rank they are
+ * the same map.
+ */
+class DofMap {
+public:
+  /**
+   * @brief Construct the map for `nglob` mesh points with `ncomp` components
+   * per point.
+   *
+   * @param nglob Number of unique global mesh points of the medium
+   * @param ncomp Number of field components per point (3 for 3D elastic)
+   * @param comm Teuchos communicator; must have exactly one rank
+   */
+  DofMap(const int nglob, const int ncomp,
+         const Teuchos::RCP<const Teuchos::Comm<int>> &comm)
+      : nglob_(nglob), ncomp_(ncomp), comm_(comm) {
+    if (comm_->getSize() > 1) {
+      throw std::runtime_error(
+          "specfem::linear_system::DofMap: distributed assembly on " +
+          std::to_string(comm_->getSize()) +
+          " ranks is not implemented yet. SPECFEM++ numbers global points "
+          "per rank, and the cross-rank GID negotiation is a follow-up of "
+          "issue #1982. Run on a single rank.");
+    }
+    const Tpetra::global_size_t num_entries =
+        static_cast<Tpetra::global_size_t>(num_global_dofs());
+    const global_ordinal_type index_base = 0;
+    owned_map_ = Teuchos::rcp(new map_type(num_entries, index_base, comm_));
+    // At one rank every dof is owned; a distributed build replaces this with
+    // a map that additionally holds the shared-interface dofs of neighbor
+    // ranks.
+    overlap_map_ = owned_map_;
+  }
+
+  /**
+   * @brief Build the DofMap of one medium from an assembled 3D simulation.
+   *
+   * Reads `nglob` from the forward simulation field and the component count
+   * from the element attributes, using `Tpetra::getDefaultComm()`.
+   *
+   * @tparam MediumTag Medium whose degrees of freedom the map describes
+   * @param assembly Assembly with constructed fields
+   * @return DofMap for the medium
+   */
+  template <specfem::element::medium_tag MediumTag>
+  static DofMap from_assembly(
+      const specfem::assembly::assembly<specfem::element::dimension_tag::dim3>
+          &assembly) {
+    const auto &field =
+        assembly.fields
+            .get_simulation_field<specfem::simulation::field_type::forward>();
+    const int nglob = field.get_nglob<MediumTag>();
+    constexpr int ncomp =
+        specfem::element::attributes<specfem::element::dimension_tag::dim3,
+                                     MediumTag>::components;
+    return DofMap(nglob, ncomp, Tpetra::getDefaultComm());
+  }
+
+  /**
+   * @brief Global id of component `icomp` at mesh point `iglob`.
+   *
+   * Component-blocked layout `gid = icomp * nglob + iglob` -- the single
+   * source of truth for the global dof ordering (see the class docs).
+   *
+   * @param iglob Per-medium global point index in `[0, nglob())`
+   * @param icomp Field component in `[0, ncomp())`
+   * @return Tpetra global dof id in `[0, num_global_dofs())`
+   */
+  inline global_ordinal_type gid(const int iglob, const int icomp) const {
+    return static_cast<global_ordinal_type>(icomp) * nglob_ + iglob;
+  }
+
+  /// Number of unique global mesh points of the medium
+  inline int nglob() const { return nglob_; }
+
+  /// Number of field components per mesh point
+  inline int ncomp() const { return ncomp_; }
+
+  /// Total number of degrees of freedom: `ncomp * nglob`
+  inline global_ordinal_type num_global_dofs() const {
+    return static_cast<global_ordinal_type>(ncomp_) * nglob_;
+  }
+
+  /// Uniquely-owned row map: contiguous `[0, num_global_dofs())`
+  inline Teuchos::RCP<const map_type> owned_map() const { return owned_map_; }
+
+  /**
+   * @brief Overlap (column) map: owned dofs plus, in a future distributed
+   * build, the shared-interface dofs of neighbor ranks. Equal to
+   * @ref owned_map at one rank.
+   */
+  inline Teuchos::RCP<const map_type> overlap_map() const {
+    return overlap_map_;
+  }
+
+  /// Communicator the maps are defined on
+  inline Teuchos::RCP<const Teuchos::Comm<int>> comm() const { return comm_; }
+
+private:
+  int nglob_ = 0;                               ///< Points of the medium
+  int ncomp_ = 0;                               ///< Components per point
+  Teuchos::RCP<const Teuchos::Comm<int>> comm_; ///< Communicator
+  Teuchos::RCP<const map_type> owned_map_;      ///< Uniquely-owned rows
+  Teuchos::RCP<const map_type> overlap_map_;    ///< Owned + shared dofs
+};
+
+} // namespace linear_system
+} // namespace specfem
+
+#endif // SPECFEM_ENABLE_TRILINOS

core/specfem/linear_system/tpetra_assembler.cpp

@@ -0,0 +1,231 @@
+#include "specfem/linear_system/tpetra_assembler.hpp"
+
+#ifdef SPECFEM_ENABLE_TRILINOS
+
+#include "specfem/assembly/assembly.hpp"
+#include "specfem/datatype/element_index_range.hpp"
+#include "specfem/element/to_string.hpp"
+#include "specfem/linear_system/element_stiffness.hpp"
+#include "specfem/tags.hpp"
+#include <Kokkos_Core.hpp>
+#include <Teuchos_ArrayView.hpp>
+#include <algorithm>
+#include <cstddef>
+#include <sstream>
+#include <stdexcept>
+#include <vector>
+
+template <typename Tags>
+specfem::linear_system::StiffnessAssembler<Tags>::StiffnessAssembler(
+    const AssemblyType &assembly, const int batch_size)
+    : assembly_(assembly),
+      dof_map_(DofMap::from_assembly<Tags::medium_tag>(assembly)),
+      batch_size_(batch_size) {
+
+  if (batch_size_ < 1) {
+    throw std::runtime_error(
+        "specfem::linear_system::StiffnessAssembler: batch_size must be at "
+        "least 1.");
+  }
+
+  specfem::linear_system::validate_stiffness_scope<Tags>(assembly_);
+
+  // Single-medium milestone: matrix blocks coupling different media
+  // (fluid-solid) are deferred, so reject mixed meshes outright rather than
+  // silently assembling an operator that ignores the coupling.
+  const int nspec = assembly_.element_types.nspec;
+  for (int ispec = 0; ispec < nspec; ++ispec) {
+    const auto element_medium = assembly_.element_types.get_medium_tag(ispec);
+    if (element_medium != medium_tag) {
+      std::ostringstream message;
+      message << "specfem::linear_system::StiffnessAssembler: element " << ispec
+              << " has medium '" << specfem::element::to_string(element_medium)
+              << "'; only single-medium '"
+              << specfem::element::to_string(medium_tag)
+              << "' meshes are supported. Fluid-solid coupling blocks are "
+                 "deferred (issue #1982).";
+      throw std::runtime_error(message.str());
+    }
+  }
+}
+
+template <typename Tags>
+std::vector<specfem::linear_system::global_ordinal_type>
+specfem::linear_system::StiffnessAssembler<Tags>::element_column_gids(
+    const int ispec) const {
+  const auto &field = assembly_.fields.template get_simulation_field<
+      specfem::simulation::field_type::forward>();
+  const int ngllz = field.ngllz;
+  const int nglly = field.nglly;
+  const int ngllx = field.ngllx;
+  const int npoints = ngllz * nglly * ngllx;
+
+  std::vector<global_ordinal_type> cols(static_cast<std::size_t>(ncomp) *
+                                        npoints);
+  for (int iz = 0; iz < ngllz; ++iz) {
+    for (int iy = 0; iy < nglly; ++iy) {
+      for (int ix = 0; ix < ngllx; ++ix) {
+        const int iglob =
+            field.template get_iglob<false, medium_tag>(ispec, iz, iy, ix);
+        // Same ordering as local_dof_index<NGLL> on the cubic GLL grid.
+        const int point = (iz * nglly + iy) * ngllx + ix;
+        for (int icomp = 0; icomp < ncomp; ++icomp) {
+          cols[static_cast<std::size_t>(icomp) * npoints + point] =
+              dof_map_.gid(iglob, icomp);
+        }
+      }
+    }
+  }
+  return cols;
+}
+
+template <typename Tags>
+Teuchos::RCP<const specfem::linear_system::crs_graph_type>
+specfem::linear_system::StiffnessAssembler<Tags>::build_graph() const {
+  const auto &field = assembly_.fields.template get_simulation_field<
+      specfem::simulation::field_type::forward>();
+  const int ngllz = field.ngllz;
+  const int nglly = field.nglly;
+  const int ngllx = field.ngllx;
+  const int npoints = ngllz * nglly * ngllx;
+  const int ndof_e = ncomp * npoints;
+
+  const auto elements =
+      assembly_.element_types.get_elements_on_host(medium_tag);
+
+  // Pass 1: per-row allocation upper bound. A row interacts with every dof
+  // of every element sharing its mesh point; per-element duplicates are
+  // inserted raw in pass 2 (fillComplete merges them), so the bound must --
+  // and does exactly -- cover the raw insert count.
+  std::vector<std::size_t> adjacent(dof_map_.nglob(), 0);
+  for (int i = 0; i < elements.size(); ++i) {
+    const int ispec = elements(i);
+    for (int iz = 0; iz < ngllz; ++iz) {
+      for (int iy = 0; iy < nglly; ++iy) {
+        for (int ix = 0; ix < ngllx; ++ix) {
+          ++adjacent[field.template get_iglob<false, medium_tag>(ispec, iz, iy,
+                                                                 ix)];
+        }
+      }
+    }
+  }
+
+  std::vector<std::size_t> entries_per_row(
+      static_cast<std::size_t>(dof_map_.num_global_dofs()), 0);
+  for (int iglob = 0; iglob < dof_map_.nglob(); ++iglob) {
+    for (int icomp = 0; icomp < ncomp; ++icomp) {
+      entries_per_row[static_cast<std::size_t>(dof_map_.gid(iglob, icomp))] =
+          adjacent[iglob] * static_cast<std::size_t>(ndof_e);
+    }
+  }
+
+  auto graph = Teuchos::rcp(
+      new crs_graph_type(dof_map_.overlap_map(),
+                         Teuchos::ArrayView<const std::size_t>(
+                             entries_per_row.data(), entries_per_row.size())));
+
+  // Pass 2: every dof of an element couples to every other dof of that
+  // element, so each element contributes its full column list to each of its
+  // rows.
+  for (int i = 0; i < elements.size(); ++i) {
+    const auto cols = element_column_gids(elements(i));
+    for (int r = 0; r < ndof_e; ++r) {
+      graph->insertGlobalIndices(cols[r], ndof_e, cols.data());
+    }
+  }
+
+  graph->fillComplete(dof_map_.owned_map(), dof_map_.owned_map());
+  return graph;
+}
+
+template <typename Tags>
+void specfem::linear_system::StiffnessAssembler<Tags>::fill_matrix(
+    crs_matrix_type &matrix) const {
+  const auto &field = assembly_.fields.template get_simulation_field<
+      specfem::simulation::field_type::forward>();
+  const int npoints = field.ngllz * field.nglly * field.ngllx;
+  const int ndof_e = ncomp * npoints;
+
+  const auto elements =
+      assembly_.element_types.get_elements_on_host(medium_tag);
+  const int nelements = elements.size();
+
+  // One block buffer reused across batches; only batch_size_ dense element
+  // blocks ever exist at a time -- no global dense matrix.
+  Kokkos::View<type_real ***, Kokkos::LayoutRight,
+               Kokkos::DefaultExecutionSpace>
+      k_e("specfem::linear_system::element_stiffness_blocks",
+          std::min(batch_size_, std::max(nelements, 1)), ndof_e, ndof_e);
+  auto h_k_e = Kokkos::create_mirror_view(k_e);
+
+  for (int offset = 0; offset < nelements; offset += batch_size_) {
+    const int batch_count = std::min(batch_size_, nelements - offset);
+    const auto batch = specfem::datatype::subview(
+        elements, Kokkos::pair<int, int>(offset, offset + batch_count));
+
+    specfem::linear_system::compute_element_stiffness<Tags>(assembly_, batch,
+                                                            k_e);
+    Kokkos::deep_copy(h_k_e, k_e);
+
+    for (int e = 0; e < batch_count; ++e) {
+      const int ispec = batch(e);
+      const auto cols = element_column_gids(ispec);
+      for (int r = 0; r < ndof_e; ++r) {
+        const int updated = matrix.sumIntoGlobalValues(
+            cols[r], ndof_e, &h_k_e(e, r, 0), cols.data());
+        if (updated != ndof_e) {
+          std::ostringstream message;
+          message << "specfem::linear_system::StiffnessAssembler: row "
+                     "scatter of element "
+                  << ispec << " updated " << updated << " of " << ndof_e
+                  << " entries; the graph does not match the element "
+                     "connectivity.";
+          throw std::runtime_error(message.str());
+        }
+      }
+    }
+  }
+}
+
+template <typename Tags>
+Teuchos::RCP<specfem::linear_system::crs_matrix_type>
+specfem::linear_system::StiffnessAssembler<Tags>::assemble() const {
+  const auto graph = build_graph();
+
+  // Matrix on the fill-complete static graph: values start at zero and only
+  // sumIntoGlobalValues is allowed, which is exactly what the scatter uses.
+  auto matrix = Teuchos::rcp(new crs_matrix_type(graph));
+
+  fill_matrix(*matrix);
+
+  matrix->fillComplete(dof_map_.owned_map(), dof_map_.owned_map());
+
+  // MPI-ready seam: a distributed build fills the matrix on the overlap map
+  // and Export(ADD)s into the owned map here, reproducing the matrix-free
+  // assembly sum across ranks. Unreachable at one rank (DofMap rejects
+  // larger communicators), where overlap == owned and the matrix is final.
+  if (!dof_map_.overlap_map()->isSameAs(*dof_map_.owned_map())) {
+    throw std::runtime_error(
+        "specfem::linear_system::StiffnessAssembler: overlap map differs "
+        "from owned map, but the distributed Export(ADD) step is not "
+        "implemented yet (follow-up of issue #1982).");
+  }
+
+  return matrix;
+}
+
+namespace specfem::linear_system_impl {
+/// Tag bundle for the only combination explicitly instantiated for the
+/// linear system (issue #1982).
+using elastic_isotropic_tags =
+    specfem::tags::Tags<specfem::element::dimension_tag::dim3,
+                        specfem::element::medium_tag::elastic,
+                        specfem::element::property_tag::isotropic,
+                        specfem::element::attenuation_tag::none>;
+} // namespace specfem::linear_system_impl
+
+// Explicit instantiation: 3D elastic isotropic
+template class specfem::linear_system::StiffnessAssembler<
+    specfem::linear_system_impl::elastic_isotropic_tags>;
+
+#endif // SPECFEM_ENABLE_TRILINOS