Write and read mesh using ADIOS2

cpp/demo/checkpointing/CMakeLists.txt

@@ -0,0 +1,40 @@
+# This file was generated by running
+#
+# python cmake/scripts/generate-cmakefiles.py from dolfinx/cpp
+#
+cmake_minimum_required(VERSION 3.19)
+
+set(PROJECT_NAME demo_checkpointing)
+project(${PROJECT_NAME} LANGUAGES C CXX)
+
+# Set C++20 standard
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+if(NOT TARGET dolfinx)
+  find_package(DOLFINX REQUIRED)
+endif()
+
+set(CMAKE_INCLUDE_CURRENT_DIR ON)
+
+add_executable(${PROJECT_NAME} main.cpp)
+target_link_libraries(${PROJECT_NAME} dolfinx)
+
+# Do not throw error for 'multi-line comments' (these are typical in rst which
+# includes LaTeX)
+include(CheckCXXCompilerFlag)
+check_cxx_compiler_flag("-Wno-comment" HAVE_NO_MULTLINE)
+set_source_files_properties(
+  main.cpp
+  PROPERTIES
+    COMPILE_FLAGS
+    "$<$<BOOL:${HAVE_NO_MULTLINE}>:-Wno-comment -Wall -Wextra -pedantic -Werror>"
+)
+
+# Test targets (used by DOLFINx testing system)
+set(TEST_PARAMETERS2 -np 2 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}")
+set(TEST_PARAMETERS3 -np 3 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}")
+add_test(NAME ${PROJECT_NAME}_mpi_2 COMMAND "mpirun" ${TEST_PARAMETERS2})
+add_test(NAME ${PROJECT_NAME}_mpi_3 COMMAND "mpirun" ${TEST_PARAMETERS3})
+add_test(NAME ${PROJECT_NAME}_serial COMMAND ${PROJECT_NAME})

cpp/demo/checkpointing/main.cpp

@@ -0,0 +1,37 @@
+// ```text
+// Copyright (C) 2024 Abdullah Mujahid
+// This file is part of DOLFINx (https://www.fenicsproject.org)
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+// ```
+
+// # Checkpointing
+//
+
+#include <adios2.h>
+#include <dolfinx.h>
+#include <dolfinx/io/ADIOS2_utils.h>
+#include <dolfinx/io/checkpointing.h>
+#include <mpi.h>
+
+using namespace dolfinx;
+using namespace dolfinx::io;
+
+int main(int argc, char* argv[])
+{
+  dolfinx::init_logging(argc, argv);
+  MPI_Init(&argc, &argv);
+
+  // Create mesh and function space
+  auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet);
+  auto mesh = std::make_shared<mesh::Mesh<float>>(mesh::create_rectangle<float>(
+      MPI_COMM_WORLD, {{{0.0, 0.0}, {1.0, 1.0}}}, {4, 4},
+      mesh::CellType::quadrilateral, part));
+
+  auto writer = ADIOS2Engine(mesh->comm(), "mesh.bp", "mesh-write", "BP5",
+                             adios2::Mode::Write);
+
+  io::checkpointing::write(writer, mesh);
+
+  MPI_Finalize();
+  return 0;
+}

cpp/dolfinx/io/ADIOS2_utils.cpp

@@ -0,0 +1,36 @@
+// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken and Garth N. Wells
+//
+// This file is part of DOLFINX (https://www.fenicsproject.org)
+//
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+
+#ifdef HAS_ADIOS2
+
+#include "ADIOS2_utils.h"
+#include <adios2.h>
+#include <mpi.h>
+
+using namespace dolfinx::io;
+
+//-----------------------------------------------------------------------------
+ADIOS2Engine::ADIOS2Engine(MPI_Comm comm, const std::filesystem::path& filename,
+                           std::string tag, std::string engine,
+                           const adios2::Mode mode)
+
+    : _adios(std::make_unique<adios2::ADIOS>(comm)),
+      _io(std::make_unique<adios2::IO>(_adios->DeclareIO(tag)))
+{
+  _io->SetEngine(engine);
+  _engine = std::make_unique<adios2::Engine>(_io->Open(filename, mode));
+}
+//-----------------------------------------------------------------------------
+ADIOS2Engine::~ADIOS2Engine() { close(); }
+//-----------------------------------------------------------------------------
+void ADIOS2Engine::close()
+{
+  assert(_engine);
+  if (*_engine)
+    _engine->Close();
+}
+
+#endif

cpp/dolfinx/io/ADIOS2_utils.h

