Checkpointing in FEniCSx¶

Jørgen Schartum Dokken¶

Simula Research Laboratory¶

FEniCS 23'

Input/Output formats for Finite Element software¶

FEM software¶

  • Most finite element software use their own format for inputting/outputting meshes and functions
  • No general consensus on what format to use
  • Most formats are based on VTK in some shape or form
  • Many different file formats support vtk
    • .pvd-files (xml-based)
    • .xdmf-files (xml+binary (.h5)
    • .bp-files (binary files in folder)

What is checkpointing?¶

Checkpointing refers to the ability to store the state of a computation in a way that allows it be continued at a later time without changing the computation’s behavior [1]

[1] Schulz, M. (2011). Checkpointing. In: Padua, D. (eds) Encyclopedia of Parallel Computing. Springer, Boston, MA. 10.1007/978-0-387-09766-4_62

Why separate post-processing and checkpointing?¶

The post-processing landscape keeps on changing

  • Plain text formats are not scalable
  • Formats go in and out of fashion
  • Requirements for visualization and internal representation of data differs

A basis function for Degree 1 Nédélec (first kind) on a triangle from DefElement.

Types of checkpointing¶

  • Snapshot checkpointing
    • Only works while the program is running
    • FEniCS-setting: Mesh has to be the same
  • Recoverable checkpointing
    • Should work on any system
    • Any number of processes

Implementing a snapshot checkpoint¶

In [1]:
import ipyparallel as ipp
cluster = ipp.Cluster(engines="mpi", n=3)
rc = cluster.start_and_connect_sync()

Starting 3 engines with <class 'ipyparallel.cluster.launcher.MPIEngineSetLauncher'>
In [2]:
%%px
from mpi4py import MPI
from dolfinx import mesh, fem
import adios2
from pathlib import Path
import numpy as np
print(f"{MPI.COMM_WORLD.rank=}, {MPI.COMM_WORLD.size=}")

[stdout:0] 
MPI.COMM_WORLD.rank=0, MPI.COMM_WORLD.size=3
[stdout:1] 
MPI.COMM_WORLD.rank=1, MPI.COMM_WORLD.size=3
[stdout:2] 
MPI.COMM_WORLD.rank=2, MPI.COMM_WORLD.size=3

Implementing a snapshot checkpoint (write)¶

In [3]:
%%px

def snapshot_checkpoint_write(uh: fem.Function, filename: Path):
    # Create ADIOS IO
    adios = adios2.ADIOS(uh.function_space.mesh.comm)
    io_name = "SnapshotCheckPoint"
    io = adios.DeclareIO(io_name)
    io.SetEngine("BP4")

    dofmap = uh.function_space.dofmap
    num_dofs_local = dofmap.index_map.size_local * dofmap.index_map_bs
    local_dofs = uh.x.array[:num_dofs_local].copy()

    file = io.Open(str(filename), adios2.Mode.Write)
    file.BeginStep()
    dofs = io.DefineVariable("dofs", local_dofs, count=[num_dofs_local])
    file.Put(dofs, local_dofs, adios2.Mode.Sync)
    file.EndStep()
    file.Close()
    adios.RemoveIO(io_name)

Implementing a snapshot checkpoint (write)¶

In [4]:
%%px
domain = mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
V = fem.FunctionSpace(domain, ("Lagrange", 5))
u = fem.Function(V)
u.interpolate(lambda x: x[0]**5 + 3*x[1]*x[0]**2)
scp_file = Path("u.bp")
snapshot_checkpoint_write(u, scp_file)
In [5]:
import subprocess
from pathlib import Path
scp_file = Path("u.bp")
_ = subprocess.run(["bpls", "-l", "-D", str(scp_file.absolute())])

  double   dofs  [3]*{__} = 2.99914e-53 / 4
        step 0: 
          block 0: [0:871] = 0.000425685 / 3.1
          block 1: [0:862] = 0.0866018 / 4
          block 2: [0:865] = 2.99914e-53 / 0.49024

Implementing a snapshot checkpoint (read)¶

In [6]:
%%px
def snapshot_checkpoint_read(uh: fem.Function, filename: Path):
    adios = adios2.ADIOS(uh.function_space.mesh.comm)
    io_name = "SnapshotCheckPoint"
    io = adios.DeclareIO(io_name)
    io.SetEngine("BP4")

    file = io.Open(str(filename), adios2.Mode.Read)
    file.BeginStep()
    in_variable = io.InquireVariable("dofs")
    in_variable.SetBlockSelection(uh.function_space.mesh.comm.rank)
    file.Get(in_variable, uh.x.array, adios2.Mode.Sync)
    file.EndStep()
    file.Close()
    adios.RemoveIO(io_name)
    uh.x.scatter_forward()
In [7]:
%%px
v = fem.Function(V)
snapshot_checkpoint_read(v, scp_file)
assert(np.allclose(u.x.array, v.x.array))

ADIOS4DOLFINx¶

  • Implemented in Python with ADIOS2 [2]
  • Available at: https://github.com/jorgensd/adios4dolfinx/
  • Snapshot checkpointing
  • Recoverable checkpointing
In [8]:
%%px
import adios4dolfinx as adx
domain = mesh.create_unit_cube(MPI.COMM_WORLD, 10, 7, 3)
V = fem.FunctionSpace(domain, ("N1curl", 3))
u = fem.Function(V)
u.interpolate(lambda x: (x[0], x[1]*x[1], np.sin(x[2])))
cp_file = Path("u_curl.bp")
adx.snapshot_checkpoint(u, cp_file, adios2.Mode.Write)
v = fem.Function(V)
adx.snapshot_checkpoint(v, cp_file, adios2.Mode.Read)
assert np.allclose(v.x.array, u.x.array)

[2] W.F. Godoy, et al. ADIOS 2: The Adaptable Input Output System. A framework for high-performance data management, SoftwareX, Volume 12, 2020, 10.1016/j.softx.2020.100561

How to do recoverable checkpointing?¶

Required information:

  • Mesh (topology, geometry, base-element)
  • FunctionSpace (dofmap)
  • Function coefficient (dofs)
In [9]:
%%px
domain = mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, mesh.CellType.quadrilateral)
V = fem.FunctionSpace(domain, ("NCE", 3))
u = fem.Function(V)
u.interpolate(lambda x: (x[0]**5, 3*x[1]*x[0]**2))
checkpoint_file = Path("function_checkpoint.bp")
adx.write_mesh(domain, checkpoint_file, engine="BP4")
adx.write_function(u, checkpoint_file, engine="BP4")

