Add support for distributed compilation in DaceProgram (#1551)

dace/frontend/python/parser.py

@@ -21,6 +21,12 @@
 except ImportError:
     from typing_compat import get_origin, get_args
 
+try:
+    import mpi4py
+    from dace.sdfg.utils import distributed_compile
+except ImportError:
+    mpi4py = None
+
 ArgTypes = Dict[str, Data]
 
 
@@ -443,9 +449,12 @@ def __call__(self, *args, **kwargs):
                 sdfg.simplify()
 
         with hooks.invoke_sdfg_call_hooks(sdfg) as sdfg:
-            # Compile SDFG (note: this is done after symbol inference due to shape
-            # altering transformations such as Vectorization)
-            binaryobj = sdfg.compile(validate=self.validate)
+            if not mpi4py:
+                # Compile SDFG (note: this is done after symbol inference due to shape
+                # altering transformations such as Vectorization)
+                binaryobj = sdfg.compile(validate=self.validate)
+            else:
+                binaryobj = distributed_compile(sdfg, mpi4py.MPI.COMM_WORLD, validate=self.validate)
 
             # Recreate key and add to cache
             cachekey = self._cache.make_key(argtypes, specified, self.closure_array_keys, self.closure_constant_keys,

dace/sdfg/utils.py

@@ -1355,7 +1355,7 @@ def load_precompiled_sdfg(folder: str):
                               csdfg.ReloadableDLL(os.path.join(folder, 'build', f'lib{sdfg.name}.{suffix}'), sdfg.name))
 
 
-def distributed_compile(sdfg: SDFG, comm) -> csdfg.CompiledSDFG:
+def distributed_compile(sdfg: SDFG, comm, validate: bool = True) -> csdfg.CompiledSDFG:
     """
     Compiles an SDFG in rank 0 of MPI communicator ``comm``. Then, the compiled SDFG is loaded in all other ranks.
 
@@ -1371,7 +1371,7 @@ def distributed_compile(sdfg: SDFG, comm) -> csdfg.CompiledSDFG:
 
     # Rank 0 compiles SDFG.
     if rank == 0:
-        func = sdfg.compile()
+        func = sdfg.compile(validate=validate)
         folder = sdfg.build_folder
 
     # Broadcasts build folder.