@@ -0,0 +1,69 @@
+// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken and Garth N. Wells
+//
+// This file is part of DOLFINX (https://www.fenicsproject.org)
+//
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+
+#pragma once
+
+#ifdef HAS_ADIOS2
+
+#include <adios2.h>
+#include <cassert>
+#include <filesystem>
+#include <mpi.h>
+
+/// @file ADIOS2_utils.h
+/// @brief Utils for ADIOS2
+
+namespace dolfinx::io
+{
+
+/// ADIOS2-based writers/readers
+class ADIOS2Engine
+{
+public:
+  /// @brief Create an ADIOS2-based engine writer/reader
+  /// @param[in] comm The MPI communicator
+  /// @param[in] filename Name of output file
+  /// @param[in] tag The ADIOS2 object name
+  /// @param[in] engine ADIOS2 engine type. See
+  /// https://adios2.readthedocs.io/en/latest/engines/engines.html.
+  /// @param[in] mode ADIOS2 mode, default is Write or Read
+  ADIOS2Engine(MPI_Comm comm, const std::filesystem::path& filename,
+               std::string tag, std::string engine = "BP5",
+               const adios2::Mode mode = adios2::Mode::Write);
+
+  /// @brief Move constructor
+  ADIOS2Engine(ADIOS2Engine&& engine) = default;
+
+  /// @brief Copy constructor
+  ADIOS2Engine(const ADIOS2Engine&) = delete;
+
+  /// @brief Destructor
+  ~ADIOS2Engine();
+
+  /// @brief Move assignment
+  ADIOS2Engine& operator=(ADIOS2Engine&& engine) = default;
+
+  // Copy assignment
+  ADIOS2Engine& operator=(const ADIOS2Engine&) = delete;
+
+  /// @brief  Close the file
+  void close();
+
+  /// @brief  Get the IO object
+  std::shared_ptr<adios2::IO> io() { return _io; }
+
+  /// @brief  Close the Engine object
+  std::shared_ptr<adios2::Engine> engine() { return _engine; }
+
+protected:
+  std::shared_ptr<adios2::ADIOS> _adios;
+  std::shared_ptr<adios2::IO> _io;
+  std::shared_ptr<adios2::Engine> _engine;
+};
+
+} // namespace dolfinx::io
+
+#endif

cpp/dolfinx/io/CMakeLists.txt

