From 91b6c04272f6654f94ae07c1d1d8e5896e351159 Mon Sep 17 00:00:00 2001 From: "Jack S. Hale" Date: Thu, 26 Feb 2026 15:37:03 +0100 Subject: [PATCH 1/5] Expore jit MPI comm through to form --- python/dolfinx/fem/forms.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 3c994db79c..32adf36fb9 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -324,6 +324,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. @@ -333,6 +334,8 @@ def form( dtype: Scalar type to use for the compiled form. form_compiler_options: See :func:`ffcx_jit ` jit_options: See :func:`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 @@ -369,8 +372,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 From 700971a1c93b76d4f104594f96961e0955250d60 Mon Sep 17 00:00:00 2001 From: "Jack S. Hale" Date: Thu, 26 Feb 2026 15:53:12 +0100 Subject: [PATCH 2/5] Also add to experimental jit compiler. --- python/dolfinx/fem/forms.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 32adf36fb9..33f0b7e5bd 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -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. @@ -262,6 +263,8 @@ def mixed_topology_form( dtype: Scalar type to use for the compiled form. form_compiler_options: See :func:`ffcx_jit ` jit_options: See :func:`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 @@ -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 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, From 029f1c1b7a715b6c6370bfe733dae84734826052 Mon Sep 17 00:00:00 2001 From: "Jack S. Hale" Date: Thu, 26 Feb 2026 16:20:04 +0100 Subject: [PATCH 3/5] Apply suggestion from @jorgensd MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jørgen Schartum Dokken --- python/dolfinx/fem/forms.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 33f0b7e5bd..7126887b3f 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -292,10 +292,7 @@ def mixed_topology_form( 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 + comm = mesh.comm if jit_comm is None else jit_comm ufcx_forms = [] modules = [] From d7e8d2ca22d8f904c34ceabf5bf799444c0e01a9 Mon Sep 17 00:00:00 2001 From: "Jack S. Hale" Date: Thu, 26 Feb 2026 16:20:55 +0100 Subject: [PATCH 4/5] And @jorgensd suggestion again --- python/dolfinx/fem/forms.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index 7126887b3f..cae2891ddc 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -378,10 +378,7 @@ 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 + comm = mesh.comm if jit_comm is None else jit_comm ufcx_form, module, code = jit.ffcx_jit( comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options From 6f85aea10ec94758db5eb42f765457c4aa7b6e4c Mon Sep 17 00:00:00 2001 From: "Jack S. Hale" Date: Thu, 26 Feb 2026 16:36:18 +0100 Subject: [PATCH 5/5] Typo --- python/dolfinx/fem/forms.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py index cae2891ddc..1de7765544 100644 --- a/python/dolfinx/fem/forms.py +++ b/python/dolfinx/fem/forms.py @@ -378,7 +378,7 @@ def _form(form): msh = domain.ufl_cargo() if msh is None: raise RuntimeError("Expecting to find a Mesh in the form.") - comm = mesh.comm if jit_comm is None else jit_comm + comm = msh.comm if jit_comm is None else jit_comm ufcx_form, module, code = jit.ffcx_jit( comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options