Required data-structures¶

  uint32_t  CellPermutations  {100} = 0 / 12
        step 0: 
          block 0: [ 0:32] = 0 / 8
          block 1: [33:65] = 0 / 12
          block 2: [66:99] = 0 / 8
  int64_t   Dofmap            {2400} = 0 / 1859
        step 1: 
          block 0: [   0: 791] = 0 / 1373
          block 1: [ 792:1583] = 189 / 1520
          block 2: [1584:2399] = 504 / 1859
  double    Points            {121, 2} = 0 / 1
        step 0: 
          block 0: [  0: 39, 0:1] = 0 / 1
          block 1: [ 40: 80, 0:1] = 0 / 1
          block 2: [ 81:120, 0:1] = 0.3 / 1
  int64_t   Topology          {100, 4} = 0 / 120
        step 0: 
          block 0: [ 0:32, 0:3] = 0 / 93
          block 1: [33:65, 0:3] = 17 / 117
          block 2: [66:99, 0:3] = 38 / 120
  double    Values            {1860} = -0.165 / 0.285
        step 1: 
          block 0: [   0: 620] = -0.00945 / 0.025175
          block 1: [ 621:1232] = -0.105 / 0.135
          block 2: [1233:1859] = -0.165 / 0.285
  int64_t   XDofmap           {101} = 0 / 2400
        step 1: 
          block 0: [  0: 33] = 0 / 792
          block 1: [ 33: 66] = 792 / 1584
          block 2: [ 66:100] = 1584 / 2400