@@ -1,7 +1,9 @@
 set(HEADERS_io
     ${CMAKE_CURRENT_SOURCE_DIR}/dolfinx_io.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2_utils.h
     ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.h
     ${CMAKE_CURRENT_SOURCE_DIR}/cells.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/checkpointing.h
     ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.h
     ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.h
     ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.h
@@ -14,8 +16,10 @@ set(HEADERS_io
 
 target_sources(
   dolfinx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cells.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/checkpointing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.cpp

cpp/dolfinx/io/checkpointing.cpp

@@ -0,0 +1,151 @@
+// Copyright (C) 2024 Abdullah Mujahid
+//
+// This file is part of DOLFINX (https://www.fenicsproject.org)
+//
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+
+#ifdef HAS_ADIOS2
+
+#include "checkpointing.h"
+#include "ADIOS2_utils.h"
+#include <adios2.h>
+#include <basix/finite-element.h>
+#include <dolfinx/mesh/Mesh.h>
+#include <mpi.h>
+
+using namespace dolfinx;
+using namespace dolfinx::io;
+
+/// @file checkpointing.h
+/// @brief ADIOS2 based checkpointing
+namespace
+{
+std::map<basix::element::lagrange_variant, std::string> lagrange_variants{
+    {basix::element::lagrange_variant::unset, "unset"},
+    {basix::element::lagrange_variant::equispaced, "equispaced"},
+    {basix::element::lagrange_variant::gll_warped, "gll_warped"},
+};
+
+template <std::floating_point T>
+void _write(ADIOS2Engine& adios2engine,
+            std::shared_ptr<dolfinx::mesh::Mesh<T>> mesh)
+{
+
+  auto io = adios2engine.io();
+  auto writer = adios2engine.engine();
+
+  const dolfinx::mesh::Geometry<T>& geometry = mesh->geometry();
+  auto topology = mesh->topology();
+
+  std::int16_t mesh_dim = geometry.dim();
+
+  auto imap = mesh->geometry().index_map();
+  std::uint64_t num_nodes_global = imap->size_global();
+  std::uint32_t num_nodes_local = imap->size_local();
+  std::uint64_t offset = imap->local_range()[0];
+
+  const std::shared_ptr<const dolfinx::common::IndexMap> topo_imap
+      = topology->index_map(mesh_dim);
+  std::uint64_t num_cells_global = topo_imap->size_global();
+  std::uint32_t num_cells_local = topo_imap->size_local();
+  std::uint64_t cell_offset = topo_imap->local_range()[0];
+
+  auto cmap = mesh->geometry().cmap();
+  auto geom_layout = cmap.create_dof_layout();
+  std::uint32_t num_dofs_per_cell
+      = geom_layout.num_entity_closure_dofs(mesh_dim);
+
+  const std::vector<int64_t> input_global_indices
+      = geometry.input_global_indices();
+  const std::span<const int64_t> input_global_indices_span(
+      input_global_indices.begin(), num_nodes_local);
+  const std::span<const T> mesh_x = geometry.x();
+
+  auto connectivity = topology->connectivity(mesh_dim, 0);
+  auto topology_array = connectivity->array();
+  auto topology_offsets = connectivity->offsets();
+
+  const std::span<const int32_t> topology_array_span(
+      topology_array.begin(), topology_offsets[num_cells_local]);
+  std::vector<std::int64_t> topology_array_global(
+      topology_offsets[num_cells_local]);
+
+  std::iota(topology_array_global.begin(), topology_array_global.end(), 0);
+
+  imap->local_to_global(topology_array_span, topology_array_global);
+
+  for (std::size_t i = 0; i < num_cells_local + 1; ++i)
+    topology_offsets[i] += cell_offset * num_dofs_per_cell;
+
+  const std::span<const int32_t> topology_offsets_span(topology_offsets.begin(),
+                                                       num_cells_local + 1);
+
+  io->DefineAttribute<std::string>("name", mesh->name);
+  io->DefineAttribute<std::int16_t>("dim", geometry.dim());
+  io->DefineAttribute<std::string>("cell_type",
+                                   dolfinx::mesh::to_string(cmap.cell_shape()));
+  io->DefineAttribute<std::int32_t>("degree", cmap.degree());
+  io->DefineAttribute<std::string>("lagrange_variant",
+                                   lagrange_variants[cmap.variant()]);
+
+  adios2::Variable<std::uint64_t> var_num_nodes
+      = io->DefineVariable<std::uint64_t>("num_nodes");
+  adios2::Variable<std::uint64_t> var_num_cells
+      = io->DefineVariable<std::uint64_t>("num_cells");
+  adios2::Variable<std::uint32_t> var_num_dofs_per_cell
+      = io->DefineVariable<std::uint32_t>("num_dofs_per_cell");
+
+  adios2::Variable<std::int64_t> var_input_global_indices
+      = io->DefineVariable<std::int64_t>(
+          "input_global_indices", {num_nodes_global}, {offset},
+          {num_nodes_local}, adios2::ConstantDims);
+
+  adios2::Variable<T> var_x
+      = io->DefineVariable<T>("x", {num_nodes_global, 3}, {offset, 0},
+                              {num_nodes_local, 3}, adios2::ConstantDims);
+
+  adios2::Variable<std::int64_t> var_topology_array
+      = io->DefineVariable<std::int64_t>(
+          "topology_array", {num_cells_global * num_dofs_per_cell},
+          {cell_offset * num_dofs_per_cell},
+          {num_cells_local * num_dofs_per_cell}, adios2::ConstantDims);
+
+  adios2::Variable<std::int32_t> var_topology_offsets
+      = io->DefineVariable<std::int32_t>(
+          "topology_offsets", {num_cells_global + 1}, {cell_offset},
+          {num_cells_local + 1}, adios2::ConstantDims);
+
+  writer->BeginStep();
+  writer->Put(var_num_nodes, num_nodes_global);
+  writer->Put(var_num_cells, num_cells_global);
+  writer->Put(var_num_dofs_per_cell, num_dofs_per_cell);
+  writer->Put(var_input_global_indices, input_global_indices_span.data());
+  writer->Put(var_x, mesh_x.subspan(0, num_nodes_local * 3).data());
+  writer->Put(var_topology_array, topology_array_global.data());
+  writer->Put(var_topology_offsets, topology_offsets_span.data());
+  writer->EndStep();
+}
+
+} // namespace
+
+using namespace dolfinx::io::checkpointing;
+
+//-----------------------------------------------------------------------------
+void dolfinx::io::checkpointing::write(
+    ADIOS2Engine& adios2engine,
+    std::shared_ptr<dolfinx::mesh::Mesh<float>> mesh)
+{
+
+  _write(adios2engine, mesh);
+}
+
+//-----------------------------------------------------------------------------
+void dolfinx::io::checkpointing::write(
+    ADIOS2Engine& adios2engine,
+    std::shared_ptr<dolfinx::mesh::Mesh<double>> mesh)
+{
+
+  _write(adios2engine, mesh);
+}
+
+#endif

cpp/dolfinx/io/checkpointing.h

@@ -0,0 +1,39 @@
+// Copyright (C) 2024 Abdullah Mujahid
+//
+// This file is part of DOLFINX (https://www.fenicsproject.org)
+//
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+
+#pragma once
+
+#ifdef HAS_ADIOS2
+
+#include "ADIOS2_utils.h"
+#include <adios2.h>
+#include <basix/finite-element.h>
+#include <dolfinx/mesh/Mesh.h>
+#include <mpi.h>
+
+/// @file checkpointing.h
+/// @brief ADIOS2 based checkpointing
+
+namespace dolfinx::io::checkpointing
+{
+
+/// @brief Write mesh to a file.
+///
+/// @param[in] adios2engine ADIOS2Engine
+/// @param[in] mesh Mesh of type float to write to the file
+void write(ADIOS2Engine& adios2engine,
+           std::shared_ptr<dolfinx::mesh::Mesh<float>> mesh);
+
+/// @brief Write mesh to a file.
+///
+/// @param[in] adios2engine ADIOS2Engine
+/// @param[in] mesh Mesh of type double to write to the file
+void write(ADIOS2Engine& adios2engine,
+           std::shared_ptr<dolfinx::mesh::Mesh<double>> mesh);
+
+} // namespace dolfinx::io::checkpointing
+
+#endif