mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/ops/operations/comm_ops.py

@@ -51,17 +51,12 @@ class ReduceOp:
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with multiple devices.
 
         >>> import numpy as np
@@ -94,7 +89,7 @@ class ReduceOp:
 
 def check_collective_target_dtype(data_name, data_dtype, prim_name):
     """Check if data type is valid."""
-    default_target_dtypes = (mstype.int8, mstype.uint8, mstype.int32, mstype.float16, mstype.bfloat16, mstype.float32)
+    default_target_dtypes = (mstype.int8, mstype.int32, mstype.float16, mstype.float32, mstype.bfloat16)
     gpu_target_dtypes = (mstype.bool_, mstype.int8, mstype.int32, mstype.int64, mstype.uint32, mstype.uint64,
                          mstype.float16, mstype.float32, mstype.float64)
 
@@ -113,15 +108,16 @@ def check_hcom_group_valid(group, prim_name=None):
 
 class AllReduce(Primitive):
     """
-    Reduces the tensor data across all devices in such a way that all devices will get the same final result.
+    Reduces tensors across all devices in such a way that all devices will get the same final result,
+    returns the tensor which is all reduced.
 
     Note:
         The tensors must have the same shape and format in all processes of the collection.
 
     Args:
-        op (str): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
+        op (str, optional): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
                   On the CPU, only 'sum' is supported. Default: ``ReduceOp.SUM`` .
-        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
                   means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
 
     Inputs:
@@ -132,8 +128,8 @@ class AllReduce(Primitive):
         The contents depend on the specified operation.
 
     Raises:
-        TypeError: If any of `op` and `group` is not a str,
-                   or fusion is not an integer, or the input's dtype is bool.
+        TypeError: If any of `op` and `group` is not a str or the input's dtype is bool.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -142,17 +138,12 @@ class AllReduce(Primitive):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with 2 devices.
 
         >>> import numpy as np
@@ -160,7 +151,7 @@ class AllReduce(Primitive):
         >>> from mindspore import Tensor
         >>> from mindspore.ops import ReduceOp
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>>
         >>> init()
         >>> class Net(nn.Cell):
@@ -180,36 +171,125 @@ class AllReduce(Primitive):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - AllReduce
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#allreduce>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#allreduce>`_
 
     """
 
     @prim_attr_register
     def __init__(self, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize AllReduce."""
+        self.group = _get_group(group)
         if not isinstance(op, type(ReduceOp.SUM)):
             raise TypeError(f"For '{self.name}', the 'op' must be str, but got {type(op).__name__}.")
-        if not isinstance(_get_group(group), str):
+        if not isinstance(self.group, str):
             raise TypeError(f"For '{self.name}', the 'group' must be str, "
-                            f"but got {type(_get_group(group)).__name__}.")
-        check_hcom_group_valid(group, prim_name=self.name)
+                            f"but got {type(self.group).__name__}.")
+        check_hcom_group_valid(self.group, prim_name=self.name)
         self.op = op
-        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('index', 0)
         self.add_prim_attr('no_eliminate', True)
 
 
+class Reduce(PrimitiveWithInfer):
+    """
+    Reduces tensors across the processes in the specified communication group, sends the result
+    to the target dest_rank(local rank), and returns the tensor which is sent to the target process.
+
+    Note:
+        Only process with destination rank receives the reduced output.
+        Support PyNative mode and Graph mode, but Graph mode only supports scenes with a graph compilation level of O0.
+        Other processes only get a tensor with shape [1], which has no mathematical meaning.
+
+    Args:
+        dest_rank (int): The target process(local rank) in the specific group that receives the reduced output.
+        op (str, optional): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
+                  On the CPU, only 'sum' is supported. Default: ``ReduceOp.SUM`` .
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
+                  means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Outputs:
+        Tensor. Return the tensor in the specific rank of the process after reduction.
+        The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Raises:
+        TypeError: If the type of the first input parameter is not Tensor,
+                or any of `op` and `group` is not a str.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method without any third-party
+            or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 4 devices.
+
+        >>> from mindspore import ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> import numpy as np
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class ReduceNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(Net, self).__init__()
+        >>>         self.reduce = ops.Reduce(dest_rank=1)
+        >>>
+        >>>     def construct(self, x):
+        >>>         out = self.reduce(x)
+        >>>         return out
+        >>> input = Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> net = ReduceNet()
+        >>> output = net(input)
+        >>> print(output)
+        Process with rank 1: [[4. 4. 4. 4. 4. 4. 4. 4.]
+                             [4. 4. 4. 4. 4. 4. 4. 4.]],
+        Other proesses: [0.].
+    """
+
+    @prim_attr_register
+    def __init__(self, dest_rank, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
+        validator.check_value_type('group', _get_group(group), (str,), self.name)
+        validator.check_value_type('op', op, (type(ReduceOp.SUM),), self.name)
+        self.dest_rank = dest_rank
+        self.op = op
+        self.group = _get_group(group)
+        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('dest_rank', dest_rank)
+
+    def infer_shape(self, x_shape):
+        # The process with dest_rank returns the reduced output.
+        # Other processes only gets a tensor with shape [1], which has no mathematical meaning.
+        if self.dest_rank == get_rank():
+            return x_shape
+        return [1]
+
+    def infer_dtype(self, x_dtype):
+        return x_dtype
+
+
 class AllGather(PrimitiveWithInfer):
     """
-    Gathers tensors from the specified communication group.
+    Gathers tensors from the specified communication group and returns the tensor which is all gathered.
 
     Note:
         - The tensors must have the same shape and format in all processes of the collection.
-        - Currently only supports GRAPH_MODE and it should be called in Cell.
 
     Args:
-        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
             means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
 
     Inputs:
@@ -223,6 +303,7 @@ class AllGather(PrimitiveWithInfer):
         TypeError: If `group` is not a str.
         ValueError: If the local rank id of the calling process in the group
                     is larger than the group's rank size.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -231,22 +312,17 @@ class AllGather(PrimitiveWithInfer):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with 2 devices.
 
         >>> import numpy as np
         >>> import mindspore as ms
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import mindspore.nn as nn
         >>> from mindspore.communication import init
         >>> from mindspore import Tensor
@@ -272,19 +348,20 @@ class AllGather(PrimitiveWithInfer):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - AllGather
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#allgather>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#allgather>`_
 
     """
 
     @prim_attr_register
     def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize AllGather."""
-        validator.check_value_type('group', _get_group(group), (str,), self.name)
-        self.rank = get_rank(_get_group(group))
-        self.rank_size = get_group_size(_get_group(group))
+        self.group = _get_group(group)
+        validator.check_value_type('group', self.group, (str,), self.name)
+        self.rank = get_rank(self.group)
+        self.rank_size = get_group_size(self.group)
         validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
-        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('mean_flag', False)
         self.add_prim_attr('no_eliminate', True)
@@ -427,7 +504,8 @@ class _HostAllGather(PrimitiveWithInfer):
 
 class ReduceScatter(Primitive):
     r"""
-    Reduces and scatters tensors from the specified communication group.
+    Reduces and scatters tensors from the specified communication group
+    and returns the tensor which is reduced and scattered.
 
     Note:
         The tensors must have the same shape and format in all processes of the collection.
@@ -448,6 +526,7 @@ class ReduceScatter(Primitive):
     Raises:
         TypeError: If any of operation and group is not a string.
         ValueError: If the first dimension of the input cannot be divided by the rank_size.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
         ``Ascend`` ``GPU``
@@ -456,17 +535,12 @@ class ReduceScatter(Primitive):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with 2 devices.
 
         >>> import mindspore as ms
@@ -474,7 +548,7 @@ class ReduceScatter(Primitive):
         >>> from mindspore.communication import init
         >>> from mindspore.ops import ReduceOp
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import numpy as np
         >>>
         >>> ms.set_context(mode=ms.GRAPH_MODE)
@@ -498,7 +572,7 @@ class ReduceScatter(Primitive):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - ReduceScatter
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#reducescatter>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#reducescatter>`_
 
     """
 
@@ -506,11 +580,12 @@ class ReduceScatter(Primitive):
     def __init__(self, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize ReduceScatter."""
         validator.check_value_type('op', op, (type(ReduceOp.SUM),), self.name)
-        validator.check_value_type('group', _get_group(group), (str,), self.name)
+        self.group = _get_group(group)
+        validator.check_value_type('group', self.group, (str,), self.name)
         self.op = op
-        self.rank_size = get_group_size(_get_group(group))
+        self.rank_size = get_group_size(self.group)
         self.add_prim_attr('rank_size', self.rank_size)
-        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('no_eliminate', True)
 
@@ -577,12 +652,12 @@ class Broadcast(PrimitiveWithInfer):
         The tensors must have the same shape and format in all processes of the collection.
 
     Args:
-        root_rank (int): Source rank. Required in all processes except the one
-                   that is sending the data.
+        root_rank (int): Specifies the rank(global rank) of the process that broadcast the tensor.
+            And only process `root_rank` will broadcast the tensor.
         group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` .
 
     Inputs:
-        - **input_x** (tuple[Tensor]) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
 
     Outputs:
         tuple[Tensor], Tensor has the same shape of the input, i.e., :math:`(x_1, x_2, ..., x_R)`.
@@ -598,24 +673,19 @@ class Broadcast(PrimitiveWithInfer):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
-            This example should be run with multiple devices.
+            This example should be run with 2 devices.
 
         >>> import mindspore as ms
         >>> from mindspore import Tensor
         >>> from mindspore.communication import init
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import numpy as np
         >>>
         >>> ms.set_context(mode=ms.GRAPH_MODE)
@@ -638,7 +708,7 @@ class Broadcast(PrimitiveWithInfer):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - Broadcast
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#broadcast>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#broadcast>`_
 
     """
 
@@ -651,16 +721,6 @@ class Broadcast(PrimitiveWithInfer):
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('no_eliminate', True)
 
-    def infer_shape(self, x_shape):
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        if not isinstance(x_dtype, tuple):
-            raise TypeError(f"For '{self.name}', the 'input_x' must be a tuple, but got {type(x_dtype).__name__}!")
-        for _ele in x_dtype:
-            check_collective_target_dtype('tuple input_x', _ele, self.name)
-        return x_dtype
-
 
 class _AllSwap(PrimitiveWithCheck):
     """
@@ -728,11 +788,11 @@ class NeighborExchange(Primitive):
         The user needs to preset
         communication environment variables before running the following example, please check the details on the
         official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.ops.primitive.html#communication-operator>`_.
+        <https://www.mindspore.cn/docs/en/master/api_python/mindspore.ops.primitive.html#communication-operator>`_.
 
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to.
@@ -758,7 +818,7 @@ class NeighborExchange(Primitive):
         >>> from mindspore import Tensor
         >>> from mindspore.communication import init
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import numpy as np
         >>> class Net(nn.Cell):
         ...     def __init__(self):
@@ -781,7 +841,7 @@ class NeighborExchange(Primitive):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - NeighborExchange
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#neighborexchange>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#neighborexchange>`_
 
     """
 
@@ -814,7 +874,7 @@ class AlltoAll(PrimitiveWithInfer):
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         split_count (int): On each process, divide blocks into split_count number.
@@ -843,17 +903,12 @@ class AlltoAll(PrimitiveWithInfer):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with 8 devices.
 
         >>> import os
@@ -861,7 +916,7 @@ class AlltoAll(PrimitiveWithInfer):
         >>> from mindspore import Tensor
         >>> from mindspore.communication import init
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import numpy as np
         >>> class Net(nn.Cell):
         ...     def __init__(self):
@@ -883,7 +938,7 @@ class AlltoAll(PrimitiveWithInfer):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - AlltoAll
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#alltoall>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#alltoall>`_
 
     """
 
@@ -905,20 +960,20 @@ class AlltoAll(PrimitiveWithInfer):
         if self.split_count != rank_size:
             raise ValueError(f"For '{self.name}', the 'split_count' must be equal to 'rank_size', "
                              f"but got 'split_count': {self.split_count}, 'rank_size': {rank_size}.")
-        if x_shape[self.split_dim] % self.split_count != 0:
-            raise ValueError(f"For '{self.name}', the 'split_count' must be divisible by 'rank_size', "
-                             f"but got 'split_count' {self.split_count}, 'rank_size' {x_shape[self.split_dim]}.")
-        x_shape[self.concat_dim] = x_shape[self.concat_dim] * self.split_count
-        x_shape[self.split_dim] = int(x_shape[self.split_dim] / self.split_count)
+        if x_shape[self.split_dim] >= 0 and x_shape[self.split_dim] % self.split_count != 0:
+            raise ValueError(f"For '{self.name}', the 'x_shape[self.split_dim]' must be divisible by 'split_count', "
+                             f"but got 'x_shape[self.split_dim]' {x_shape[self.split_dim]}, "
+                             f"'split_count' {self.split_count}.")
+        if x_shape[self.concat_dim] >= 0:
+            x_shape[self.concat_dim] = x_shape[self.concat_dim] * self.split_count
+        if x_shape[self.split_dim] >= 0:
+            x_shape[self.split_dim] = int(x_shape[self.split_dim] / self.split_count)
         return x_shape
 
     def infer_dtype(self, x_dtype):
         check_collective_target_dtype('x', x_dtype, self.name)
         return x_dtype
 
-    def __call__(self, tensor):
-        raise NotImplementedError
-
 
 class NeighborExchangeV2(Primitive):
     r"""
@@ -931,7 +986,7 @@ class NeighborExchangeV2(Primitive):
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to. 8 rank_ids represents 8 directions, if one
@@ -967,24 +1022,19 @@ class NeighborExchangeV2(Primitive):
         .. note::
             Before running the following examples, you need to configure the communication environment variables.
 
-            For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `rank table Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
-
-            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
-            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
-
             This example should be run with 2 devices.
 
         >>> import os
         >>> import mindspore as ms
         >>> from mindspore.communication import init
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> import numpy as np
         >>>
         >>> class Net0(nn.Cell):
@@ -998,7 +1048,6 @@ class NeighborExchangeV2(Primitive):
         ...     def construct(self, x):
         ...         out = self.neighbor_exchangev2(x)
         ...         return out
-        >>>
         ... class Net1(nn.Cell):
         ...     def __init__(self):
         ...         super(Net1, self).__init__()
@@ -1028,7 +1077,7 @@ class NeighborExchangeV2(Primitive):
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - NeighborExchangeV2
-          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#neighborexchangev2>`_
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#neighborexchangev2>`_
 
     """
 
@@ -1057,6 +1106,414 @@ class NeighborExchangeV2(Primitive):
         raise NotImplementedError
 
 
+class CollectiveScatter(Primitive):
+    r"""
+    Scatter tensor evently across the processes in the specified communication group.
+
+    Note:
+        The interface behavior only support Tensor input and scatter evenly.
+        Only the tensor in process `src_rank` (global rank) will do scatter.
+
+    Args:
+        src_rank (int, optional): Specifies the rank of the process that send the tensor.
+            And only process `src_rank` will send the tensor.
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``.
+
+    Inputs:
+        - **input_x** (Tensor) - The input tensor to be scattered. The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Outputs:
+        Tensor, the shape of output is :math:`(x_1/src\_rank, x_2, ..., x_R)`. The dimension 0 of data is equal to
+        the dimension of input tensor divided by `src`, and the other dimension keep the same.
+
+    Raises:
+        TypeError: If `group` is not a str.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+        ValueError: If the local rank id of the calling process in the group
+            is larger than the group's rank size.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor
+        >>> from mindspore.communication.management import init, get_rank
+        >>> from mindspore import ops
+        >>> # Launch 2 processes.
+        >>> init()
+        >>> class CollectiveScatterNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(CollectiveScatter, self).__init__()
+        >>>         self.collective_scatter = ops.CollectiveScatter(src_rank=0)
+        >>>
+        >>>     def construct(self, x):
+        >>>         return self.collective_scatter(x)
+        >>>
+        >>> input = Tensor(np.arange(8).reshape([4, 2]).astype(np.float32))
+        >>> net = CollectiveScatterNet()
+        >>> output = net(input)
+        >>> print(output)
+        Process with rank 0: [[0. 1.],
+                              [2. 3.]]
+        Process with rank 1: [[4. 5.],
+                              [6. 7.]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - CollectiveScatter
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#reducescatter>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, src_rank=0, group=GlobalComm.WORLD_COMM_GROUP):
+        validator.check_value_type('group', _get_group(group), (str,), self.name)
+        self.rank_id = get_rank(_get_group(group))
+        self.src_rank = src_rank
+        self.rank_size = get_group_size(_get_group(group))
+        validator.check('rank', self.rank_id, 'rank_size', self.rank_size, validator.LT, self.name)
+        self.add_prim_attr('rank_id', self.rank_id)
+        self.add_prim_attr('src_rank', self.src_rank)
+        self.add_prim_attr('rank_size', self.rank_size)
+        self.add_prim_attr('group', _get_group(group))
+
+
+class CollectiveGather(Primitive):
+    r"""
+    Gathers tensors from the specified communication group. The operation will gather the tensor
+    from processes according to dimension 0.
+
+    Note:
+        Only the tensor in process `dest_rank` (global rank) will keep the gathered tensor. The other process
+        will keep a tensor with shape [1], which has no mathematical meaning.
+
+    Args:
+        dest_rank(int): Specifies the rank of the process that receive the tensor.
+            And only process `dest_rank` will receive the gathered tensor.
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``.
+
+    Inputs:
+        - **input_x** (Tensor) - The tensor to be gathered. The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Outputs:
+        Tensor, the shape of output is :math:`(\sum x_1, x_2, ..., x_R)`. The dimension 0 of data is equal to
+        sum of the dimension of input tensor, and the other dimension keep the same.
+
+    Raises:
+        TypeError: If `group` is not a str.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+        ValueError: If the local rank id of the calling process in the group
+                    is larger than the group's rank size.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 4 devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>> # Launch 2 processes.
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> init()
+        >>> class CollectiveGatherNet(nn.Cell):
+        ...     def __init__(self):
+        ...         super(CollectiveGatherNet, self).__init__()
+        ...         self.collective_gather = ops.CollectiveGather(dest_rank=0)
+        ...
+        ...     def construct(self, x):
+        ...         return self.collective_gather(x)
+        ...
+        >>> input = Tensor(np.arange(4).reshape([2, 2]).astype(np.float32))
+        >>> net = CollectiveGatherNet()
+        >>> output = net(input)
+        >>> print(output)
+        Process with rank 0: [[0. 1.],
+                              [2. 3.],
+                              [0. 1.],
+                              [2. 3.]]
+        Process with rank 1: [0.]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - CollectiveGather
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#collectivegather>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, dest_rank, group=GlobalComm.WORLD_COMM_GROUP):
+        """Initialize Gather."""
+        validator.check_value_type('group', _get_group(group), (str,), self.name)
+        self.rank_id = get_rank(_get_group(group))
+        self.dest_rank = dest_rank
+        self.rank_size = get_group_size(_get_group(group))
+        validator.check('rank', self.rank_id, 'rank_size', self.rank_size, validator.LT, self.name)
+        self.add_prim_attr('rank_size', self.rank_size)
+        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('dest_rank', self.dest_rank)
+        self.add_prim_attr('rank_id', self.rank_id)
+
+
+class Barrier(PrimitiveWithInfer):
+    """
+    Synchronizes all processes in the specified group. Once the process call this operation, it will be blocked until
+    all processes call this operation. After all processes finish calling the operations, the blocked processes
+    will be waken and continue their task.
+
+    Args:
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``.
+
+    Raises:
+        TypeError: If `group` is not a str.
+        RuntimeError: If backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>> # Launch 4 processes.
+        >>> init()
+        >>> class BarrierNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(BarrierNet, self).__init__()
+        >>>         self.barrier = ops.Barrier()
+        >>>
+        >>>     def construct(self):
+        >>>         self.barrier()
+        >>> net = BarrierNet()
+        >>> net()
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - Barrier
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#barrier>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
+        self.group = group
+        self.add_prim_attr("side_effect_mem", True)
+
+    def infer_shape(self):
+        return [1]
+
+    def infer_dtype(self):
+        return mstype.float32
+
+
+class Send(PrimitiveWithInfer):
+    """
+    Send tensors to the specified dest_rank.
+
+    Note:
+        Send and Receive must be used in combination and have same sr_tag.
+
+    Args:
+        sr_tag (int): The tag to identify the send/recv message. The message will
+                      be received by the Receive op with the same "sr_tag".
+        dest_rank (int): A required integer identifying the destination rank.
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``.
+        group_back (str, optional): The communication group for backpropagation.
+                                    Default: ``GlobalComm.WORLD_COMM_GROUP``.
+
+    Inputs:
+        - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Raises:
+        TypeError: If `group` is not a str.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+        ValueError: If the local rank id of the calling process in the group
+                    is larger than the group's rank size.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>>
+        >>> init()
+        >>> class SendNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(SendNet, self).__init__()
+        >>>         self.depend = ops.Depend()
+        >>>         self.send = ops.Send(st_tag=0, dest_rank=8, group="hccl_world_group")
+        >>>
+        >>>     def construct(self, x):
+        >>>         out = self.depend(x, self.send(x))
+        >>>         return out
+        >>>
+        >>> input_ = Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> net = Net()
+        >>> output = net(input_)
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - Send
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#send>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, sr_tag, dest_rank, group=GlobalComm.WORLD_COMM_GROUP, group_back=GlobalComm.WORLD_COMM_GROUP):
+        self.rank = dest_rank
+        self.sr_tag = sr_tag
+        self.group = _get_group(group)
+        self.add_prim_attr("no_eliminate", True)
+
+    def infer_shape(self, x_shape):
+        self.add_prim_attr("shape", x_shape)
+        return x_shape
+
+    def infer_dtype(self, x_dtype):
+        return x_dtype
+
+
+class Receive(PrimitiveWithInfer):
+    """
+    Receive tensors from src_rank.
+
+    Note:
+        Send and Receive must be used in combination and have same sr_tag.
+
+    Args:
+        sr_tag (int): A required integer identifying the send/recv message tag. The message will
+                      will be send by the Send op with the same "sr_tag".
+        src_rank (int): A required integer identifying the source rank.
+        shape (list[int]): A required list identifying the shape of the tensor to be received.
+        dtype (Type): A required Type identifying the type of the tensor to be received. The supported types:
+                       int8/int16/int32/float16/float32.
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``.
+        group_back (str, optional): The communication group for backpropagation.
+                                    Default: ``GlobalComm.WORLD_COMM_GROUP``.
+
+    Outputs:
+        Tensor, output has the same shape as the Tensor sent by `Send` operation.
+
+    Raises:
+        TypeError: If `group` is not a str.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+        ValueError: If the local rank id of the calling process in the group
+                    is larger than the group's rank size.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init
+        >>> from mindspore import Tensor
+        >>> from mindspore import ops
+        >>>
+        >>> init()
+        >>> class ReceiveNet(nn.Cell):
+        >>>     def __init__(self):
+        >>>         super(ReceiveNet, self).__init__()
+        >>>         self.recv = ops.Receive(sr_tag=0, src_rank=0, shape=[2, 8], dtype=ms.float32,
+        >>>                               group="hccl_world_group")
+        >>>
+        >>>     def construct(self):
+        >>>         out = self.recv()
+        >>>         return out
+        >>>
+        >>> net = Net()
+        >>> output = net()
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - Receive
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#receive>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, sr_tag, src_rank, shape, dtype, group=GlobalComm.WORLD_COMM_GROUP,
+                 group_back=GlobalComm.WORLD_COMM_GROUP):
+        self.rank = src_rank
+        self.tag = sr_tag
+        self.shape = shape
+        self.dtype = dtype
+        self.group = _get_group(group)
+        self.add_prim_attr("no_eliminate", True)
+        valid_type = [mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16,
+                      mstype.int8, mstype.int16, mstype.int32, mstype.int64,
+                      mstype.uint8, mstype.uint16, mstype.uint32, mstype.uint64]
+        args = {"dtype": dtype}
+        validator.check_scalar_or_tensor_types_same(args, valid_type, self.name)
+
+    def infer_shape(self, x_shape=None):
+        return self.get_attr_dict()['shape']
+
+    def infer_dtype(self, x_dtype=None):
+        return self.get_attr_dict()['dtype']
+
+
 class _MirrorOperator(PrimitiveWithInfer):
     """
     Auto parallel virtual operator. Do nothing in forward, do all reduce and mean in backward. It is only for
@@ -1306,8 +1763,223 @@ class _GetTensorSlice(PrimitiveWithInfer):
         """Initialize _GetTensorSlice."""
         self.add_prim_attr('order_enforce_skip', True)
 
-    def infer_value(self, x, dev_mat, tensor_map):
+    def infer_value(self, x, dev_mat, tensor_map, slice_shape, full_shape):
         from mindspore.parallel._tensor import _load_tensor
         validator.check_value_type("dev_mat", dev_mat, [tuple], self.name)
         validator.check_value_type("tensor_map", tensor_map, [tuple], self.name)
-        return Tensor(_load_tensor(x, dev_mat, tensor_map), x.dtype)
+        tensor_slice = _load_tensor(x, dev_mat, tensor_map, full_shape)
+        if tensor_slice.shape != slice_shape:
+            tensor_slice = tensor_slice.reshape(slice_shape)
+        return Tensor(tensor_slice, x.dtype)
+
+
+class BatchISendIRecv(PrimitiveWithInfer):
+    """
+    Batch send and recv tensors asynchronously.
+
+    Note:
+        - The ``isend`` and ``irecv`` in ``op_types`` between ranks need to match each other.
+        - ``isend`` and ``irecv`` in a batch can only be used in the same communication group.
+
+    Args:
+        op_types(Union[tuple[str], list[str]]): "isend" or "irecv" to indicate the order and number of communication.
+        remote_ranks(Union[tuple[int], list[int]]): src or dst rank that matches the op_types.
+        receive_shapes(Union[tuple[int], list[int]]): receive tensor shapes that matches "irecv" in op_types.
+        receive_types(Union[tuple[mindspore.dtype], list[mindspore.dtype]]): receive tensor dtype
+          that matches "irecv" in op_types.
+        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``, which
+          means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+
+    Inputs:
+        - **input_x** (Union[tuple[Tensor], list[Tensor], tuple(None)]) -
+          The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
+
+    Outputs:
+        tuple(Tensor). Output tensors is corresponding to ``op_types``:
+        At ``"isend"`` position, output tensor is a fake tensor with scalar, which has no meaning.
+        At ``"irecv"`` position, output tensor is a tensor received from remote end.
+
+
+    Raises:
+        TypeError: If ``group`` is not a str.
+        TypeError: If ``op_types``, ``receive_shapes``, ``receive_dtypes``, ``remote_ranks`` are not tuple or list.
+        ValueError: If the length of ``receive_shapes`` and ``receive_dtypes`` are not the same.
+        ValueError: If the length of ``op_types`` and ``remote_ranks`` are not the same.
+        RuntimeError: If the length of input tensors and ``"isend"`` count are not the same.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init, get_rank
+        >>> from mindspore import Tensor
+        >>>
+        >>> init()
+        >>> rank = get_rank()
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         if rank == 0:
+        ...             remote_rank = [1, 1]
+        ...         else:
+        ...             remote_rank = [0, 0]
+        ...         self.batchisendirecv = ops.BatchISendIRecv(("isend", "irecv"), remote_rank, [()], (ms.float32,))
+        ...
+        ...     def construct(self, x):
+        ...         if isinstance(x, Tensor):
+        ...             x = (x,)
+        ...         return self.batchisendirecv(x)
+        ...
+        >>> send_x = Tensor(rank + 1).astype(ms.float32)
+        >>> net = Net()
+        >>> output = net(send_x)
+        >>> print(output)
+        rank 0:
+        (Tensor(shape=[], dtype=Float32, value= 0), Tensor(shape=[], dtype=Float32, value= 2))
+        rank 1:
+        (Tensor(shape=[], dtype=Float32, value= 0), Tensor(shape=[], dtype=Float32, value= 1))
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - BatchISendIRecv
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#allgather>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, op_types, remote_ranks, receive_shapes=None,
+                 receive_dtypes=None, group=GlobalComm.WORLD_COMM_GROUP):
+        if receive_shapes is None:
+            receive_shapes = ()
+        else:
+            validator.check_value_type("receive_shapes", receive_shapes, [tuple, list], self.name)
+
+        if receive_dtypes is None:
+            receive_dtypes = ()
+        else:
+            validator.check_value_type("receive_dtypes", receive_dtypes, [tuple, list], self.name)
+
+        validator.check_value_type("op_types", op_types, [tuple, list], self.name)
+        validator.check_value_type("remote_ranks", remote_ranks, [tuple, list], self.name)
+
+        if len(receive_shapes) != len(receive_dtypes):
+            raise ValueError("length of receive_shapes and receive_shapes must be the same, "
+                             f"but got receive_shapes: {len(receive_shapes)} "
+                             f" and receive_shapes: {receive_dtypes}")
+
+        if len(op_types) != len(remote_ranks):
+            raise ValueError("length of op_types and remote_ranks must be the same.")
+
+        if group is None:
+            group = GlobalComm.WORLD_COMM_GROUP
+        self.add_prim_attr('group', group)
+        self.add_prim_attr('op_types', op_types)
+        self.add_prim_attr('remote_ranks', remote_ranks)
+        self.add_prim_attr('receive_shapes', receive_shapes)
+        self.add_prim_attr('receive_dtypes', receive_dtypes)
+        self.add_prim_attr('no_eliminate', True)
+
+
+class AlltoAllV(PrimitiveWithInfer):
+    """
+    AllToAll which support uneven split.
+
+    Note:
+        - Only support flatten tensor as input. input tensor should be flattened and
+          concatenated before call this primitive.
+
+    Args:
+        send_numel_list(Union[tuple[int], list[int]]): split numel to scatter to different remote rank.
+        recv_numel_list(Union[tuple[int], list[int]]): split numel to gather from different remote rank.
+        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``, which
+          means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+        TODO:
+
+    Inputs:
+        - **input_x** (Tensor) - flatten tensor to scatter. The shape of tensor is :math:`(x_1)`.
+
+    Outputs:
+        Tensor. flattened and concatenated tensor gather from remote ranks.
+        If gather result is empty, it will return a Tensor with value 0, which has no actual meaning.
+
+    Raises:
+        TypeError: If 'send_numel_list'  or 'recv_numel_list' is not type of tuple and list.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+
+            Please see the `msrun start up
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import ops
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init, get_rank
+        >>> from mindspore import Tensor
+        >>>
+        >>> init()
+        >>> rank = get_rank()
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         if rank == 0:
+        ...             self.all_to_all = ops.AlltoAllV([1, 2], [1, 2])
+        ...         else:
+        ...             self.all_to_all = ops.AlltoAllV([2, 1], [2, 1])
+        ...
+        ...     def construct(self, x):
+        ...         return self.all_to_all(x)
+        ...
+        >>> if rank == 0:
+        >>>    send_tensor = Tensor([0, 1, 2.])
+        >>> elif rank == 1:
+        >>>    send_tensor = Tensor([3, 4, 5.])
+        >>> net = Net()
+        >>> output = net(send_tensor)
+        >>> print(output)
+        rank 0:
+        [0. 3. 4]
+        rank 1:
+        [1. 2. 5]
+
+    """
+
+    @prim_attr_register
+    def __init__(self, send_numel_list, recv_numel_list, group=None, split_sizes_empty=False):
+        validator.check_value_type("send_numel_list", send_numel_list, [tuple, list], self.name)
+        validator.check_value_type("recv_numel_list", recv_numel_list, [tuple, list], self.name)
+        self.group = GlobalComm.WORLD_COMM_GROUP if group is None else _get_group(group)
+        self.send_numel_list = send_numel_list
+        self.recv_numel_list = recv_numel_list
+        self.split_sizes_empty = split_sizes_empty
+        self.rank_size = get_group_size(self.group)
+
+        self.add_prim_attr('group', self.group)
+        self.add_prim_attr('send_numel_list', send_numel_list)
+        self.add_prim_attr('recv_numel_list', recv_numel_list)