Skip to content

Expose JIT MPI communicator #4107

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jhale merged 8 commits into from
Feb 27, 2026
Merged

Expose JIT MPI communicator #4107

Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
91b6c04
Expore jit MPI comm through to form
jhale Feb 26, 2026
700971a
Also add to experimental jit compiler.
jhale Feb 26, 2026
029f1c1
Apply suggestion from @jorgensd
jhale Feb 26, 2026
d7e8d2c
And @jorgensd suggestion again
jhale Feb 26, 2026
6f85aea
Typo
jhale Feb 26, 2026
e9b3730
Fix typing
jhale Feb 27, 2026
f1bbaea
Improve error handling in JIT - now fails if cache doesn't load.
jhale Feb 27, 2026
1ce8db4
Tweaks.
jhale Feb 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions python/dolfinx/fem/forms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,7 @@ def mixed_topology_form(
dtype: npt.DTypeLike = default_scalar_type,
form_compiler_options: dict | None = None,
jit_options: dict | None = None,
jit_comm: MPI.IntraComm | None = None,
entity_maps: Sequence[_EntityMap] | None = None,
):
"""Create a mixed-topology from from an array of Forms.
Expand All @@ -262,6 +263,8 @@ def mixed_topology_form(
dtype: Scalar type to use for the compiled form.
form_compiler_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>`
jit_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>`.
jit_comm: MPI communicator used when compiling the form. If
``None``, then ``form.mesh.comm``.
entity_maps: If any trial functions, test functions, or
coefficients in the form are not defined over the same mesh
as the integration domain (the domain associated with the
Expand All @@ -287,13 +290,19 @@ def mixed_topology_form(
assert all([d is data[0] for d in data if d is not None])

mesh = domain.ufl_cargo()
if mesh is None:
raise RuntimeError("Expecting to find a Mesh in the form.")
if jit_comm is None:
comm = mesh.comm
else:
comm = jit_comm
Comment thread
jhale marked this conversation as resolved.
Outdated

ufcx_forms = []
modules = []
codes = []
for form in forms:
ufcx_form, module, code = jit.ffcx_jit(
mesh.comm,
comm,
form,
form_compiler_options=form_compiler_options,
jit_options=jit_options,
Expand Down Expand Up @@ -324,6 +333,7 @@ def form(
dtype: npt.DTypeLike = default_scalar_type,
form_compiler_options: dict | None = None,
jit_options: dict | None = None,
jit_comm: MPI.IntraComm | None = None,
entity_maps: Sequence[_EntityMap] | None = None,
):
"""Create a Form or list of Forms.
Expand All @@ -333,6 +343,8 @@ def form(
dtype: Scalar type to use for the compiled form.
form_compiler_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>`
jit_options: See :func:`ffcx_jit <dolfinx.jit.ffcx_jit>`.
jit_comm: MPI communicator used when compiling the form. If
`None`, then `form.mesh.comm`.
entity_maps: If any trial functions, test functions, or
coefficients in the form are not defined over the same mesh
as the integration domain (the domain associated with the
Expand Down Expand Up @@ -369,8 +381,13 @@ def _form(form):
msh = domain.ufl_cargo()
if msh is None:
raise RuntimeError("Expecting to find a Mesh in the form.")
if jit_comm is None:
comm = msh.comm
else:
comm = jit_comm

ufcx_form, module, code = jit.ffcx_jit(
msh.comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options
comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options
)

# For each argument in form extract its function space
Expand Down
Loading