Required data-structures¶

  string   CellType         attr   = "quadrilateral"
  int32_t  Degree           attr   = 1
  int32_t  LagrangeVariant  attr   = 1

Cell-permutations are used to consistently orient any finite element in parallel [3]¶

  uint32_t  CellPermutations  {100} = 0 / 12
        step 0: 
          block 0: [ 0:32] = 0 / 8
          block 1: [33:65] = 0 / 12
          block 2: [66:99] = 0 / 8

[3] M.W. Scroggs, J.S. Dokken, C.N. Richardson, and G.N. Wells. 2022. Construction of Arbitrary Order Finite Element Degree-of-Freedom Maps on Polygonal and Polyhedral Cell Meshes ACM Trans. Math. Softw. 48, 2, 10.1145/3524456

The dofmap is needed to map the coefficients to the correct cell after partitioning¶

  int64_t   Dofmap            {2400} = 0 / 1859
        step 1: 
          block 0: [   0: 791] = 0 / 1373
          block 1: [ 792:1583] = 189 / 1520
          block 2: [1584:2399] = 504 / 1859
  int64_t   XDofmap           {101} = 0 / 2400
        step 1: 
          block 0: [  0: 33] = 0 / 792
          block 1: [ 33: 66] = 792 / 1584
          block 2: [ 66:100] = 1584 / 2400

All communication is done using MPI neighbourhood communicators¶

  • Avoids all MPI_Alltoall and MPI_Alltoallv communications
  • Scales better with increasing number of processes[4]

[4] C. Richardson, G.N. Wells Scalable MPI in FEniCS-dolfinx FEniCS 22'

Further work¶

  • Move implementation into C++ and DOLFINx
    • First step in PR 2618
  • Add capability to store MeshTags in checkpointing files

Summary¶

  • A module for snapshot and proper checkpoints for DOLFINx
  • Works for any finite element in DOLFINx in serial and parallel
  • Leverages ADIOS2 for parallel IO

Example of reading the checkpoint of 2 processes¶

In [16]:
new_cluster = ipp.Cluster(engines="mpi", n=2)
rc = new_cluster.start_and_connect_sync(activate=True)

Starting 2 engines with <class 'ipyparallel.cluster.launcher.MPIEngineSetLauncher'>

INFO:ipyparallel.cluster.cluster.1686580696-v7k6:Starting 2 engines with <class 'ipyparallel.cluster.launcher.MPIEngineSetLauncher'>
In [17]:
%%px
import adios4dolfinx as adx
import dolfinx
from pathlib import Path
from mpi4py import MPI
import ufl
checkpoint_file = Path("function_checkpoint.bp")
domain = adx.read_mesh(MPI.COMM_WORLD, checkpoint_file, "BP4", dolfinx.mesh.GhostMode.none)
V = dolfinx.fem.FunctionSpace(domain, ("NCE", 3))
u = dolfinx.fem.Function(V)
adx.read_function(u, checkpoint_file, "BP4")
x = ufl.SpatialCoordinate(domain)
u_exact = ufl.as_vector((x[0]**5, 3*x[1]*x[0]**2))
L2_error = dolfinx.fem.form(ufl.inner(u-u_exact, u-u_exact)*ufl.dx)
print(f"L^2-error {dolfinx.fem.assemble_scalar(L2_error)}")

%px:   0%|          | 0/2 [00:00<?, ?tasks/s]
[stdout:0] 
L^2-error 3.5469155884721383e-09
[stdout:1] 
L^2-error 3.546915588472256e-09