mpi4py

dace/dtypes.py

@@ -1200,6 +1200,7 @@ def isconstant(var):
 complex64 = typeclass(numpy.complex64)
 complex128 = typeclass(numpy.complex128)
 string = stringtype()
+MPI_Request = opaque('MPI_Request')
 
 
 @undefined_safe_enum

dace/frontend/python/newast.py

@@ -1303,6 +1303,14 @@ def defined(self):
         # Add SDFG arrays, in case a replacement added a new output
         result.update(self.sdfg.arrays)
 
+        # MPI-related stuff
+        result.update({k: self.sdfg.process_grids[v] for k, v in self.variables.items() if v in self.sdfg.process_grids})
+        try:
+            from mpi4py import MPI
+            result.update({k: v for k, v in self.globals.items() if isinstance(v, MPI.Comm)})
+        except:
+            pass
+
         return result
 
     def _add_state(self, label=None):
@@ -4356,8 +4364,11 @@ def visit_Call(self, node: ast.Call, create_callbacks=False):
             # Add object as first argument
             if modname in self.variables.keys():
                 arg = self.variables[modname]
-            else:
+            elif modname in self.scope_vars.keys():
                 arg = self.scope_vars[modname]
+            else:
+                # Fallback to (name, object)
+                arg = (modname, self.defined[modname])
             args.append(arg)
         # Otherwise, try to find a default implementation for the SDFG
         elif not found_ufunc:
@@ -4667,7 +4678,9 @@ def _gettype(self, opnode: ast.AST) -> List[Tuple[str, str]]:
 
         result = []
         for operand in operands:
-            if isinstance(operand, str) and operand in self.sdfg.arrays:
+            if isinstance(operand, str) and operand in self.sdfg.process_grids:
+                result.append((operand, type(self.sdfg.process_grids[operand]).__name__))
+            elif isinstance(operand, str) and operand in self.sdfg.arrays:
                 result.append((operand, type(self.sdfg.arrays[operand])))
             elif isinstance(operand, str) and operand in self.scope_arrays:
                 result.append((operand, type(self.scope_arrays[operand])))

dace/frontend/python/preprocessing.py

@@ -1503,6 +1503,62 @@ def find_disallowed_statements(node: ast.AST):
     return None
 
 
