Bugs + additional features in checkpoints

LICENSE

@@ -1,19 +1,8 @@
 BSD 3-Clause License
 
-Copyright (c) 2024, Princeton University (PU), and
-Princeton Plasma Physics Lab (PPPL).
+Copyright (c) 2021-present, Entity development team.
 All rights reserved.
 
-Core developers (alphabetical order):
-* Alexander Chernoglazov
-* Benjamin Crinquand
-* Alisa Galishnikova
-* Hayk Hakobyan
-* Jens Mahlmann
-* Sasha Philippov
-* Arno Vanthieghem
-* Muni Zhou
-
 Redistribution and use in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:
 

dev/nix/adios2.nix

@@ -19,10 +19,12 @@ let
     BUILD_TESTING = "OFF";
     ADIOS2_BUILD_EXAMPLES = "OFF";
     ADIOS2_USE_MPI = if mpi then "ON" else "OFF";
+    ADIOS2_HAVE_HDF5_VOL = if mpi then "ON" else "OFF";
     CMAKE_BUILD_TYPE = "Release";
-  } // (if !mpi then { ADIOS2_HAVE_HDF5_VOL = "OFF"; } else { });
+  };
+  stdenv = pkgs.gcc13Stdenv;
 in
-pkgs.stdenv.mkDerivation {
+stdenv.mkDerivation {
   pname = "${name}${if hdf5 then "-hdf5" else ""}${if mpi then "-mpi" else ""}";
   version = "${version}";
   src = pkgs.fetchgit {
@@ -36,12 +38,10 @@ pkgs.stdenv.mkDerivation {
     perl
   ];
 
-  propagatedBuildInputs =
-    [
-      pkgs.gcc13
-    ]
-    ++ (if hdf5 then (if mpi then [ pkgs.hdf5-mpi ] else [ pkgs.hdf5 ]) else [ ])
-    ++ (if mpi then [ pkgs.openmpi ] else [ ]);
+  propagatedBuildInputs = [
+    pkgs.gcc13
+  ] ++ (if hdf5 then (if mpi then [ pkgs.hdf5-mpi ] else [ pkgs.hdf5-cpp ]) else [ ]);
+  # ++ (if mpi then [ pkgs.openmpi ] else [ ]);
 
   configurePhase = ''
     cmake -B build $src ${

dev/nix/kokkos.nix

@@ -20,11 +20,18 @@ let
     "CUDA" = with pkgs.cudaPackages; [
       cudatoolkit
       cuda_cudart
+      pkgs.gcc13
     ];
     "NONE" = [
       pkgs.gcc13
     ];
   };
+  getArch =
+    _:
+    if gpu != "NONE" && arch == "NATIVE" then
+      throw "Please specify an architecture when the GPU support is enabled. Available architectures: https://kokkos.org/kokkos-core-wiki/keywords.html#architectures"
+    else
+      arch;
   cmakeExtraFlags = {
     "HIP" = [
       "-D Kokkos_ENABLE_HIP=ON"
@@ -39,13 +46,6 @@ let
     ];
     "NONE" = [ ];
   };
-  getArch =
-    _:
-    if gpu != "NONE" && arch == "NATIVE" then
-      throw "Please specify an architecture when the GPU support is enabled. Available architectures: https://kokkos.org/kokkos-core-wiki/keywords.html#architectures"
-    else
-      arch;
-
 in
 pkgs.stdenv.mkDerivation rec {
   pname = "${name}";

dev/nix/shell.nix

@@ -42,7 +42,8 @@ pkgs.mkShell {
     zlib
     cmake
 
-    clang-tools
+    llvmPackages_18.clang-tools
+    libgcc
 
     adios2Pkg
     kokkosPkg

input.example.toml

@@ -474,6 +474,24 @@
   #   @note: 0 = disable checkpointing
   #   @note: -1 = keep all checkpoints
   keep = ""
+  # Write a checkpoint once after a fixed walltime 
+  #   @type: string
+  #   @default: "00:00:00" (disabled)
+  #   @note: The format is "HH:MM:SS"
+  #   @note: Useful for long-running simulations to ensure that ...
+  #   @note: ... the simulation can be resumed after allocation expires
+  #   @note: Empty string means disables this (still writes checkpoints on interval)
+  #   @note: This does not exit the simulation
+  walltime = ""
+  # Parent directory to write checkpoints to
+  #   @type: string
+  #   @default: "<simname>.ckpt"
+  #   @note: The directory is created if it does not exist
+  write_path = ""
+  # Parent directory to use when resuming from a checkpoint
+  #   @type: string
+  #   @default: inherits from `write_path`
+  read_path = ""
 
   # @inferred:
   # - is_resuming

src/archetypes/energy_dist.h

@@ -392,7 +392,7 @@ namespace arch {
 
       Inline void operator()(const coord_t<M::Dim>& x_Code,
                              vec_t<Dim::_3D>&       v,
-                             spidx_t                sp = 0) const override {
+                             spidx_t = 0) const override {
         if (cmp::AlmostZero(temperature)) {
           v[0] = ZERO;
           v[1] = ZERO;

src/checkpoint/reader.cpp

@@ -110,7 +110,7 @@ namespace checkpoint {
 
   void ReadParticlePayloads(adios2::IO&        io,
                             adios2::Engine&    reader,
-                            spidx_t     s,
+                            spidx_t            s,
                             array_t<real_t**>& array,
                             std::size_t        nplds,
                             npart_t            count,
@@ -130,64 +130,32 @@ namespace checkpoint {
     }
   }
 
-  template void ReadFields<Dim::_1D, 3>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_1D, 3>&);
-  template void ReadFields<Dim::_2D, 3>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_2D, 3>&);
-  template void ReadFields<Dim::_3D, 3>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_3D, 3>&);
-  template void ReadFields<Dim::_1D, 6>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_1D, 6>&);
-  template void ReadFields<Dim::_2D, 6>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_2D, 6>&);
-  template void ReadFields<Dim::_3D, 6>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        const adios2::Box<adios2::Dims>&,
-                                        ndfield_t<Dim::_3D, 6>&);
+#define CHECKPOINT_FIELDS(D, N)                                                \
+  template void ReadFields<D, N>(adios2::IO&,                                  \
+                                 adios2::Engine&,                              \
+                                 const std::string&,                           \
+                                 const adios2::Box<adios2::Dims>&,             \
+                                 ndfield_t<D, N>&);
+  CHECKPOINT_FIELDS(Dim::_1D, 3)
+  CHECKPOINT_FIELDS(Dim::_2D, 3)
+  CHECKPOINT_FIELDS(Dim::_3D, 3)
+  CHECKPOINT_FIELDS(Dim::_1D, 6)
+  CHECKPOINT_FIELDS(Dim::_2D, 6)
+  CHECKPOINT_FIELDS(Dim::_3D, 6)
+#undef CHECKPOINT_FIELDS
 
-  template void ReadParticleData<int>(adios2::IO&,
-                                      adios2::Engine&,
-                                      const std::string&,
-                                      spidx_t,
-                                      array_t<int*>&,
-                                      npart_t,
-                                      npart_t);
-  template void ReadParticleData<float>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        spidx_t,
-                                        array_t<float*>&,
-                                        npart_t,
-                                        npart_t);
-  template void ReadParticleData<double>(adios2::IO&,
-                                         adios2::Engine&,
-                                         const std::string&,
-                                         spidx_t,
-                                         array_t<double*>&,
-                                         npart_t,
-                                         npart_t);
-  template void ReadParticleData<short>(adios2::IO&,
-                                        adios2::Engine&,
-                                        const std::string&,
-                                        spidx_t,
-                                        array_t<short*>&,
-                                        npart_t,
-                                        npart_t);
+#define CHECKPOINT_PARTICLE_DATA(T)                                            \
+  template void ReadParticleData<T>(adios2::IO&,                               \
+                                    adios2::Engine&,                           \
+                                    const std::string&,                        \
+                                    spidx_t,                                   \
+                                    array_t<T*>&,                              \
+                                    npart_t,                                   \
+                                    npart_t);
+  CHECKPOINT_PARTICLE_DATA(int)
+  CHECKPOINT_PARTICLE_DATA(float)
+  CHECKPOINT_PARTICLE_DATA(double)
+  CHECKPOINT_PARTICLE_DATA(short)
+#undef CHECKPOINT_PARTICLE_DATA
 
 } // namespace checkpoint

src/checkpoint/tests/checkpoint-nompi.cpp

@@ -12,14 +12,13 @@
 
 #include <filesystem>
 #include <iostream>
-#include <stdexcept>
 
 using namespace ntt;
 using namespace checkpoint;
 
 void cleanup() {
   namespace fs = std::filesystem;
-  fs::path temp_path { "checkpoints" };
+  fs::path temp_path { "chck" };
   fs::remove_all(temp_path);
 }
 
@@ -58,18 +57,21 @@ auto main(int argc, char* argv[]) -> int {
         CreateRangePolicy<Dim::_3D>({ i1min, i2min, i3min },
                                     { i1max, i2max, i3max }),
         Lambda(index_t i1, index_t i2, index_t i3) {
-          field1(i1, i2, i3, 0) = i1 + i2 + i3;
-          field1(i1, i2, i3, 1) = i1 * i2 / i3;
-          field1(i1, i2, i3, 2) = i1 / i2 * i3;
-          field1(i1, i2, i3, 3) = i1 + i2 - i3;
-          field1(i1, i2, i3, 4) = i1 * i2 + i3;
-          field1(i1, i2, i3, 5) = i1 / i2 - i3;
-          field2(i1, i2, i3, 0) = -(i1 + i2 + i3);
-          field2(i1, i2, i3, 1) = -(i1 * i2 / i3);
-          field2(i1, i2, i3, 2) = -(i1 / i2 * i3);
-          field2(i1, i2, i3, 3) = -(i1 + i2 - i3);
-          field2(i1, i2, i3, 4) = -(i1 * i2 + i3);
-          field2(i1, i2, i3, 5) = -(i1 / i2 - i3);
+          const auto i1_        = static_cast<real_t>(i1);
+          const auto i2_        = static_cast<real_t>(i2);
+          const auto i3_        = static_cast<real_t>(i3);
+          field1(i1, i2, i3, 0) = i1_ + i2_ + i3_;
+          field1(i1, i2, i3, 1) = i1_ * i2_ / i3_;
+          field1(i1, i2, i3, 2) = i1_ / i2_ * i3_;
+          field1(i1, i2, i3, 3) = i1_ + i2_ - i3_;
+          field1(i1, i2, i3, 4) = i1_ * i2_ + i3_;
+          field1(i1, i2, i3, 5) = i1_ / i2_ - i3_;
+          field2(i1, i2, i3, 0) = -(i1_ + i2_ + i3_);
+          field2(i1, i2, i3, 1) = -(i1_ * i2_ / i3_);
+          field2(i1, i2, i3, 2) = -(i1_ / i2_ * i3_);
+          field2(i1, i2, i3, 3) = -(i1_ + i2_ - i3_);
+          field2(i1, i2, i3, 4) = -(i1_ * i2_ + i3_);
+          field2(i1, i2, i3, 5) = -(i1_ / i2_ - i3_);
         });
       Kokkos::parallel_for(
         "fillPrtl1",
@@ -88,11 +90,12 @@ auto main(int argc, char* argv[]) -> int {
     }
 
     adios2::ADIOS adios;
+    const path_t  checkpoint_path { "chck" };
 
     {
       // write checkpoint
-      Writer writer;
-      writer.init(&adios, 0, 0.0, 1);
+      Writer writer {};
+      writer.init(&adios, checkpoint_path, 0, 0.0, 1);
 
       writer.defineFieldVariables(SimEngine::GRPIC,
                                   { nx1_gh, nx2_gh, nx3_gh },
@@ -127,7 +130,7 @@ auto main(int argc, char* argv[]) -> int {
       array_t<real_t*> u2_read { "u_2", npart2 };
 
       adios2::IO     io     = adios.DeclareIO("checkpointRead");
-      adios2::Engine reader = io.Open("checkpoints/step-00000000.bp",
+      adios2::Engine reader = io.Open(checkpoint_path / "step-00000000.bp",
                                       adios2::Mode::Read);
       reader.BeginStep();