chore: import upstream snapshot with attribution

python/paddle/distributed/auto_parallel/static/process_mesh_v2.py

@@ -0,0 +1,145 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from paddle.framework import core
+
+
+class ProcessMesh(core.ProcessMesh):
+    r"""
+    The class `ProcessMesh` describes the topology of logical processes.
+
+    Args:
+        mesh (list|numpy.array): an N-dimensional array describes the topology
+            of logical processes.
+        dim_names (list, optional): the i-th element of this list gives the name of the
+            i-th dimension.
+
+    Returns:
+        None
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> paddle.enable_static()
+
+            >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]])
+            >>> assert mesh.shape == [2, 3]
+            >>> assert mesh.process_ids == [2, 4, 5, 0, 1, 3]
+
+    """
+
+    def __init__(self, mesh, dim_names=None):
+        if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
+            raise ValueError(
+                'The mesh must be an instance of list or np.ndarray.'
+            )
+        if isinstance(mesh, list):
+            mesh = np.array(mesh)
+
+        self._mesh = mesh
+
+        self._shape = list(self._mesh.shape)
+
+        self._process_ids = self._mesh.flatten().tolist()
+        if not all(isinstance(p, int) for p in self._process_ids):
+            raise ValueError("All elements of the mesh must be integer")
+
+        if min(self._process_ids) < 0:
+            raise ValueError('All elements of the mesh must be >= 0.')
+
+        unique_process_ids = set(self._process_ids)
+        if len(unique_process_ids) != len(self._process_ids):
+            raise ValueError('All elements of the mesh must be unique.')
+
+        if dim_names is not None:
+            assert len(dim_names) == len(self._shape), (
+                "The length of dims_names must be same as the shape of the mesh."
+            )
+            self._dim_names = dim_names
+        else:
+            self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
+
+        # Follow the requirement for using pybind11
+        core.ProcessMesh.__init__(
+            self, self._shape, self._process_ids, self._dim_names
+        )
+
+    @property
+    def mesh(self):
+        return self._mesh
+
+
+def compute_compatible_process_mesh(process_meshes):
+    """Compute the compatible process mesh given a list of process meshes."""
+    if not process_meshes:
+        return None
+
+    def _compute_compatible_of_two_process_meshes(pm1, pm2):
+        if pm1 is None:
+            return True, pm2
+        if pm2 is None:
+            return True, pm1
+        if pm1 == pm2:
+            return True, pm1
+        if pm1.process_ids == pm2.process_ids:
+            if len(pm1.shape) >= len(pm2.shape):
+                return True, pm1
+            else:
+                return True, pm2
+        process_set1 = set(pm1.process_ids)
+        process_set2 = set(pm2.process_ids)
+        if process_set1.issubset(process_set2):
+            return True, pm2
+        if process_set2.issubset(process_set1):
+            return True, pm1
+        return False, None
+
+    compatible_result = None
+    for process_mesh in process_meshes:
+        (
+            compatible,
+            compatible_result,
+        ) = _compute_compatible_of_two_process_meshes(
+            compatible_result, process_mesh
+        )
+        if not compatible:
+            return None
+    if compatible_result.empty():
+        return None
+    if isinstance(compatible_result, core.ProcessMesh):
+        mesh = np.array(compatible_result.process_ids).reshape(
+            compatible_result.shape
+        )
+        return ProcessMesh(mesh, compatible_result.dim_names)
+    elif isinstance(compatible_result, ProcessMesh):
+        return ProcessMesh(compatible_result.mesh, compatible_result.dim_names)
+    else:
+        raise ValueError("Unrecognized ProcessMesh.")
+
+
+def merge_process_mesh(process_meshes):
+    """Merge a list of process meshes."""
+    merged_process_mesh = None
+    merged_process_ids = set()
+    for process_mesh in process_meshes:
+        if process_mesh is not None:
+            process_ids = set(process_mesh.process_ids)
+            merged_process_ids = merged_process_ids.union(process_ids)
+    if len(merged_process_ids) != 0:
+        merged_process_mesh = ProcessMesh(list(merged_process_ids))
+    return merged_process_mesh