+class MPIResolver(ast.NodeTransformer):
+    """ Resolves mpi4py-related constants, e.g., mpi4py.MPI.COMM_WORLD. """
+    def __init__(self, globals: Dict[str, Any]):
+        from mpi4py import MPI
+        self.globals = globals
+        self.MPI = MPI
+        self.parent = None
+
+    def visit(self, node):
+        node.parent = self.parent
+        self.parent = node
+        node = super().visit(node)
+        if isinstance(node, ast.AST):
+            self.parent = node.parent
+        return node
+
+    def visit_Name(self, node: ast.Name) -> Union[ast.Name, ast.Attribute]:
+        self.generic_visit(node)
+        if node.id in self.globals:
+            obj = self.globals[node.id]
+            if isinstance(obj, self.MPI.Comm):
+                lattr = ast.Attribute(ast.Name(id='mpi4py', ctx=ast.Load), attr='MPI')
+                if obj is self.MPI.COMM_NULL:
+                    newnode = ast.copy_location(ast.Attribute(value=lattr, attr='COMM_NULL'), node)
+                    newnode.parent = node.parent
+                    return newnode
+        return node
+
+    def visit_Attribute(self, node: ast.Attribute) -> ast.Attribute:
+        self.generic_visit(node)
+        if isinstance(node.attr, str) and node.attr == 'Request':
+            try:
+                val = astutils.evalnode(node, self.globals)
+                if val is self.MPI.Request and not isinstance(node.parent, ast.Attribute):
+                    newnode = ast.copy_location(
+                        ast.Attribute(value=ast.Name(id='dace', ctx=ast.Load), attr='MPI_Request'), node)
+                    newnode.parent = node.parent
+                    return newnode
+            except SyntaxError:
+                pass
+        return node
+
+
+class ModuloConverter(ast.NodeTransformer):
+    """ Converts a % b expressions to (a + b) % b for C/C++ compatibility. """
+
+    def visit_BinOp(self, node: ast.BinOp) -> ast.BinOp:
+        if isinstance(node.op, ast.Mod):
+            left = self.generic_visit(node.left)
+            right = self.generic_visit(node.right)
+            newleft = ast.copy_location(ast.BinOp(left=left, op=ast.Add(), right=copy.deepcopy(right)), left)
+            node.left = newleft
+            return node
+        return self.generic_visit(node)
+
+
 def preprocess_dace_program(f: Callable[..., Any],
                             argtypes: Dict[str, data.Data],
                             global_vars: Dict[str, Any],
@@ -1544,6 +1600,12 @@ def preprocess_dace_program(f: Callable[..., Any],
         newmod = global_vars[mod]
         #del global_vars[mod]
         global_vars[modval] = newmod
+
+    try:
+        src_ast = MPIResolver(global_vars).visit(src_ast)
+    except ModuleNotFoundError:
+        pass
+    src_ast = ModuloConverter().visit(src_ast)
 
     # Resolve constants to their values (if they are not already defined in this scope)
     # and symbols to their names

dace/frontend/python/replacements.py

@@ -282,26 +282,37 @@ def _numpy_full(pv: ProgramVisitor,
                 sdfg: SDFG,
                 state: SDFGState,
                 shape: Shape,
-                fill_value: Union[sp.Expr, Number],
+                fill_value: Union[sp.Expr, Number, data.Scalar],
                 dtype: dace.typeclass = None):
     """ Creates and array of the specified shape and initializes it with
         the fill value.
     """
+    is_data = False
     if isinstance(fill_value, (Number, np.bool_)):
         vtype = dtypes.DTYPE_TO_TYPECLASS[type(fill_value)]
     elif isinstance(fill_value, sp.Expr):
         vtype = _sym_type(fill_value)
     else:
-        raise mem_parser.DaceSyntaxError(pv, None, "Fill value {f} must be a number!".format(f=fill_value))
+        is_data = True
+        vtype = sdfg.arrays[fill_value].dtype
     dtype = dtype or vtype
     name, _ = sdfg.add_temp_transient(shape, dtype)
 
-    state.add_mapped_tasklet(
-        '_numpy_full_', {"__i{}".format(i): "0: {}".format(s)
-                         for i, s in enumerate(shape)}, {},
-        "__out = {}".format(fill_value),
-        dict(__out=dace.Memlet.simple(name, ",".join(["__i{}".format(i) for i in range(len(shape))]))),
-        external_edges=True)
+    if is_data:
+        state.add_mapped_tasklet(
+            '_numpy_full_', {"__i{}".format(i): "0: {}".format(s)
+                            for i, s in enumerate(shape)},
+            dict(__inp=dace.Memlet(data=fill_value, subset='0')),
+            "__out = __inp",
+            dict(__out=dace.Memlet.simple(name, ",".join(["__i{}".format(i) for i in range(len(shape))]))),
+            external_edges=True)
+    else:
+        state.add_mapped_tasklet(
+            '_numpy_full_', {"__i{}".format(i): "0: {}".format(s)
+                            for i, s in enumerate(shape)}, {},
+            "__out = {}".format(fill_value),
+            dict(__out=dace.Memlet.simple(name, ",".join(["__i{}".format(i) for i in range(len(shape))]))),
+            external_edges=True)
 
     return name
 

dace/libraries/mpi/nodes/__init__.py

@@ -10,5 +10,6 @@
 from .reduce import Reduce
 from .allreduce import Allreduce
 from .allgather import Allgather
+from .alltoall import Alltoall
 from .dummy import Dummy
 from .redistribute import Redistribute

dace/libraries/mpi/nodes/alltoall.py

@@ -0,0 +1,84 @@
+# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.
+import dace.library
+import dace.properties
+import dace.sdfg.nodes
+from dace.transformation.transformation import ExpandTransformation
+from .. import environments
+from dace.libraries.mpi.nodes.node import MPINode
+
+
+@dace.library.expansion
+class ExpandAlltoallMPI(ExpandTransformation):
+
+    environments = [environments.mpi.MPI]
+
+    @staticmethod
+    def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
+        (inbuffer, in_count_str), (outbuffer, out_count_str) = node.validate(parent_sdfg, parent_state)
+        in_mpi_dtype_str = dace.libraries.mpi.utils.MPI_DDT(inbuffer.dtype.base_type)
+        out_mpi_dtype_str = dace.libraries.mpi.utils.MPI_DDT(outbuffer.dtype.base_type)
+
+        if inbuffer.dtype.veclen > 1:
+            raise (NotImplementedError)
+
+        comm = "MPI_COMM_WORLD"
+        if node.grid:
+            comm = f"__state->{node.grid}_comm"
+
+        code = f"""
+            int size;
+            MPI_Comm_size({comm}, &size);
+            int sendrecv_amt = {in_count_str} / size;
+            MPI_Alltoall(_inbuffer, sendrecv_amt, {in_mpi_dtype_str}, \
+                        _outbuffer, sendrecv_amt, {out_mpi_dtype_str}, \
+                        {comm});
+            """
+        tasklet = dace.sdfg.nodes.Tasklet(node.name,
+                                          node.in_connectors,
+                                          node.out_connectors,
+                                          code,
+                                          language=dace.dtypes.Language.CPP)
+        return tasklet
+
+
+@dace.library.node
+class Alltoall(MPINode):
+
+    # Global properties
+    implementations = {
+        "MPI": ExpandAlltoallMPI,
+    }
+    default_implementation = "MPI"
+
+    grid = dace.properties.Property(dtype=str, allow_none=True, default=None)
+
+    def __init__(self, name, grid=None, *args, **kwargs):
+        super().__init__(name, *args, inputs={"_inbuffer"}, outputs={"_outbuffer"}, **kwargs)
+        self.grid = grid
+
+    def validate(self, sdfg, state):
+        """
+        :return: A three-tuple (buffer, root) of the three data descriptors in the
+                 parent SDFG.
+        """
+
+        inbuffer, outbuffer = None, None
+        for e in state.out_edges(self):
+            if e.src_conn == "_outbuffer":
+                outbuffer = sdfg.arrays[e.data.data]
+        for e in state.in_edges(self):
+            if e.dst_conn == "_inbuffer":
+                inbuffer = sdfg.arrays[e.data.data]
+
+        in_count_str = "XXX"
+        out_count_str = "XXX"
+        for _, src_conn, _, _, data in state.out_edges(self):
+            if src_conn == '_outbuffer':
+                dims = [str(e) for e in data.subset.size_exact()]
+                out_count_str = "*".join(dims)
+        for _, _, _, dst_conn, data in state.in_edges(self):
+            if dst_conn == '_inbuffer':
+                dims = [str(e) for e in data.subset.size_exact()]
+                in_count_str = "*".join(dims)
+
+        return (inbuffer, in_count_str), (outbuffer, out_count_str)

dace/libraries/mpi/nodes/bcast.py

@@ -42,11 +42,16 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
         if isinstance(buffer, dace.data.Scalar):
             ref = "&"
 
+        init = ""
         comm = "MPI_COMM_WORLD"
         if node.grid:
             comm = f"__state->{node.grid}_comm"
+        elif node.fcomm:
+            init = f"MPI_Comm __comm = MPI_Comm_f2c({node.fcomm});"
+            comm = "__comm"
 
         code = f"""
+            {init}
             MPI_Bcast({ref}_inbuffer, {count_str}, {mpi_dtype_str}, _root, {comm});
             _outbuffer = _inbuffer;"""
         tasklet = dace.sdfg.nodes.Tasklet(node.name,
@@ -67,10 +72,12 @@ class Bcast(MPINode):
     default_implementation = "MPI"
 
     grid = dace.properties.Property(dtype=str, allow_none=True, default=None)
+    fcomm = dace.properties.Property(dtype=str, allow_none=True, default=None)
 
-    def __init__(self, name, grid=None, *args, **kwargs):
+    def __init__(self, name, grid=None, fcomm=None, *args, **kwargs):
         super().__init__(name, *args, inputs={"_inbuffer", "_root"}, outputs={"_outbuffer"}, **kwargs)
         self.grid = grid
+        self.fcomm = fcomm
 
     def validate(self, sdfg, state):
         """

dace/libraries/mpi/nodes/irecv.py

@@ -20,6 +20,11 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
 
         if buffer.dtype.veclen > 1:
             raise NotImplementedError
+
+        comm = "MPI_COMM_WORLD"
+        if node.grid:
+            comm = f"__state->{node.grid}_comm"
+
         code = ""
         if ddt is not None:
             code = f"""static MPI_Datatype newtype;
@@ -33,7 +38,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
             mpi_dtype_str = "newtype"
             count_str = "1"
         buffer_offset = 0  #this is here because the frontend already changes the pointer
-        code += f"MPI_Irecv(_buffer, {count_str}, {mpi_dtype_str}, _src, _tag, MPI_COMM_WORLD, _request);"
+        code += f"MPI_Irecv(_buffer, {count_str}, {mpi_dtype_str}, int(_src), int(_tag), {comm}, _request);"
         if ddt is not None:
             code += f"""// MPI_Type_free(&newtype);
             """
@@ -58,8 +63,11 @@ class Irecv(MPINode):
     }
     default_implementation = "MPI"
 
-    def __init__(self, name, *args, **kwargs):
+    grid = dace.properties.Property(dtype=str, allow_none=True, default=None)
+
+    def __init__(self, name, grid=None, *args, **kwargs):
         super().__init__(name, *args, inputs={"_src", "_tag"}, outputs={"_buffer", "_request"}, **kwargs)
+        self.grid = grid
 
     def validate(self, sdfg, state):
         """

dace/libraries/mpi/nodes/isend.py

@@ -20,6 +20,10 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
 
         if buffer.dtype.veclen > 1:
             raise NotImplementedError
+
+        comm = "MPI_COMM_WORLD"
+        if node.grid:
+            comm = f"__state->{node.grid}_comm"
 
         code = ""
 
@@ -40,7 +44,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
             mpi_dtype_str = "newtype"
             count_str = "1"
         buffer_offset = 0
-        code += f"MPI_Isend(&(_buffer[{buffer_offset}]), {count_str}, {mpi_dtype_str}, _dest, _tag, MPI_COMM_WORLD, _request);"
+        code += f"MPI_Isend(&(_buffer[{buffer_offset}]), {count_str}, {mpi_dtype_str}, int(_dest), int(_tag), {comm}, _request);"
         if ddt is not None:
             code += f"""// MPI_Type_free(&newtype);
             """
@@ -69,13 +73,12 @@ class Isend(MPINode):
     }
     default_implementation = "MPI"
 
-    # Object fields
-    n = dace.properties.SymbolicProperty(allow_none=True, default=None)
-
+    grid = dace.properties.Property(dtype=str, allow_none=True, default=None)
     nosync = dace.properties.Property(dtype=bool, default=False, desc="Do not sync if memory is on GPU")
 
-    def __init__(self, name, *args, **kwargs):
+    def __init__(self, name, grid=None, *args, **kwargs):
         super().__init__(name, *args, inputs={"_buffer", "_dest", "_tag"}, outputs={"_request"}, **kwargs)
+        self.grid = grid
 
     def validate(self, sdfg, state):
         """
@@ -94,10 +97,11 @@ def validate(self, sdfg, state):
             if e.src_conn == "_request":
                 req = sdfg.arrays[e.data.data]
 
-        if dest.dtype.base_type != dace.dtypes.int32:
-            raise ValueError("Source must be an integer!")
-        if tag.dtype.base_type != dace.dtypes.int32:
-            raise ValueError("Tag must be an integer!")
+        # TODO: Should we expect any integer type here and cast to int32 later?. Investigate further in the future.
+        # if dest.dtype.base_type != dace.dtypes.int32:
+        #     raise ValueError("Destination must be an integer!")
+        # if tag.dtype.base_type != dace.dtypes.int32:
+        #     raise ValueError("Tag must be an integer!")
 
         count_str = "XXX"
         for _, _, _, dst_conn, data in state.in_edges(self):