Expose JIT MPI communicator

[jhale]

Allows the compilation of forms on a communicator different to the mesh. The hypothesis is that this may be useful for setting up JIT cache on HPC runs across multiple nodes with shared memory-backed /tmp - this is still to test.

Example of use:

from mpi4py import MPI
from dolfinx.mesh import create_unit_square
from dolfinx.fem import functionspace, form
import ufl

mesh = create_unit_square(MPI.COMM_WORLD, 10, 10)
V = functionspace(mesh, ("CG", 1))
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
a = ufl.inner(u, v)*ufl.dx

# See https://docs.open-mpi.org/en/v5.0.x/man-openmpi/man3/MPI_Comm_split_type.3.html#split-types
shared_mem_comm = MPI.COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED, key=MPI.COMM_WORLD.rank)
a_dolfinx = form(a, jit_comm=shared_mem_comm)

The default behaviour (use the communicator from the form's mesh) is retained when jit_comm=None, which is also the default argument.

[jorgensd]

I think we should do it as a one-liner. Otherwise good

[jhale]

I've moved to a broadcast on the check for root success, and added a code path to check the other ranks actually managed to load from the cache.

[jhale]

Tested on our HPC, seems to work:

#!/bin/bash -l
#SBATCH -p batch
#SBATCH --time=00:00:30
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=32
#SBATCH --output=logs/%x-%j.out
source $HOME/spack/share/spack/setup-env.sh
export XDG_CACHE_HOME=/tmp/jhale/cache
spack env activate $HOME/fenics-shared-mem-test
module load tools/binutils/2.40-GCCcore-13.2.0
srun python dolfinx-shared-mem-comm.py
srun ls -F /tmp/jhale/cache/fenics

from mpi4py import MPI
from dolfinx.mesh import create_unit_square
from dolfinx.fem import functionspace, form, assemble_matrix
import ufl

shared_mem_comm = MPI.COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED, key=MPI.COMM_WORLD.rank)

mesh = create_unit_square(MPI.COMM_WORLD, 32, 32)
V = functionspace(mesh, ("CG", 1))
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)

a = ufl.inner(u, v)*ufl.dx
a_dolfinx = form(a, jit_comm=shared_mem_comm)
MPI.COMM_WORLD.Barrier()

print("JIT done.")

A = assemble_matrix(a_dolfinx)

print("Assembly done.")

Couldn't detect any issues without barrier, but seems sensible to add it.