chore: import upstream snapshot with attribution

paddle/phi/core/distributed/CMakeLists.txt

@@ -0,0 +1,34 @@
+add_subdirectory(check)
+add_subdirectory(store)
+add_subdirectory(auto_parallel)
+add_subdirectory(collective)
+
+set(DISTRIBUTED_COMMON_SRCS comm_context_manager.cc utils.cc)
+
+if(WITH_NCCL OR WITH_RCCL)
+  list(APPEND DISTRIBUTED_COMMON_SRCS comm_task_manager.cc)
+  list(APPEND DISTRIBUTED_COMMON_SRCS nccl_comm_context.cc nccl_comm_task.cc
+       nccl_tools.cc)
+  list(APPEND DISTRIBUTED_COMMON_SRCS nccl_config.cc)
+endif()
+
+if(WITH_GLOO)
+  list(APPEND DISTRIBUTED_COMMON_SRCS gloo_utils.cc gloo_comm_context.cc)
+endif()
+
+if(WITH_CUSTOM_DEVICE)
+  list(APPEND DISTRIBUTED_COMMON_SRCS xccl_comm_context.cc)
+endif()
+
+if(WITH_XPU_BKCL)
+  list(APPEND DISTRIBUTED_COMMON_SRCS bkcl_comm_context.cc)
+endif()
+
+if(WITH_FLAGCX)
+  list(APPEND DISTRIBUTED_COMMON_SRCS flagcx_tools.cc)
+  if(NOT WITH_XPU)
+    list(APPEND DISTRIBUTED_COMMON_SRCS flagcx_comm_context.cc)
+  endif()
+endif()
+
+collect_srcs(core_srcs SRCS ${DISTRIBUTED_COMMON_SRCS})

paddle/phi/core/distributed/auto_parallel/CMakeLists.txt

@@ -0,0 +1,16 @@
+proto_library(auto_parallel_proto SRCS auto_parallel.proto)
+
+collect_srcs(
+  core_srcs
+  SRCS
+  device_mesh.cc
+  process_mesh.cc
+  dist_attr.cc
+  dist_mapper.cc
+  dist_tensor.cc
+  dist_meta_tensor.cc
+  proto_helper.cc
+  placement_types.cc
+  inferspmd_utils.cc)
+
+add_subdirectory(reshard)

paddle/phi/core/distributed/auto_parallel/auto_parallel.proto

@@ -0,0 +1,183 @@
+/* Copyright (c) 2023 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. */
+
+syntax = "proto2";
+
+package phi.distributed.auto_parallel;
+
+// ProcessMesh is used to organize processes and like n-dimension array.
+message ProcessMeshProto {
+  // The size of each dimension.
+  repeated int64 shape = 1;
+
+  // These process ids are stored by a row-major way.
+  // There are no duplicate process ids within one process mesh.
+  repeated int64 process_ids = 2;
+
+  // The name of each dimension.
+  repeated string dim_names = 3;
+
+}
+
+message MeshDimsProto {
+  repeated int64 mesh_dims = 1;
+}
+
+// This distributed attribute describes how to distribute the corresponding tensor,
+// and store any other information needed by auto parallel.
+message TensorDistAttrProto {
+  // The process mesh where a tensor is distributed.
+  optional ProcessMeshProto process_mesh = 1;
+
+  // The length of dims_mapping is same as the length of the tensor shape.
+  // The i-th dimension of the tensor will be sharded by the dims_mapping[i]-th dimension
+  // of the above process mesh. If dims_mapping[i] is -1, the i-th dimension of the tensor
+  // will not be sharded. For example, given a tensor shape [2, 6, 12], a process mesh
+  // shape [2, 3] and a dims_mapping [-1, 1, 0], each sharded tensor will have a shape [2, 2, 6].
+  repeated MeshDimsProto dims_mapping = 2;
+
+  // The batch dimension of the corresponding tensor.
+  optional int64 batch_dim = 3;
+
+  // If the dynamic_dims[i] is True, the i-th dimension of the corresponding tensor
+  // is dynamic changed. Otherwise, the i-th dimension of the tensor is static determined.
+  repeated bool dynamic_dims = 4;
+
+  // This field is used to distinguish vars which are in same process_mesh and in different vpp chunk
+  optional int64 chunk_id = 5;
+}
+
+// This distributed attribute describes how to distribute the corresponding operator,
+// and store any other information needed by auto parallel.
+message OperatorDistAttrProto {
+  message TensorDistAttrMappingEntryProto {
+    optional string name = 1;
+    optional TensorDistAttrProto tensor_dist_attr = 2;
+  }
+  // The key of this map is the input tensor name and the value is the distributed attribute
+  // of the input tensor required by this corresponding operator.
+  // The distributed attribute of the actual tensor may be not the same as that within
+  // the distributed attribute of the operator.
+  repeated TensorDistAttrMappingEntryProto input_dist_attrs = 1;
+
+  // The key of this map is the output tensor name and the value is the distributed attribute
+  // of the output tensor required by this corresponding operator.
+  // The distributed attribute of the actual tensor may be not the same as that within
+  // the distributed attribute of the operator.
+  repeated TensorDistAttrMappingEntryProto output_dist_attrs = 2;
+
+  // The process mesh where a op is distributed.
+  optional ProcessMeshProto process_mesh = 3;
+
+  // A operator ideally has a distributed operator which may have multiple distributed implementations.
+  // This filed is usually same as the operator type. However, some operators such as the element-wise operators
+  // may shared the same distributed operator, the field is use for this scenario.
+  optional string impl_type = 4;
+
+  // This field tells which distributed implementations of this corresponding operator
+  // will be selected for the actual computation.
+  optional int64 impl_idx = 5;
+
+  // This field is used to distinguish ops which are in same process_mesh and in different vpp chunk
+  optional int64 chunk_id = 6;
+}
+
+// This proto describes the capability of one device such as the computation and memory.
+message DeviceCapabilityProto {
+  optional double single_precision_flops = 1;
+
+  optional double double_precision_flops = 2;
+
+  optional double memory_size_in_bytes = 3;
+
+  optional double clock_rate_in_ghz = 4;
+}
+
+// This proto represents a device.
+message DeviceProto {
+  // The global id of this device within the cluster.
+  optional int64 global_id = 1;
+
+  // The local id of this device within the machine.
+  optional int64 local_id = 2;
+
+  // The id of the machine own this device.
+  optional int64 machine_id = 3;
+
+  // The id of the machine has this device.
+  optional string type = 4;
+
+  // The capability of this device.
+  optional DeviceCapabilityProto capability = 5;
+}
+
+// This proto describes the capability of the link between two devices.
+message LinkCapabilityProto {
+  optional int64 bandwidth = 1; // Bytes/s
+  optional int64 latency = 2;
+}
+
+message LinkProto {
+  // The global id of the source device.
+  optional int64 source_id = 1;
+
+  // The global id of the source device.
+  optional int64 target_id = 2;
+
+  // Represent the link type.
+  optional string type = 3;
+
+  // The capability of this link.
+  optional LinkCapabilityProto capability = 4;
+}
+
+// DeviceMesh is used to organize devices and like n-dimension array.
+message DeviceMeshProto {
+  // The global id of this mesh.
+  optional string name = 1;
+
+  // The size of each dimension.
+  repeated int64 shape = 2;
+
+  // These device ids are stored by a row-major way.
+  // There are no duplicate device ids within one device mesh.
+  repeated int64 device_ids = 3;
+
+  // The name of each dimension.
+  repeated string dim_names = 4;
+
+  // The devices of this mesh.
+  repeated DeviceProto devices = 5;
+
+  // The links are between devices.
+  repeated LinkProto links = 6;
+}
+
+// Record the mapping between the logical processes and the physical devices.
+message DistributedMapperProto {
+  // The device meshes used by this distributed computation,
+  // which may be shared by different multiple device meshes.
+  repeated DeviceMeshProto device_meshes = 1;
+
+  message MapperEntryProto {
+    optional int64 process_id = 1;
+    optional string device_mesh_name = 2;
+    repeated int64 device_ids = 3;
+  }
+
+  // The mapping from process ids to device ids.
+  // It is also possible for one process to use multiple devices.
+  // It is possible for one device shared by multiple processes.
+  repeated MapperEntryProto process_id_to_device_ids = 2;
+}

paddle/phi/core/distributed/auto_parallel/device_mesh.cc

@@ -0,0 +1,384 @@
+/* Copyright (c) 2022 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. */
+
+#include <algorithm>
+#include <iterator>
+
+#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+namespace phi::distributed::auto_parallel {
+
+std::string DeviceCapability::to_string() const {
+  std::string str;
+  str += "{sflops: " + to_string_with_precision(single_precision_flops) + ", ";
+  str += "dflops: " + to_string_with_precision(double_precision_flops) + ", ";
+  str += "memory: " + to_string_with_precision(memory_size_in_bytes) + ", ";
+  str += "rate: " + to_string_with_precision(clock_rate_in_ghz) + "}";
+  return str;
+}
+
+DeviceCapability DeviceCapability::from_proto(
+    const DeviceCapabilityProto &proto) {
+  DeviceCapability capability;
+  capability.single_precision_flops = proto.single_precision_flops();
+  capability.double_precision_flops = proto.double_precision_flops();
+  capability.memory_size_in_bytes = proto.memory_size_in_bytes();
+  capability.clock_rate_in_ghz = proto.clock_rate_in_ghz();
+  return capability;
+}
+
+void DeviceCapability::to_proto(DeviceCapabilityProto *proto) const {
+  proto->set_single_precision_flops(single_precision_flops);
+  proto->set_double_precision_flops(double_precision_flops);
+  proto->set_memory_size_in_bytes(memory_size_in_bytes);
+  proto->set_clock_rate_in_ghz(clock_rate_in_ghz);
+}
+
+std::string Device::to_string() const {
+  std::string str = "{global_id: " + std::to_string(global_id_) + ", ";
+  str += "local_id: " + std::to_string(local_id_) + ", ";
+  str += "machine_id: " + std::to_string(machine_id_) + ", ";
+  str += "type: " + type_ + ", ";
+  str += "capability: " + capability_.to_string() + "}";
+  return str;
+}
+
+Device Device::from_proto(const DeviceProto &proto) {
+  Device device;
+  device.global_id_ = proto.global_id();
+  device.local_id_ = proto.local_id();
+  device.machine_id_ = proto.machine_id();
+  device.type_ = proto.type();
+  device.capability_ = DeviceCapability::from_proto(proto.capability());
+  return device;
+}
+
+void Device::to_proto(DeviceProto *proto) const {
+  proto->set_global_id(global_id_);
+  proto->set_local_id(local_id_);
+  proto->set_machine_id(machine_id_);
+  proto->set_type(type_);
+  proto->mutable_capability()->CopyFrom(distributed::to_proto(capability_));
+}
+
+bool operator==(const Device &lhs, const Device &rhs) {
+  if (lhs.global_id() != rhs.global_id()) {
+    return false;
+  }
+  if (lhs.local_id() != rhs.local_id()) {
+    return false;
+  }
+  if (lhs.machine_id() != rhs.machine_id()) {
+    return false;
+  }
+  if (lhs.type() != rhs.type()) {
+    return false;
+  }
+  return true;
+}
+
+std::string LinkCapability::to_string() const {
+  std::string str;
+  str += "{bandwidth: " + to_string_with_precision(bandwidth) + ",";
+  str += "latency: " + to_string_with_precision(latency) + "}";
+  return str;
+}
+
+LinkCapability LinkCapability::from_proto(const LinkCapabilityProto &proto) {
+  LinkCapability capability;
+  capability.bandwidth = proto.bandwidth();
+  capability.latency = proto.latency();
+  return capability;
+}
+
+void LinkCapability::to_proto(LinkCapabilityProto *proto) const {
+  proto->set_bandwidth(bandwidth);
+  proto->set_latency(latency);
+}
+
+std::string Link::to_string() const {
+  std::string str = "{source_id:" + std::to_string(source_id_) + ",";
+  str += "target_id:" + std::to_string(target_id_) + ",";
+  str += "type:" + type_ + ",";
+  str += "capability:" + capability_.to_string() + "}";
+  return str;
+}
+
+Link Link::from_proto(const LinkProto &proto) {
+  Link link;
+  link.source_id_ = proto.source_id();
+  link.target_id_ = proto.target_id();
+  link.type_ = proto.type();
+  link.capability_ = LinkCapability::from_proto(proto.capability());
+  return link;
+}
+
+void Link::to_proto(LinkProto *proto) const {
+  proto->set_source_id(source_id_);
+  proto->set_target_id(target_id_);
+  proto->set_type(type_);
+  proto->mutable_capability()->CopyFrom(distributed::to_proto(capability_));
+}
+
+bool operator==(const Link &lhs, const Link &rhs) {
+  if (lhs.source_id() != rhs.source_id()) {
+    return false;
+  }
+  if (lhs.target_id() != rhs.target_id()) {
+    return false;
+  }
+  if (lhs.type() != rhs.type()) {
+    return false;
+  }
+  return true;
+}
+
+bool Machine::contains(int64_t device_id) const {
+  if (devices_.count(device_id) == 1) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+void Machine::add_device(const Device &device) {
+  if (id() == -1) {
+    set_id(device.machine_id());
+  } else {
+    PADDLE_ENFORCE_EQ(device.machine_id(),
+                      id(),
+                      errors::InvalidArgument(
+                          "The machine id [%d] of the device should be equal "
+                          "to this machine id [%d].",
+                          device.machine_id(),
+                          id_));
+  }
+  devices_[device.global_id()] = &device;
+}
+
+void Machine::add_link(const Link &link) {
+  PADDLE_ENFORCE_EQ(contains(link.source_id()),
+                    true,
+                    errors::InvalidArgument(
+                        "The source device id of the added link [%s] "
+                        "cannot be found in the device_ids. Please add the "
+                        "source device before adding this link",
+                        std::to_string(link.source_id())));
+  links_[link.source_id()][link.target_id()] = &link;
+}
+
+std::string Machine::to_string() const {
+  std::string str = "{devices: [";
+  for (const auto &device : devices_) {
+    str += device.second->to_string() + ", ";
+  }
+  str.replace(str.size() - 2, 2, "], ");
+
+  str += "links: [";
+  for (const auto &item : links_) {
+    str += "{";
+    str += "source_id: " + std::to_string(item.first) + ", neighbors: [";
+    for (const auto &link : item.second) {
+      str += link.second->to_string() + ", ";
+    }
+    str.replace(str.size() - 2, 2, "]}, ");
+  }
+  str.replace(str.size() - 4, 4, "]}");
+  return str;
+}
+
+DeviceMesh::DeviceMesh(const std::string &name,
+                       const std::vector<int64_t> &shape,
+                       const std::vector<int64_t> &device_ids,
+                       const std::vector<std::string> &dim_names) {
+  name_ = name;
+  shape_ = shape;
+  int64_t size = this->size();
+
+  PADDLE_ENFORCE_EQ(
+      size,
+      device_ids.size(),
+      errors::InvalidArgument("The size %d of this device mesh must be "
+                              "equal to the size %d of its device ids.",
+                              size,
+                              device_ids.size()));
+  PADDLE_ENFORCE_EQ(
+      has_duplicates(device_ids),
+      false,
+      errors::InvalidArgument("The device ids [%s] must be unique.",
+                              str_join(device_ids)));
+  device_ids_ = device_ids;
+
+  PADDLE_ENFORCE_EQ(
+      shape_.size(),
+      dim_names.size(),
+      errors::InvalidArgument(
+          "The size %d of mesh shape must be equal to the size %d "
+          "of the dimension names.",
+          shape_.size(),
+          dim_names.size()));
+  PADDLE_ENFORCE_EQ(has_duplicates(dim_names),
+                    false,
+                    errors::InvalidArgument(
+                        "The names [%s] of each dimension must be unique.",
+                        str_join(dim_names)));
+  dim_names_ = dim_names;
+}
+
+int64_t DeviceMesh::size() const {
+  if (shape_.empty()) return 0;
+  int64_t size = 1;
+  for (const int64_t dim_size : shape_) size *= dim_size;
+  return size;
+}
+
+bool DeviceMesh::contains(int64_t device_id) const {
+  auto result =
+      std::find(std::begin(device_ids_), std::end(device_ids_), device_id);
+  if (result != std::end(device_ids_)) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+void DeviceMesh::add_device(const Device &device) {
+  PADDLE_ENFORCE_EQ(
+      contains(device.global_id()),
+      true,
+      errors::InvalidArgument(
+          "The added device id [%s] cannot be found in the device_ids.",
+          std::to_string(device.global_id())));
+  // Operator [] will create a new object if it cannot find one.
+  // So we add the default constructor for Device and Machine
+  // to make sure the new object can be created.
+  devices_[device.global_id()] = device;
+  machines_[device.machine_id()].add_device(devices_[device.global_id()]);
+}
+
+void DeviceMesh::add_link(const Link &link) {
+  PADDLE_ENFORCE_EQ(
+      contains(link.source_id()),
+      true,
+      errors::InvalidArgument("The source id of the added link [%s] "
+                              "cannot be found in the device_ids.",
+                              std::to_string(link.source_id())));
+  PADDLE_ENFORCE_EQ(
+      contains(link.target_id()),
+      true,
+      errors::InvalidArgument("The source id of the added link [%s] "
+                              "cannot be found in the device_ids.",
+                              std::to_string(link.target_id())));
+  // Operator [] will create a new object if it cannot find one.
+  // So we add the default constructor for Device and Machine
+  // to make sure the new object can be created.
+  links_[link.source_id()][link.target_id()] = link;
+  const Device &source_device = devices_[link.source_id()];
+  machines_[source_device.machine_id()].add_link(
+      links_[link.source_id()][link.target_id()]);
+}
+
+std::string DeviceMesh::to_string() const {
+  std::string mesh_str = "{name: " + name_ + ", ";
+  mesh_str += "shape: [" + str_join(shape_) + "], ";
+  mesh_str += "device_ids: [" + str_join(device_ids_) + "], ";
+  mesh_str += "dim_names: [" + str_join(dim_names_) + "], ";
+  mesh_str += "\ndevices: [\n";
+  for (const auto &device : devices_) {
+    mesh_str += "  " + device.second.to_string() + ",\n";
+  }
+  mesh_str.replace(mesh_str.size() - 2, 2, "],");
+
+  mesh_str += "\nlinks: [\n";
+  for (const auto &item : links_) {
+    mesh_str += "  {";
+    mesh_str += "source_id: " + std::to_string(item.first) + ", neighbors: [";
+    for (const auto &link : item.second) {
+      mesh_str += link.second.to_string() + ", ";
+    }
+    mesh_str.replace(mesh_str.size() - 2, 2, "]},\n");
+  }
+  mesh_str.replace(mesh_str.size() - 4, 4, "]}");
+  return mesh_str;
+}
+
+DeviceMesh DeviceMesh::from_proto(const DeviceMeshProto &proto) {
+  DeviceMesh mesh;
+
+  mesh.name_ = proto.name();
+
+  mesh.shape_.resize(proto.shape_size());
+  for (int i = 0; i < proto.shape_size(); ++i) {
+    mesh.shape_[i] = proto.shape(i);
+  }
+
+  mesh.device_ids_.resize(proto.device_ids_size());
+  for (int i = 0; i < proto.device_ids_size(); ++i) {
+    mesh.device_ids_[i] = proto.device_ids(i);
+  }
+
+  mesh.dim_names_.resize(proto.dim_names_size());
+  for (int i = 0; i < proto.dim_names_size(); ++i) {
+    mesh.dim_names_[i] = proto.dim_names(i);
+  }
+
+  for (int i = 0; i < proto.devices_size(); ++i) {
+    mesh.add_device(Device::from_proto(proto.devices(i)));
+  }
+
+  for (int i = 0; i < proto.links_size(); ++i) {
+    mesh.add_link(Link::from_proto(proto.links(i)));
+  }
+
+  return mesh;
+}
+
+void DeviceMesh::to_proto(DeviceMeshProto *proto) const {
+  proto->set_name(name_);
+
+  for (const auto &i : shape_) {
+    proto->add_shape(i);
+  }
+
+  for (const auto &i : device_ids_) {
+    proto->add_device_ids(i);
+  }
+
+  for (const auto &i : dim_names_) {
+    proto->add_dim_names(i);
+  }
+
+  for (const auto &device : devices_) {
+    proto->mutable_devices()->Add()->CopyFrom(
+        distributed::to_proto(device.second));
+  }
+
+  for (const auto &neighbors : links_) {
+    for (const auto &link : neighbors.second) {
+      proto->mutable_links()->Add()->CopyFrom(
+          distributed::to_proto(link.second));
+    }
+  }
+}
+
+bool operator==(const DeviceMesh &lhs, const DeviceMesh &rhs) {
+  // Use the unique name to do the fast comparison
+  if (lhs.name() != rhs.name()) {
+    return false;
+  }
+  return true;
+}
+
+}  // namespace phi::distributed::auto_parallel

paddle/phi/core/distributed/auto_parallel/device_mesh.h

@@ -0,0 +1,307 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <iostream>
+#include <iterator>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace distributed {
+namespace auto_parallel {
+
+class DeviceCapabilityProto;
+class DeviceProto;
+class LinkCapabilityProto;
+class LinkProto;
+class DeviceMeshProto;
+
+struct PADDLE_API DeviceCapability {
+  double single_precision_flops = 0.0;
+  double double_precision_flops = 0.0;
+  double memory_size_in_bytes = 0.0;
+  double clock_rate_in_ghz = 0.0;
+
+  // DeviceCapability from_string(const std::string& str);
+  std::string to_string() const;
+
+  static DeviceCapability from_proto(const DeviceCapabilityProto& proto);
+  void to_proto(DeviceCapabilityProto* proto) const;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const DeviceCapability& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+class PADDLE_API Device {
+ public:
+  Device() = default;
+  Device(int64_t global_id,
+         int64_t local_id,
+         int64_t machine_id,
+         const std::string& type)
+      : global_id_(global_id),
+        local_id_(local_id),
+        machine_id_(machine_id),
+        type_(type) {}
+
+  int64_t global_id() const { return global_id_; }
+  int64_t local_id() const { return local_id_; }
+  int64_t machine_id() const { return machine_id_; }
+  const std::string& type() const { return type_; }
+
+  const DeviceCapability& capability() const { return capability_; }
+  void set_capability(const DeviceCapability& capability) {
+    capability_ = capability;
+  }
+
+  // Device from_string(const std::string& mesh_str);
+  std::string to_string() const;
+
+  static Device from_proto(const DeviceProto& proto);
+  void to_proto(DeviceProto* proto) const;
+
+ private:
+  int64_t global_id_;
+  int64_t local_id_;
+  int64_t machine_id_;
+  std::string type_;
+  DeviceCapability capability_;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const Device& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+PADDLE_API bool operator==(const Device& lhs, const Device& rhs);
+
+inline bool operator!=(const Device& lhs, const Device& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+struct PADDLE_API LinkCapability {
+  int64_t bandwidth = 0.0;  // Bytes/s
+  int64_t latency = 0.0;
+
+  // LinkCapability from_string(const std::string& str);
+  std::string to_string() const;
+
+  static LinkCapability from_proto(const LinkCapabilityProto& proto);
+  void to_proto(LinkCapabilityProto* proto) const;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const LinkCapability& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+class PADDLE_API Link {
+ public:
+  Link() = default;
+
+  Link(int64_t source_id, int64_t target_id, const std::string& type)
+      : source_id_(source_id), target_id_(target_id), type_(type) {}
+
+  int64_t source_id() const { return source_id_; }
+  int64_t target_id() const { return target_id_; }
+  const std::string& type() const { return type_; }
+
+  const LinkCapability& capability() const { return capability_; }
+  void set_capability(const LinkCapability& capability) {
+    capability_ = capability;
+  }
+
+  // Link from_string(const std::string& str);
+  std::string to_string() const;
+
+  static Link from_proto(const LinkProto& proto);
+  void to_proto(LinkProto* proto) const;
+
+ private:
+  int64_t source_id_;
+  int64_t target_id_;
+  std::string type_;
+  LinkCapability capability_;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const Link& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+PADDLE_API bool operator==(const Link& lhs, const Link& rhs);
+
+inline bool operator!=(const Link& lhs, const Link& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+class PADDLE_API Machine {
+ public:
+  Machine() = default;
+
+  explicit Machine(int64_t id) : id_(id) {}
+
+  int64_t id() const { return id_; }
+
+  void set_id(int64_t id) { id_ = id; }
+
+  const std::unordered_map<int64_t, const Device*>& devices() const {
+    return devices_;
+  }
+
+  const std::unordered_map<int64_t, std::unordered_map<int64_t, const Link*>>&
+  links() const {
+    return links_;
+  }
+
+  const Device& device(int64_t global_id) const {
+    return *devices_.at(global_id);
+  }
+
+  const Link& link(int64_t source_id, int64_t target_id) const {
+    return *links_.at(source_id).at(target_id);
+  }
+
+  bool contains(int64_t device_id) const;
+
+  void add_device(const Device& device);
+
+  void add_link(const Link& link);
+
+  // Machine from_string(const std::string& str);
+  std::string to_string() const;
+
+ private:
+  int64_t id_ = -1;
+  std::unordered_map<int64_t, const Device*> devices_;
+  std::unordered_map<int64_t, std::unordered_map<int64_t, const Link*>> links_;
+};
+
+class PADDLE_API DeviceMesh {
+ public:
+  DeviceMesh() = default;
+
+  DeviceMesh(const std::string& name,
+             const std::vector<int64_t>& shape,
+             const std::vector<int64_t>& device_ids,
+             const std::vector<std::string>& dim_names);
+
+  const std::string& name() const { return name_; }
+
+  void set_name(const std::string& name) { name_ = name; }
+
+  const std::vector<int64_t>& shape() const { return shape_; }
+
+  const std::vector<int64_t>& device_ids() const { return device_ids_; }
+
+  const std::vector<std::string>& dim_names() const { return dim_names_; }
+
+  std::string device_type() const {
+    if (empty()) return "UNKNOWN";
+    if (devices_.empty())
+      return "UNKNOWN";
+    else
+      return std::begin(devices_)->second.type();
+  }
+
+  const std::unordered_map<int64_t, Device>& devices() const {
+    return devices_;
+  }
+
+  const std::unordered_map<int64_t, std::unordered_map<int64_t, Link>>& links()
+      const {
+    return links_;
+  }
+
+  const std::unordered_map<int64_t, Machine>& machines() const {
+    return machines_;
+  }
+
+  const Device& device(int64_t global_id) const {
+    return devices_.at(global_id);
+  }
+
+  const Link& link(int64_t source_id, int64_t target_id) const {
+    return links_.at(source_id).at(target_id);
+  }
+
+  const Machine& machine(int64_t machine_id) const {
+    return machines_.at(machine_id);
+  }
+
+  int64_t size() const;
+  int64_t ndim() const { return shape_.size(); }
+
+  int64_t dim_size(int64_t dim) const {
+    int64_t cdim = canonical_dim(dim, shape_.size());
+    return shape_[cdim];
+  }
+
+  int64_t dim_size(const std::string& dim_name) const {
+    for (std::size_t i = 0; i < dim_names_.size(); ++i) {
+      if (dim_names_[i] == dim_name) {
+        return shape_[i];
+      }
+    }
+    PADDLE_THROW(errors::InvalidArgument(
+        "Cannot find the dimension of %s in this device mesh.", dim_name));
+  }
+
+  bool empty() const { return (shape_.empty() || device_ids_.empty()); }
+  bool contains(int64_t device_id) const;
+
+  void add_device(const Device& device);
+  void add_link(const Link& link);
+
+  // DeviceMesh from_string(const std::string& mesh_str);
+  std::string to_string() const;
+
+  static DeviceMesh from_proto(const DeviceMeshProto& proto);
+  void to_proto(DeviceMeshProto* proto) const;
+
+ private:
+  std::string name_;
+  std::vector<int64_t> shape_;
+  std::vector<int64_t> device_ids_;
+  std::vector<std::string> dim_names_;
+  std::unordered_map<int64_t, Device> devices_;
+  std::unordered_map<int64_t, std::unordered_map<int64_t, Link>> links_;
+  std::unordered_map<int64_t, Machine> machines_;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const DeviceMesh& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+PADDLE_API bool operator==(const DeviceMesh& lhs, const DeviceMesh& rhs);
+
+inline bool operator!=(const DeviceMesh& lhs, const DeviceMesh& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+}  // namespace auto_parallel
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/dist_attr.cc

@@ -0,0 +1,605 @@
+/* Copyright (c) 2022 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. */
+#include <algorithm>
+#include <iostream>
+#include <iterator>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
+
+namespace phi::distributed {
+using phi::distributed::auto_parallel::str_join;
+using phi::distributed::auto_parallel::TensorDistAttrProto;
+
+// partial is not allow annotated by user by now.
+std::vector<std::string> TensorDistAttr::fields_{
+    "process_mesh", "dims_mapping", "batch_dim", "chunk_id", "dynamic_dims"};
+
+TensorDistAttr::TensorDistAttr(const std::vector<int64_t>& tensor_shape) {
+  set_default_dims_mapping(tensor_shape);
+  set_default_dynamic_dims(tensor_shape);
+}
+
+TensorDistAttr::TensorDistAttr(const TensorDistAttr& dist_attr) {
+  copy_from(dist_attr);
+}
+
+TensorDistAttr& TensorDistAttr::operator=(const TensorDistAttr& dist_attr) {
+  if (this == &dist_attr) return *this;
+  TensorDistAttr tmp(dist_attr);
+  std::swap(this->process_mesh_, tmp.process_mesh_);
+  std::swap(this->dims_mapping_, tmp.dims_mapping_);
+  std::swap(this->batch_dim_, tmp.batch_dim_);
+  std::swap(this->chunk_id_, tmp.chunk_id_);
+  std::swap(this->dynamic_dims_, tmp.dynamic_dims_);
+  std::swap(this->annotated_, tmp.annotated_);
+  std::swap(this->partial_status_, tmp.partial_status_);
+  std::swap(this->skip_check_mesh_, tmp.skip_check_mesh_);
+  std::swap(this->split_factor_map_, tmp.split_factor_map_);
+  return *this;
+}
+
+void TensorDistAttr::copy_from(const TensorDistAttr& dist_attr) {
+  set_process_mesh(dist_attr.process_mesh());
+  set_dims_mapping(dist_attr.multi_dims_mapping());
+  split_factor_map_ = dist_attr.split_factor();
+  set_batch_dim(dist_attr.batch_dim());
+  set_chunk_id(dist_attr.chunk_id());
+  set_dynamic_dims(dist_attr.dynamic_dims());
+  set_annotated(dist_attr.annotated());
+  set_partial_status(dist_attr.partial_status());
+  skip_check_mesh_ = dist_attr.skip_check_mesh();
+}
+
+void TensorDistAttr::set_process_mesh(const ProcessMesh& process_mesh) {
+  process_mesh_ = process_mesh;
+}
+
+void TensorDistAttr::set_dims_mapping(
+    const std::vector<int64_t>& dims_mapping) {
+  dims_mapping_proxy = dims_mapping;
+  // dynamic_dims_ and dims_mapping may be not consistent
+  if (dynamic_dims_.empty() || dims_mapping.empty()) {
+    set_default_dynamic_dims(dims_mapping);
+  }
+}
+
+void TensorDistAttr::set_dims_mapping(
+    const std::vector<std::vector<int64_t>>& dims_mapping) {
+  dims_mapping_proxy = dims_mapping;
+  // dynamic_dims_ and dims_mapping may be not consistent
+  if (dynamic_dims_.empty() || dims_mapping.empty()) {
+    set_default_dynamic_dims(std::vector<int64_t>(dims_mapping.size()));
+  }
+}
+
+void TensorDistAttr::set_batch_dim(int64_t batch_dim) {
+  batch_dim_ = batch_dim;
+}
+
+void TensorDistAttr::set_chunk_id(const int64_t& chunk_id) {
+  chunk_id_ = chunk_id;
+}
+
+void TensorDistAttr::set_dynamic_dims(const std::vector<bool>& dynamic_dims) {
+  dynamic_dims_ = dynamic_dims;
+}
+
+void TensorDistAttr::set_annotated(
+    const std::map<std::string, bool>& annotated) {
+  annotated_ = annotated;
+}
+
+int64_t TensorDistAttr::get_split_factor(int64_t mesh_dim) const {
+  PADDLE_ENFORCE_LT(mesh_dim,
+                    process_mesh_.ndim(),
+                    ::common::errors::InvalidArgument(
+                        "Expected mesh dim is less than process mesh dim size, "
+                        "but got mesh dim  %d, Process mesh ndim is %d",
+                        mesh_dim,
+                        process_mesh_.ndim()));
+  return split_factor_map_.get_split_factor(mesh_dim);
+}
+
+void TensorDistAttr::set_split_factor(int64_t mesh_dim, int64_t split_factor) {
+  split_factor_map_.set_split_factor(mesh_dim, split_factor);
+}
+
+void TensorDistAttr::clear_split_factor(int64_t mesh_dim) {
+  split_factor_map_.clear_split_factor(mesh_dim);
+}
+
+const std::set<int64_t> TensorDistAttr::partial_dims() const {
+  std::set<int64_t> keys;
+  for (auto& kv : partial_status_) {
+    keys.emplace(kv.first);
+  }
+  return keys;
+}
+
+void TensorDistAttr::set_partial_status(
+    const paddle::flat_hash_map<int64_t, ReduceType>& partial_status) {
+  partial_status_ = partial_status;
+}
+
+void TensorDistAttr::set_partial_status(const std::vector<int64_t>& dims,
+                                        const ReduceType& type) {
+  for (const auto& dim : dims) {
+    if (is_partial(dim)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Trying to Set dim %d as Partial which is already a Partial dim.",
+          dim));
+    }
+    if (is_shard(dim)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Trying to Set dim %d as Partial which is a Sharding dim.", dim));
+    }
+    partial_status_.emplace(dim, type);
+  }
+}
+
+void TensorDistAttr::clean_partial_status() { partial_status_.clear(); }
+
+void TensorDistAttr::clean_partial_dims(const std::vector<int64_t>& dims) {
+  for (const auto& dim : dims) {
+    if (partial_status_.count(dim) == 0) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Trying to clean Partial on dim %d but it is not Partial.", dim));
+    } else {
+      partial_status_.erase(dim);
+    }
+  }
+}
+
+void TensorDistAttr::set_default_dims_mapping(
+    const std::vector<int64_t>& tensor_shape) {
+  if (!tensor_shape.empty()) {
+    dims_mapping_proxy = std::vector<std::vector<int64_t>>(tensor_shape.size());
+  }
+}
+
+void TensorDistAttr::set_default_dynamic_dims(
+    const std::vector<int64_t>& tensor_shape) {
+  dynamic_dims_ = std::vector<bool>(tensor_shape.size(), false);
+}
+
+void TensorDistAttr::set_default_dynamic_dims(int64_t tensor_shape_size) {
+  dynamic_dims_ = std::vector<bool>(tensor_shape_size, false);
+}
+
+void TensorDistAttr::mark_annotated(const std::string& name) {
+  auto result = std::find(std::begin(fields_), std::end(fields_), name);
+  if (result != std::end(fields_)) {
+    annotated_[name] = true;
+  }
+}
+
+bool TensorDistAttr::verify_process_mesh(
+    const ProcessMesh& process_mesh) const {
+  VLOG(4) << "[TensorDistAttr verify_process_mesh] "
+          << process_mesh.to_string();
+  if (!process_mesh_.empty()) {
+    for (const auto& dims_mapping : dims_mapping_) {
+      for (int64_t mesh_dim : dims_mapping)
+        if (mesh_dim >= process_mesh_.ndim()) {
+          return false;
+        }
+    }
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_dims_mapping(
+    const std::vector<std::vector<int64_t>>& dims_mapping,
+    const std::vector<int64_t>& tensor_shape) const {
+  VLOG(4) << "[TensorDistAttr verify_dims_mapping] " << str_join(dims_mapping);
+  if (dims_mapping.size() != tensor_shape.size()) {
+    return false;
+  }
+  std::unordered_map<int64_t, int64_t> map;
+  if (!process_mesh_.empty()) {
+    for (const auto& mesh_dims : dims_mapping) {
+      for (int64_t mesh_dim : mesh_dims) {
+        if (mesh_dim >= process_mesh_.ndim()) {
+          return false;
+        }
+        ++map[mesh_dim];
+        if (map[mesh_dim] > 1) {
+          return false;
+        }
+      }
+    }
+  } else {
+    for (const auto& mesh_dims : dims_mapping) {
+      for (int64_t dim : mesh_dims) {
+        ++map[dim];
+        if (map[dim] > 1) {
+          return false;
+        }
+      }
+    }
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_batch_dim(
+    int64_t dim, const std::vector<int64_t>& tensor_shape) const {
+  VLOG(4) << "[TensorDistAttr verify_batch_dim] " << dim;
+  int64_t ndim = static_cast<int64_t>(tensor_shape.size());
+  if (ndim > 0) {
+    if (dim < 0) {
+      dim = dim + ndim;
+    }
+    if (dim < 0 || dim >= ndim) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_dynamic_dims(
+    const std::vector<bool>& dynamic_dims,
+    const std::vector<int64_t>& tensor_shape) const {
+  VLOG(4) << "[TensorDistAttr verify_dynamic_dims] " << str_join(dynamic_dims);
+  if (!dynamic_dims.empty() && dynamic_dims.size() != tensor_shape.size()) {
+    return false;
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_annotated(
+    const std::map<std::string, bool>& annotated) const {
+  VLOG(4) << "[TensorDistAttr verify_annotated] " << str_join(annotated);
+  for (const auto& item : annotated) {
+    auto result = std::find(std::begin(fields_), std::end(fields_), item.first);
+    if (result == std::end(fields_)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_partial_status() const {
+  VLOG(4) << "[TensorDistAttr verify_partial_status] "
+          << partial_status_string();
+  for (auto& itr : partial_status_) {
+    if (itr.first < 0 || itr.first >= process_mesh_.ndim()) {
+      return false;
+    }
+    if (itr.second < ReduceType::kRedSum || itr.second > ReduceType::kRedAll) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify(const std::vector<int64_t>& tensor_shape) const {
+  if (!verify_process_mesh(process_mesh_)) {
+    return false;
+  }
+  if (!verify_dims_mapping(dims_mapping_, tensor_shape)) {
+    return false;
+  }
+  if (!verify_batch_dim(batch_dim_, tensor_shape)) {
+    return false;
+  }
+  if (!verify_dynamic_dims(dynamic_dims_, tensor_shape)) {
+    return false;
+  }
+  if (!verify_annotated(annotated_)) {
+    return false;
+  }
+  if (!verify_partial_status()) {
+    return false;
+  }
+  return true;
+}
+
+bool TensorDistAttr::verify_dynamic(
+    const std::vector<int64_t>& tensor_shape) const {
+  if (!verify_process_mesh(process_mesh_)) {
+    return false;
+  }
+  if (!verify_dims_mapping(dims_mapping_, tensor_shape)) {
+    return false;
+  }
+  if (!verify_partial_status()) {
+    return false;
+  }
+  return true;
+}
+
+std::string TensorDistAttr::to_string() const {
+  std::string dist_str;
+  dist_str += "{process_mesh: " + process_mesh_.to_string() + ", ";
+  dist_str += "dim_mappings: [" + str_join(dims_mapping_) + "], ";
+  dist_str += "batch_dim: " + std::to_string(batch_dim_) + ", ";
+  dist_str += "chunk_id: " + std::to_string(chunk_id_) + ", ";
+  dist_str += "skip_check_mesh: " + std::to_string(skip_check_mesh_) + ", ";
+  dist_str += "dynamic_dims: [" + str_join(dynamic_dims_) + "], ";
+  dist_str += "annotated: [" + str_join(annotated_) + "], ";
+  dist_str += "partial: " + partial_status_string() + ", ";
+  dist_str += "split_factor: " + split_factor_map_.to_string() + ".}";
+  return dist_str;
+}
+
+void TensorDistAttr::from_proto(const TensorDistAttrProto& proto) {
+  process_mesh_ = ProcessMesh::from_proto(proto.process_mesh());
+  dims_mapping_.resize(proto.dims_mapping_size());
+
+  for (int i = 0; i < proto.dims_mapping_size(); ++i) {
+    const auto& mesh_dims_proto = proto.dims_mapping(i);
+    std::vector<int64_t> mesh_dims_vec;
+    for (int j = 0; j < mesh_dims_proto.mesh_dims_size(); ++j) {
+      mesh_dims_vec.push_back(mesh_dims_proto.mesh_dims(j));
+    }
+    dims_mapping_[i] = mesh_dims_vec;
+  }
+
+  batch_dim_ = proto.batch_dim();
+  chunk_id_ = proto.chunk_id();
+  dynamic_dims_.resize(proto.dynamic_dims_size());
+  for (int i = 0; i < proto.dynamic_dims_size(); ++i) {
+    dynamic_dims_[i] = proto.dynamic_dims(i);
+  }
+}
+
+void TensorDistAttr::to_proto(TensorDistAttrProto* proto) const {
+  proto->mutable_process_mesh()->CopyFrom(distributed::to_proto(process_mesh_));
+
+  for (size_t i = 0; i < dims_mapping_.size(); ++i) {
+    proto->add_dims_mapping();
+    for (size_t j = 0; j < dims_mapping_.at(i).size(); ++j) {
+      proto->mutable_dims_mapping(i)->add_mesh_dims(dims_mapping_[i][j]);
+    }
+  }
+
+  proto->set_batch_dim(batch_dim_);
+  proto->set_chunk_id(chunk_id_);
+  for (const auto& i : dynamic_dims_) {
+    proto->add_dynamic_dims(i);
+  }
+}
+
+std::string TensorDistAttr::serialize_to_string() {
+  std::string data;
+  auto proto = distributed::to_proto(*this);
+  proto.SerializeToString(&data);
+  PADDLE_ENFORCE_EQ(distributed::to_proto(*this).SerializeToString(&data),
+                    true,
+                    errors::InvalidArgument(
+                        "Failed to serialize tensor dist attr to string."));
+  return data;
+}
+
+void TensorDistAttr::parse_from_string(const std::string& data) {
+  TensorDistAttrProto proto;
+  PADDLE_ENFORCE_EQ(
+      proto.ParseFromString(data),
+      true,
+      errors::InvalidArgument(
+          "Failed to parse tensor dist attr from string: %s.", data));
+  from_proto(proto);
+}
+
+bool operator==(const TensorDistAttr& lhs, const TensorDistAttr& rhs) {
+  if (lhs.process_mesh() != rhs.process_mesh()) {
+    return false;
+  }
+  if (lhs.multi_dims_mapping() != rhs.multi_dims_mapping()) {
+    return false;
+  }
+  if (lhs.batch_dim() != rhs.batch_dim()) {
+    return false;
+  }
+  if (lhs.chunk_id() != rhs.chunk_id()) {
+    return false;
+  }
+  if (lhs.dynamic_dims() != rhs.dynamic_dims()) {
+    return false;
+  }
+  if (lhs.partial_status() != rhs.partial_status()) {
+    return false;
+  }
+  if (lhs.split_factor() != rhs.split_factor()) {
+    return false;
+  }
+  return true;
+}
+
+std::string TensorDistAttr::partial_status_string() const {
+  std::string partial_status_str = "[";
+  for (auto& itr : partial_status_) {
+    partial_status_str += "Partial(dims:" + std::to_string(itr.first) + ", " +
+                          ReduceTypeStrings[static_cast<int>(itr.second)] +
+                          "), ";
+  }
+  partial_status_str += "]";
+  return partial_status_str;
+}
+
+bool TensorDistAttr::empty() const {
+  // dims_mapping is empty when the tensor is 0-dim, but it is also be valid.
+  return process_mesh_.empty();
+}
+
+std::vector<std::shared_ptr<PlacementStatus>> TensorDistAttr::to_placement()
+    const {
+  auto ndim = process_mesh_.ndim();
+  std::vector<std::shared_ptr<PlacementStatus>> placement(
+      ndim, std::make_shared<ReplicatedStatus>());
+  for (size_t i = 0; i < dims_mapping_.size(); ++i) {
+    const auto& cur_dims = dims_mapping_.at(i);
+    for (size_t j = 0; j < cur_dims.size(); ++j) {
+      int64_t cur_dim = cur_dims.at(j);
+      PADDLE_ENFORCE_LT(
+          cur_dim,
+          ndim,
+          errors::InvalidArgument(
+              "Split axis %ld can not exceed the ndim of process_mesh %ld",
+              cur_dim,
+              ndim));
+      placement[cur_dim] = std::make_shared<ShardStatus>(i, j);
+    }
+  }
+  for (auto& itr : partial_status_) {
+    PADDLE_ENFORCE_LT(
+        itr.first,
+        ndim,
+        errors::InvalidArgument(
+            "Partial axis %ld can not exceed the ndim of process_mesh %ld",
+            itr.first,
+            ndim));
+    placement[itr.first] = std::make_shared<PartialStatus>(itr.second);
+  }
+  return placement;
+}
+
+bool TensorDistAttr::is_replicated(int64_t mesh_axis) const {
+  auto placement = ToPlacements(*this);
+  if (mesh_axis == -1) {
+    return std::all_of(
+        placement.begin(), placement.end(), [](std::shared_ptr<Placement> p) {
+          return p->is_replicated();
+        });
+  } else {
+    return placement[mesh_axis]->is_replicated();
+  }
+}
+
+bool TensorDistAttr::is_shard(int64_t mesh_axis,
+                              std::optional<int64_t> tensor_axis) const {
+  auto placement = ToPlacements(*this);
+  if (mesh_axis == -1) {
+    return std::any_of(placement.begin(),
+                       placement.end(),
+                       [tensor_axis](std::shared_ptr<Placement> p) {
+                         return p->is_shard(tensor_axis);
+                       });
+  } else {
+    return placement[mesh_axis]->is_shard(tensor_axis);
+  }
+}
+
+bool TensorDistAttr::is_partial(int64_t mesh_axis) const {
+  if (mesh_axis == -1) {
+    return !partial_status_.empty();
+  } else {
+    return partial_status_.count(mesh_axis) > 0;
+  }
+}
+
+void TensorDistAttr::set_skip_check_mesh(bool skip) { skip_check_mesh_ = skip; }
+
+bool TensorDistAttr::is_co_shard() const {
+  for (const auto& mesh_dims : dims_mapping_) {
+    if (mesh_dims.size() > 1) {
+      return true;
+    }
+  }
+  return false;
+}
+
+TensorDistAttr::DimMapProxy& TensorDistAttr::DimMapProxy::operator=(
+    const std::vector<std::vector<int64_t>>& dims_mapping) {
+  dims_mapping_2d->resize(dims_mapping.size());
+  *dims_mapping_2d = dims_mapping;
+  return *this;
+}
+
+TensorDistAttr::DimMapProxy& TensorDistAttr::DimMapProxy::operator=(
+    const std::vector<int64_t>& dims_mapping) {
+  dims_mapping_1d = dims_mapping;
+  sync_2d_map();
+  VLOG(4) << "Set 1d dims_mapping, Sync 2d. 1d "
+          << auto_parallel::str_join(dims_mapping_1d) << " , 2d  "
+          << auto_parallel::str_join(*dims_mapping_2d);
+  return *this;
+}
+
+TensorDistAttr::DimMapProxy& TensorDistAttr::DimMapProxy::operator=(
+    const DimMapProxy& other) {
+  if (this == &other) {
+    return *this;
+  }
+  *dims_mapping_2d = (*other.dims_mapping_2d);
+  return *this;
+}
+
+bool TensorDistAttr::DimMapProxy::operator==(const DimMapProxy& other) {
+  return (*dims_mapping_2d) == (*other.dims_mapping_2d);
+}
+
+bool TensorDistAttr::DimMapProxy::operator!=(const DimMapProxy& other) {
+  return !this->operator==(other);
+}
+
+TensorDistAttr::DimMapProxy::operator const std::vector<std::vector<int64_t>>&()
+    const {
+  return *dims_mapping_2d;
+}
+
+TensorDistAttr::DimMapProxy::operator const std::vector<int64_t>&() const {
+  sync_1d_map();
+  VLOG(4) << "Get 1d dims_mapping, Sync 1d. 1d is "
+          << auto_parallel::str_join(dims_mapping_1d) << ", 2d is "
+          << auto_parallel::str_join(*dims_mapping_2d);
+  return dims_mapping_1d;
+}
+
+void TensorDistAttr::DimMapProxy::sync_1d_map() const {
+  dims_mapping_1d.resize(dims_mapping_2d->size());
+  for (size_t i = 0; i < dims_mapping_2d->size(); ++i) {
+    size_t num_mesh_dim = dims_mapping_2d->at(i).size();
+    if (num_mesh_dim <= 1) {
+      dims_mapping_1d[i] =
+          (*dims_mapping_2d)[i].empty() ? -1 : (*dims_mapping_2d)[i][0];
+      continue;
+    }
+
+    int64_t max_mesh_dim = (*dims_mapping_2d)[i][0];
+    int64_t max_mesh_dim_size = process_mesh.shape()[max_mesh_dim];
+
+    for (size_t j = 1; j < num_mesh_dim; ++j) {
+      int64_t cur_mesh_dim = (*dims_mapping_2d)[i][j];
+      int64_t cur_mesh_dim_size = process_mesh.shape()[cur_mesh_dim];
+
+      if (cur_mesh_dim_size > max_mesh_dim_size) {
+        max_mesh_dim = cur_mesh_dim;
+        max_mesh_dim_size = cur_mesh_dim_size;
+      }
+    }
+
+    LOG(WARNING) << "There are " << num_mesh_dim << " shared on tensor dim "
+                 << i << ". Now fallback to sharding by mesh dim "
+                 << max_mesh_dim << ".";
+
+    dims_mapping_1d[i] = max_mesh_dim;
+  }
+}
+
+void TensorDistAttr::DimMapProxy::sync_2d_map() {
+  dims_mapping_2d->resize(dims_mapping_1d.size());
+  for (size_t i = 0; i < dims_mapping_1d.size(); ++i) {
+    if (dims_mapping_1d.at(i) == -1) {
+      (*dims_mapping_2d)[i] = {};
+    } else {
+      (*dims_mapping_2d)[i] = {dims_mapping_1d[i]};
+    }
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/dist_attr.h

@@ -0,0 +1,289 @@
+/* Copyright (c) 2023 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. */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <iostream>
+#include <map>
+#include <optional>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/common/reduce_type.h"
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/utils/flat_hash_map.h"
+#include "paddle/utils/test_macros.h"
+
+namespace phi {
+namespace distributed {
+
+namespace auto_parallel {
+class TensorDistAttrProto;
+}
+
+constexpr int kReplicateDim = -1;
+
+class PlacementStatus {
+ public:
+  virtual ~PlacementStatus() = default;
+
+  virtual bool is_shard(int64_t axis = -1) const { return false; }
+  virtual bool is_partial() const { return false; }
+  virtual bool is_replicated() const { return false; }
+};
+
+class ReplicatedStatus final : public PlacementStatus {
+ public:
+  bool is_replicated() const override { return true; }
+};
+
+class PartialStatus final : public PlacementStatus {
+ public:
+  PartialStatus(ReduceType type) : type_(type) {}
+  bool is_partial() const override { return true; }
+  ReduceType get_reduce_type() const { return type_; }
+
+ private:
+  ReduceType type_{ReduceType::kRedSum};
+};
+
+class ShardStatus final : public PlacementStatus {
+ public:
+  explicit ShardStatus(int64_t axis) : axis_(axis) {}
+  ShardStatus(int64_t axis, int64_t co_shard_order)
+      : axis_(axis), co_shard_order_(co_shard_order) {}
+  bool is_shard(int64_t axis = -1) const override {
+    if (axis == -1) {
+      return true;
+    } else {
+      return axis == axis_;
+    }
+  }
+  int64_t get_axis() const { return axis_; }
+
+ private:
+  int64_t axis_{-1};
+  int64_t co_shard_order_{0};
+};
+
+class PADDLE_API TensorDistAttr {
+ public:
+  TensorDistAttr() = default;
+
+  explicit TensorDistAttr(const std::vector<int64_t>& tensor_shape);
+
+  TensorDistAttr(const TensorDistAttr& dist_attr);
+
+  TensorDistAttr& operator=(const TensorDistAttr& dist_attr);
+
+  void copy_from(const TensorDistAttr& dist_attr);
+
+  const ProcessMesh& process_mesh() const { return process_mesh_; }
+
+  void set_process_mesh(const ProcessMesh& process_mesh);
+
+  const std::vector<int64_t>& dims_mapping() const {
+    return dims_mapping_proxy;
+  }
+
+  const std::vector<std::vector<int64_t>>& multi_dims_mapping() const {
+    return dims_mapping_proxy;
+  }
+
+  void set_dims_mapping(const std::vector<int64_t>& dims_mapping);
+
+  void set_dims_mapping(const std::vector<std::vector<int64_t>>& dims_mapping);
+
+  const auto_parallel::SplitFactor& split_factor() const {
+    return split_factor_map_;
+  }
+
+  int64_t get_split_factor(int64_t mesh_dim) const;
+
+  void set_split_factor(int64_t mesh_dim, int64_t split_factor);
+
+  void clear_split_factor(int64_t mesh_dim);
+
+  // return vector of mesh dims on which the this tensor is partial on
+  const std::set<int64_t> partial_dims() const;
+
+  const paddle::flat_hash_map<int64_t, ReduceType>& partial_status() const {
+    return partial_status_;
+  }
+
+  // by map
+  void set_partial_status(
+      const paddle::flat_hash_map<int64_t, ReduceType>& partial_status);
+
+  // by each dim
+  void set_partial_status(const std::vector<int64_t>& dims,
+                          const ReduceType& type = ReduceType::kRedSum);
+  // all
+  void clean_partial_status();
+
+  // clean by dims
+  void clean_partial_dims(const std::vector<int64_t>& dims);
+
+  void set_default_dims_mapping(const std::vector<int64_t>& tensor_shape);
+
+  int64_t batch_dim() const { return batch_dim_; }
+
+  void set_batch_dim(int64_t batch_dim);
+
+  const int64_t& chunk_id() const { return chunk_id_; }
+
+  void set_chunk_id(const int64_t& chunk_id);
+
+  const std::vector<bool>& dynamic_dims() const { return dynamic_dims_; }
+
+  void set_dynamic_dims(const std::vector<bool>& dynamic_dims);
+
+  void set_default_dynamic_dims(const std::vector<int64_t>& tensor_shape);
+
+  void set_default_dynamic_dims(int64_t tensor_shape_size);
+
+  const std::map<std::string, bool>& annotated() const { return annotated_; }
+
+  void set_annotated(const std::map<std::string, bool>& annotated);
+
+  bool is_annotated(const std::string& name) const {
+    return annotated_.count(name) == 1 && annotated_.at(name) == true;
+  }
+
+  void mark_annotated(const std::string& name);
+
+  void clear_annotated() { annotated_.clear(); }
+
+  bool verify_process_mesh(const ProcessMesh& process_mesh) const;
+
+  bool verify_dims_mapping(
+      const std::vector<std::vector<int64_t>>& dims_mapping,
+      const std::vector<int64_t>& tensor_shape) const;
+
+  bool verify_batch_dim(int64_t dim,
+                        const std::vector<int64_t>& tensor_shape) const;
+
+  bool verify_dynamic_dims(const std::vector<bool>& dynamic_dims,
+                           const std::vector<int64_t>& tensor_shape) const;
+
+  bool verify_annotated(const std::map<std::string, bool>& annotated) const;
+
+  bool verify_partial_status() const;
+
+  bool verify(const std::vector<int64_t>& tensor_shape) const;
+  bool verify_dynamic(const std::vector<int64_t>& tensor_shape) const;
+
+  // TensorDistAttr from_string(const std::string& dist_str);
+  std::string to_string() const;
+  std::string partial_status_string() const;
+
+  // in partial-support-stage-I partial will always be a runtime attribute,
+  // there is not need to serialize it. support the partial serialization in
+  // future partial-support-stage-II.
+  void from_proto(const auto_parallel::TensorDistAttrProto& proto);
+
+  void to_proto(auto_parallel::TensorDistAttrProto* proto) const;
+
+  std::string serialize_to_string();
+
+  void parse_from_string(const std::string& data);
+
+  bool empty() const;
+
+  std::vector<std::shared_ptr<PlacementStatus>> to_placement() const;
+
+  // if mesh_axis is -1, check if tensor is replicated on whole process_mesh
+  // if mesh_axis is not -1, check only on specific axis.
+  bool is_replicated(int64_t mesh_axis = -1) const;
+
+  // if mesh_axis is -1, check if tensor is shard on whole process_mesh
+  // if mesh_axis is not -1, check only on specific axis
+  // if tensor_axis is not -1, return true only if the shard axis equal to
+  // tensor_axis.
+  bool is_shard(int64_t mesh_axis = -1,
+                std::optional<int64_t> dim = std::nullopt) const;
+
+  // if mesh_axis is -1, check if tensor is partial on whole process_mesh
+  // if mesh_axis is not -1, check only on specific axis.
+  bool is_partial(int64_t mesh_axis = -1) const;
+
+  void set_skip_check_mesh(bool skip);
+
+  bool skip_check_mesh() const { return skip_check_mesh_; }
+
+  bool is_co_shard() const;
+
+ private:
+  // delete it after all 1d vector dims_mapping_ have been upgraded to 2d.
+  class PADDLE_API DimMapProxy final {
+   public:
+    DimMapProxy(std::vector<std::vector<int64_t>>* dims_mapping_2d,
+                const ProcessMesh& process_mesh)
+        : dims_mapping_2d(dims_mapping_2d), process_mesh(process_mesh) {}
+
+    DimMapProxy& operator=(
+        const std::vector<std::vector<int64_t>>& dims_mapping);
+    DimMapProxy& operator=(const std::vector<int64_t>& dims_mapping);
+    DimMapProxy& operator=(const DimMapProxy& other);
+    bool operator==(const DimMapProxy& other);
+    bool operator!=(const DimMapProxy& other);
+    operator const std::vector<std::vector<int64_t>>&() const;
+    operator const std::vector<int64_t>&() const;
+
+   private:
+    void sync_1d_map() const;
+    void sync_2d_map();
+    mutable std::vector<int64_t> dims_mapping_1d;
+    std::vector<std::vector<int64_t>>* dims_mapping_2d;
+    const ProcessMesh& process_mesh;
+  };
+
+  static std::vector<std::string> fields_;
+  ProcessMesh process_mesh_;
+  std::vector<bool> dynamic_dims_;
+  std::map<std::string, bool> annotated_;
+  int64_t batch_dim_{0};
+  int64_t chunk_id_{0};
+  // partial map would be small (less than mesh.size)
+  // iterate operation (copy and comparison) would more frequency than random
+  // element access. <key: dim on mesh, value: reduce type>
+  paddle::flat_hash_map<int64_t, ReduceType> partial_status_;
+  auto_parallel::SplitFactor split_factor_map_;
+  // The flag indicates whether to skip checking the process mesh.
+  bool skip_check_mesh_ = false;
+
+  std::vector<std::vector<int64_t>> dims_mapping_;
+  // for short time, backward compatible for existing spmd relus.
+  DimMapProxy dims_mapping_proxy{&dims_mapping_, process_mesh_};
+};
+
+inline std::ostream& operator<<(std::ostream& os, const TensorDistAttr& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+PADDLE_API bool operator==(const TensorDistAttr& lhs,
+                           const TensorDistAttr& rhs);
+
+inline bool operator!=(const TensorDistAttr& lhs, const TensorDistAttr& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/dist_mapper.cc

@@ -0,0 +1,144 @@
+/* Copyright (c) 2022 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. */
+
+#include <algorithm>
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_mapper.h"
+#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+
+namespace phi::distributed::auto_parallel {
+
+void DistributedMapper::set_process_id_to_device_ids(
+    const std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>&
+        process_id_to_device_ids) {
+  std::vector<std::string> device_mesh_names;
+  device_mesh_names.reserve(device_meshes_.size());
+  for (const auto& item : device_meshes_) {
+    device_mesh_names.push_back(item.first);
+  }
+  for (const auto& item : process_id_to_device_ids) {
+    PADDLE_ENFORCE_GE(
+        item.first,
+        0,
+        errors::InvalidArgument(
+            "The process id %d must be greater than or equal to 0.",
+            item.first));
+    std::string device_mesh_name = item.second.first;
+    const std::vector<int64_t>& device_ids = item.second.second;
+    PADDLE_ENFORCE_EQ(
+        device_meshes_.count(device_mesh_name),
+        1,
+        errors::InvalidArgument(
+            "Cannot find the device mesh %d in device_mesh ids [%s].",
+            device_mesh_name,
+            str_join(device_mesh_names)));
+    PADDLE_ENFORCE_EQ(
+        has_duplicates(device_ids),
+        false,
+        errors::InvalidArgument(
+            "The mapped device ids [%s] of process_mesh %d must be unique.",
+            str_join(device_ids),
+            item.first));
+    const DeviceMesh& device_mesh = device_meshes_[device_mesh_name];
+    const std::vector<int64_t> cur_device_ids = device_mesh.device_ids();
+    for (int64_t device_id : device_ids) {
+      bool found =
+          std::find(cur_device_ids.begin(), cur_device_ids.end(), device_id) !=
+          cur_device_ids.end();
+      PADDLE_ENFORCE_EQ(
+          found,
+          true,
+          errors::InvalidArgument(
+              "The device id %d cannot be find in the device mesh [%s].",
+              device_id,
+              str_join(cur_device_ids)));
+    }
+  }
+  process_id_to_device_ids_ = process_id_to_device_ids;
+}
+
+DistributedMapper DistributedMapper::from_proto(
+    const DistributedMapperProto& proto) {
+  DistributedMapper dist_mapper;
+  for (int i = 0; i < proto.device_meshes_size(); ++i) {
+    dist_mapper.device_meshes_[proto.device_meshes(i).name()] =
+        DeviceMesh::from_proto(proto.device_meshes(i));
+  }
+  for (int i = 0; i < proto.process_id_to_device_ids_size(); ++i) {
+    int64_t process_id = proto.process_id_to_device_ids(i).process_id();
+    std::string device_mesh_name =
+        proto.process_id_to_device_ids(i).device_mesh_name();
+    std::vector<int64_t> device_ids;
+    int num_devices = proto.process_id_to_device_ids(i).device_ids_size();
+    device_ids.reserve(num_devices);
+    for (int j = 0; j < num_devices; ++j) {
+      device_ids.push_back(proto.process_id_to_device_ids(i).device_ids(j));
+    }
+    dist_mapper.process_id_to_device_ids_[process_id].first = device_mesh_name;
+    dist_mapper.process_id_to_device_ids_[process_id].second = device_ids;
+  }
+  return dist_mapper;
+}
+
+void DistributedMapper::to_proto(DistributedMapperProto* proto) const {
+  for (const auto& item : device_meshes_) {
+    proto->mutable_device_meshes()->Add()->CopyFrom(
+        distributed::to_proto(item.second));
+  }
+  for (const auto& outer : process_id_to_device_ids_) {
+    auto proto_item = proto->mutable_process_id_to_device_ids()->Add();
+    proto_item->set_process_id(outer.first);
+    proto_item->set_device_mesh_name(outer.second.first);
+    for (const auto& inner : outer.second.second) {
+      proto_item->add_device_ids(inner);
+    }
+  }
+}
+
+std::string DistributedMapper::to_string() const {
+  std::string mapper_str = "{device_meshes: [";
+  for (const auto& item : device_meshes_) {
+    mapper_str += item.second.to_string() + ", ";
+  }
+  mapper_str.replace(mapper_str.size() - 2, 2, "]");
+
+  mapper_str += "\nprocess_id_to_device_ids: [";
+  for (const auto& item : process_id_to_device_ids_) {
+    mapper_str += "{";
+    mapper_str +=
+        "process_id: " + std::to_string(item.first) + ", device_ids: [";
+    for (const auto& device_id : item.second.second) {
+      mapper_str +=
+          "{" + item.second.first + ", " + std::to_string(device_id) + "}, ";
+    }
+    mapper_str.replace(mapper_str.size() - 2, 2, "]");
+    mapper_str += "}, ";
+  }
+  mapper_str.replace(mapper_str.size() - 2, 2, "]");
+  mapper_str += "}";
+  return mapper_str;
+}
+
+bool operator==(const DistributedMapper& lhs, const DistributedMapper& rhs) {
+  if (lhs.device_meshes() != rhs.device_meshes()) {
+    return false;
+  }
+  if (lhs.process_id_to_device_ids() != rhs.process_id_to_device_ids()) {
+    return false;
+  }
+  return true;
+}
+
+}  // namespace phi::distributed::auto_parallel

paddle/phi/core/distributed/auto_parallel/dist_mapper.h

@@ -0,0 +1,75 @@
+/* Copyright (c) 2022 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. */
+#pragma once
+
+#include <utility>
+
+#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+
+namespace phi {
+namespace distributed {
+namespace auto_parallel {
+
+class DistributedMapperProto;
+
+class PADDLE_API DistributedMapper {
+ public:
+  DistributedMapper() = default;
+
+  const std::map<std::string, DeviceMesh>& device_meshes() const {
+    return device_meshes_;
+  }
+
+  const DeviceMesh& device_mesh(const std::string& name) const {
+    return device_meshes_.at(name);
+  }
+
+  void add_device_mesh(const DeviceMesh& device_mesh) {
+    device_meshes_[device_mesh.name()] = device_mesh;
+  }
+
+  const std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>&
+  process_id_to_device_ids() const {
+    return process_id_to_device_ids_;
+  }
+
+  void set_process_id_to_device_ids(
+      const std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>&
+          process_id_to_device_ids);
+
+  // DistributedMapper from_string(const std::string& mapper_str);
+  std::string to_string() const;
+
+  static DistributedMapper from_proto(const DistributedMapperProto& proto);
+  void to_proto(DistributedMapperProto* proto) const;
+
+ private:
+  std::map<std::string, DeviceMesh> device_meshes_;
+  std::map<int64_t, std::pair<std::string, std::vector<int64_t>>>
+      process_id_to_device_ids_;
+};
+
+PADDLE_API bool operator==(const DistributedMapper& lhs,
+                           const DistributedMapper& rhs);
+
+inline std::ostream& operator<<(std::ostream& os,
+                                const DistributedMapper& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+}  // namespace auto_parallel
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.cc

@@ -0,0 +1,53 @@
+/* Copyright (c) 2023 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. */
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+
+namespace phi::distributed {
+
+DDim DistMetaTensor::dims() const {
+  // member values in tensor_ have higher priority than those in DistMetaTensor
+  if (tensor_ != nullptr) {
+    PADDLE_ENFORCE_EQ(this->is_dist(),
+                      true,
+                      common::errors::InvalidArgument(
+                          "The current MetaTensor doesn't contains "
+                          "DistTensor when call `dims` method."));
+    return MetaTensor::dims();
+  } else {
+    return dims_;
+  }
+}
+
+const distributed::TensorDistAttr& DistMetaTensor::dist_attr() const {
+  // member values in tensor_ have higher priority than those in DistMetaTensor
+  if (tensor_ != nullptr) {
+    PADDLE_ENFORCE_EQ(this->is_dist(),
+                      true,
+                      common::errors::InvalidArgument(
+                          "The current MetaTensor doesn't contains "
+                          "DistTensor when call `dist_attr` method."));
+    return static_cast<phi::distributed::DistTensor*>(tensor_)->dist_attr();
+  } else {
+    return dist_attr_;
+  }
+}
+
+bool DistMetaTensor::initialized() const {
+  return tensor_ != nullptr || dist_attr_ != TensorDistAttr();
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h

@@ -0,0 +1,72 @@
+/* Copyright (c) 2023 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. */
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/meta_tensor.h"
+
+namespace phi {
+namespace distributed {
+
+class DistMetaTensor : public MetaTensor {
+ public:
+  DistMetaTensor() : MetaTensor() {}
+
+  // supporting implicit construction is easier to use
+  DistMetaTensor(TensorBase* tensor)  // NOLINT
+      : MetaTensor(tensor) {}
+  DistMetaTensor(const TensorBase& tensor)  // NOLINT
+      : MetaTensor(tensor) {}
+  DistMetaTensor(const TensorBase* tensor)  // NOLINT
+      : MetaTensor(tensor) {}
+  DistMetaTensor(TensorBase& tensor)  // NOLINT
+      : MetaTensor(tensor) {}
+  // For static mode only
+  DistMetaTensor(const DDim& dims, const TensorDistAttr& dist_attr)
+      : dims_(dims), dist_attr_(dist_attr) {}
+
+  DistMetaTensor(DistMetaTensor&&) = default;
+  DistMetaTensor& operator=(DistMetaTensor&&) = default;
+  DistMetaTensor(const DistMetaTensor&) = default;
+  DistMetaTensor& operator=(const DistMetaTensor&) = default;
+
+  virtual ~DistMetaTensor() = default;
+
+  PADDLE_API DDim dims() const override;
+
+  const distributed::TensorDistAttr& dist_attr() const;
+
+  PADDLE_API bool initialized() const override;
+
+ private:
+  /**
+   * Note: When using the semi-automatic parallel segmentation derivation rules
+   * of the static graph, in order to facilitate the packaging of the input
+   * parameters of the construction, the DistMetaTensor is inherited and
+   * encapsulated, and the class members dims_ and dist_attr_ are added to it.
+   *
+   * The information contained in these two members is also in the tensor of the
+   * meta_tensor of the base class, and there is redundancy.
+   *
+   * We need to pay attention when using it to ensure the consistency.
+   * These two members are read-only, and their values cannot be changed
+   * after construction. To change their values, they need to be set
+   * directly in tensor_*/
+  DDim dims_;
+  TensorDistAttr dist_attr_;
+};
+
+}  // namespace distributed
+}  //  namespace phi

paddle/phi/core/distributed/auto_parallel/dist_tensor.cc

@@ -0,0 +1,340 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/api/lib/data_transform.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+
+namespace phi::distributed {
+
+inline void check_defined(const DistTensor& dist_tensor,
+                          std::string method_hint) {
+  PADDLE_ENFORCE_EQ(
+      dist_tensor.defined(),
+      true,
+      common::errors::Unimplemented(
+          "DistTensor is not defined yet when `%s` method is called.",
+          method_hint));
+}
+
+TensorDistAttr ToTensorDistAttr(const ProcessMesh& process_mesh,
+                                const Placements& placements,
+                                const DDim& dims) {
+  TensorDistAttr dist_attr(vectorize(dims));
+  // Step1: set process_mesh
+  dist_attr.set_process_mesh(process_mesh);
+
+  // Step2: set dim_mapping
+  std::vector<std::vector<int64_t>> dim_mapping(dims.size());
+  for (size_t i = 0; i < placements.size(); i++) {
+    const auto& cur_placement = placements[i];
+    if (!cur_placement->is_shard()) {
+      continue;
+    }
+    const auto& shard = dynamic_cast<const Shard&>(*cur_placement);
+    dim_mapping[shard.get_dim()].push_back(i);
+    dist_attr.set_split_factor(i, shard.get_split_factor());
+  }
+
+  auto compare_functor = [&](size_t a, size_t b) {
+    const auto& shard_a = dynamic_cast<const Shard&>(*placements[a]);
+    const auto& shard_b = dynamic_cast<const Shard&>(*placements[b]);
+    return shard_a.get_co_shard_order() < shard_b.get_co_shard_order();
+  };
+
+  for (size_t i = 0; i < dim_mapping.size(); i++) {
+    auto& mesh_dims = dim_mapping[i];
+    if (mesh_dims.size() > 1) {
+      std::sort(mesh_dims.begin(), mesh_dims.end(), compare_functor);
+    }
+  }
+  dist_attr.set_dims_mapping(dim_mapping);
+
+  // Step3: set partial_status
+  paddle::flat_hash_map<int64_t, ReduceType> partial_status;
+  for (size_t i = 0; i < placements.size(); ++i) {
+    auto& p = placements[i];
+    if (p->is_partial()) {
+      partial_status.insert({i, dynamic_cast<Partial&>(*p).get_reduce_type()});
+    }
+  }
+  dist_attr.set_partial_status(partial_status);
+
+  // Step4: mark annotated
+  dist_attr.mark_annotated("process_mesh");
+  dist_attr.mark_annotated("dims_mapping");
+  return dist_attr;
+}
+
+Placements ToPlacements(const TensorDistAttr& dist_attr) {
+  auto& process_mesh = dist_attr.process_mesh();
+  Placements placements;
+  placements.resize(process_mesh.ndim(), std::make_shared<Replicate>());
+
+  auto& partial_status = dist_attr.partial_status();
+  for (const auto& pair : partial_status) {
+    placements[pair.first] = std::make_shared<Partial>(pair.second);
+  }
+
+  const std::vector<std::vector<int64_t>>& dim_mapping =
+      dist_attr.multi_dims_mapping();
+  for (size_t t_dim = 0; t_dim < dim_mapping.size(); t_dim++) {
+    auto m_dims = dim_mapping[t_dim];
+
+    bool is_co_shard = m_dims.size() > 1;
+
+    for (size_t idx = 0; idx < m_dims.size(); idx++) {
+      int64_t mesh_dim = m_dims[idx];
+      int64_t split_factor = dist_attr.get_split_factor(mesh_dim);
+      auto& p = placements.at(mesh_dim);
+
+      if (p->is_shard()) {
+        PADDLE_THROW(common::errors::PreconditionNotMet(
+            "ProcessMesh dimension can't be mapped to two  dimension of the "
+            "same tensor: {%d} and {%d}",
+            t_dim,
+            dynamic_cast<Shard&>(*p).get_dim()));
+      } else if (p->is_partial()) {
+        PADDLE_THROW(common::errors::PreconditionNotMet(
+            "ProcessMesh dimension {%d} cannot be both shard and partial!",
+            mesh_dim));
+      }
+      // TODO(lfw): add mesh_dim >= 0 check
+      placements[mesh_dim] = is_co_shard
+                                 ? std::make_shared<CoShard>(t_dim, idx)
+                                 : std::make_shared<Shard>(t_dim, split_factor);
+    }
+  }
+  return placements;
+}
+
+DistTensor::DistTensor() : value_(std::make_shared<DenseTensor>()) {}
+
+DistTensor::DistTensor(DataType dtype)
+    : value_(std::make_shared<DenseTensor>(dtype)) {}
+
+DistTensor::DistTensor(const std::shared_ptr<DenseTensor>& global_value,
+                       const TensorDistAttr& dist_attr)
+    : global_dims_(global_value->dims()), dist_attr_(dist_attr) {
+  process_mesh_ = dist_attr_.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+
+  // If the current rank doesn't in process_mesh, we should create an
+  // uninitialized tensor only with tensor_meta.
+  if (IsCurRankInMesh(dist_attr.process_mesh())) {
+    if (!dist_attr.is_replicated()) {
+      value_ = std::make_shared<DenseTensor>();
+      // 1. create replicated global tensor
+      TensorDistAttr replicated_dist_attr(
+          common::vectorize(global_value->dims()));
+      replicated_dist_attr.set_process_mesh(dist_attr.process_mesh());
+      DistTensor replicated_tensor(global_value, replicated_dist_attr);
+
+      // 2. reshard from replicated to other state
+      VLOG(4) << "Reshard tensor: "
+              << paddle::experimental::ReshardDebugInfo(replicated_tensor,
+                                                        dist_attr);
+      auto* func = ChooseProperReshardFunction(replicated_tensor, dist_attr);
+      auto* dev_ctx = DeviceContextPool::Instance().Get(global_value->place());
+      func->Eval(dev_ctx, replicated_tensor, dist_attr, this);
+    } else {
+      value_ = global_value;
+    }
+  } else {
+    value_ = std::make_shared<DenseTensor>(
+        std::make_shared<phi::Allocation>(nullptr, 0, global_value->place()),
+        DenseTensorMeta(global_value->meta()));
+  }
+}
+
+DistTensor::DistTensor(const std::shared_ptr<DenseTensor>& local_value,
+                       const DDim& global_dims,
+                       const ProcessMesh& process_mesh,
+                       const Placements& placements)
+    : global_dims_(global_dims),
+      process_mesh_(process_mesh),
+      placements_(placements) {
+  dist_attr_ = ToTensorDistAttr(process_mesh_, placements_, global_dims_);
+  if (IsCurRankInMesh(process_mesh)) {
+    value_ = local_value;
+  } else {
+    value_ = std::make_shared<DenseTensor>(
+        std::make_shared<phi::Allocation>(nullptr, 0, local_value->place()),
+        DenseTensorMeta(local_value->dtype(), make_ddim({0})));
+  }
+}
+
+DistTensor::DistTensor(const std::shared_ptr<DenseTensor>& local_value,
+                       const DDim& global_dims,
+                       const TensorDistAttr& dist_attr)
+    : global_dims_(global_dims), dist_attr_(dist_attr) {
+  process_mesh_ = dist_attr_.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+  if (IsCurRankInMesh(process_mesh_)) {
+    value_ = local_value;
+  } else {
+    value_ = std::make_shared<DenseTensor>(
+        std::make_shared<phi::Allocation>(nullptr, 0, local_value->place()),
+        DenseTensorMeta(local_value->dtype(), global_dims_));
+  }
+}
+
+DistTensor::DistTensor(const std::shared_ptr<DenseTensor>& global_value,
+                       const ProcessMesh& process_mesh,
+                       const Placements& placements)
+    : global_dims_(global_value->dims()) {
+  process_mesh_ = process_mesh;
+  placements_ = placements;
+  // If the dims.size() == -1, the dims=[0] by default, which is not consistent
+  // and will cause ToTensorDistAttr's error.
+  if (global_dims_ == DDim()) {
+    global_dims_ = make_ddim({});
+  }
+  dist_attr_ = ToTensorDistAttr(process_mesh_, placements_, global_dims_);
+
+  // If the current rank doesn't in process_mesh, we should create an
+  // uninitialized tensor only with dist_tensor_meta_.
+  if (IsCurRankInMesh(process_mesh)) {
+    if (!dist_attr_.is_replicated()) {
+      if (global_value->initialized()) {
+        value_ = std::make_shared<DenseTensor>();
+        // 1. create replicated global tensor
+        TensorDistAttr replicated_dist_attr(
+            common::vectorize(global_value->dims()));
+        replicated_dist_attr.set_process_mesh(process_mesh);
+        DistTensor replicated_tensor(global_value, replicated_dist_attr);
+
+        // 2. reshard from replicated to other state
+        VLOG(4) << "Reshard tensor: "
+                << paddle::experimental::ReshardDebugInfo(replicated_tensor,
+                                                          dist_attr_);
+        auto* func = ChooseProperReshardFunction(replicated_tensor, dist_attr_);
+        auto* dev_ctx =
+            DeviceContextPool::Instance().Get(global_value->place());
+        func->Eval(dev_ctx, replicated_tensor, dist_attr_, this);
+      } else {
+        // For lazy init, the global value is an uninitialized tensor.
+        // Just infer the local shape of the dist tensor.
+        value_ = global_value;
+        value_->Resize(
+            InferShapeForReshardFromReplicate(global_value, dist_attr_));
+      }
+    } else {
+      value_ = global_value;
+    }
+  } else {
+    value_ = std::make_shared<DenseTensor>(
+        std::make_shared<phi::Allocation>(nullptr, 0, global_value->place()),
+        DenseTensorMeta(global_value->meta()));
+  }
+}
+
+DistTensor::DistTensor(const DDim& dims, const TensorDistAttr& dist_attr)
+    : global_dims_(dims),
+      dist_attr_(dist_attr),
+      value_(std::make_shared<DenseTensor>()) {
+  process_mesh_ = dist_attr.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+}
+
+void DistTensor::unsafe_set_dims(const DDim& dims) {
+  if (this->initialized()) {
+    VLOG(6) << "You try to set an initialized DistTensor's global dims. "
+               "Make sure you are aware of where you change its dims.";
+  }
+  global_dims_ = dims;
+}
+
+void DistTensor::unsafe_set_dist_attr(const TensorDistAttr& dist_attr) {
+  if (this->initialized()) {
+    VLOG(6) << "You try to set an initialized DistTensor's dist attr. "
+               "Make sure you are aware of where you change its dist attr.";
+  }
+  dist_attr_ = dist_attr;
+  process_mesh_ = dist_attr.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+}
+
+int64_t DistTensor::numel() const {
+  // DistTensor with uninitialized local tensor can
+  // also have numel.
+  return product(global_dims_);
+}
+
+const DDim& DistTensor::local_dims() const {
+  check_defined(*this, "local_dims");
+  return value_->dims();
+}
+
+bool DistTensor::valid() const {
+  check_defined(*this, "valid");
+  return value_->valid();
+}
+
+bool DistTensor::defined() const { return value_->holder_ != nullptr; }
+
+bool DistTensor::has_allocation() const {
+  return value_->holder_ != nullptr && (value_->holder_->ptr() || numel() != 0);
+}
+
+bool DistTensor::initialized() const {
+  return value_->holder_ != nullptr && value_->holder_->ptr();
+}
+
+DataType DistTensor::dtype() const {
+  // DistTensor with uninitialized local tensor can
+  // also have dtype.
+  return value_->dtype();
+}
+
+DataLayout DistTensor::layout() const {
+  // DistTensor with uninitialized local tensor can
+  // also have layout.
+  return value_->layout();
+}
+
+const Place& DistTensor::place() const {
+  check_defined(*this, "place");
+  return value_->holder_->place();
+}
+
+void* DistTensor::AllocateFrom(Allocator* allocator,
+                               DataType dtype,
+                               size_t requested_size,
+                               bool fake_alloc) {
+  PADDLE_THROW(common::errors::Unavailable(
+      "The DistTensor Cannot allocate memory directly and needs to perform "
+      "memory operations through its DenseTensor value."));
+  return nullptr;
+}
+
+void DistTensor::unsafe_set_skip_check_mesh(bool skip) {
+  VLOG(6) << "You try to set an initialized DistTensor's dist attr. "
+             "Make sure you are aware of where you change its dist attr.";
+  dist_attr_.set_skip_check_mesh(skip);
+}
+
+void DistTensor::clear() {
+  if (value_) {
+    value_->clear();
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/dist_tensor.h

@@ -0,0 +1,216 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <memory>
+
+#include "paddle/phi/common/reduce_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/placement_types.h"
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace distributed {
+class ReshardFunction;
+class Shard;
+class Partial;
+class Replicate;
+
+PADDLE_API TensorDistAttr ToTensorDistAttr(const ProcessMesh& process_mesh,
+                                           const Placements& placements,
+                                           const DDim& dims);
+
+PADDLE_API Placements ToPlacements(const TensorDistAttr& dist_attr);
+
+class PADDLE_API DistTensor final
+    : public phi::TensorBase,
+      public phi::TypeInfoTraits<phi::TensorBase, DistTensor> {
+ public:
+  /// \brief Careful to create dist tensor using default constructor.
+  /// this should only used in reshard for now, and the dist properties
+  /// will be set by reshard later.
+  DistTensor();
+
+  /// \brief Construct a dist tensor based dtype.
+  /// \param dtype The dtype of the current tensor.
+  explicit DistTensor(DataType dtype);
+
+  /// \brief Construct a dist tensor based dense tensor.
+  /// \param global_value The global dense tensor of the current tensor.
+  /// \param dist_attr The distributed attributes of the current tensor.
+  DistTensor(const std::shared_ptr<DenseTensor>& global_value,
+             const TensorDistAttr& dist_attr);
+
+  /// \brief Construct a dist tensor based dense tensor.
+  /// \param process_mesh The process mesh of the current tensor.
+  /// \param placements The distributed placements of the current tensor.
+  DistTensor(const std::shared_ptr<DenseTensor>& global_value,
+             const ProcessMesh& process_mesh,
+             const Placements& placements);
+
+  /// \brief Construct a dist tensor based local dense tensor.
+  /// \param global_dims The global dim of the dist tensor.
+  /// \param dist_attr The distributed attributes of the current tensor.
+  DistTensor(const std::shared_ptr<DenseTensor>& local_value,
+             const DDim& global_dims,
+             const TensorDistAttr& dist_attr);
+
+  /// \brief Construct a dist tensor based local dense tensor.
+  /// \param global_dims The global dim of the dist tensor.
+  /// \param process_mesh The process mesh of the current tensor.
+  /// \param placements The distributed placements of the current tensor.
+  DistTensor(const std::shared_ptr<DenseTensor>& local_value,
+             const DDim& global_dims,
+             const ProcessMesh& process_mesh,
+             const Placements& placements);
+
+  /// \brief Construct a empty dist tensor (for infer spmd)
+  /// \param dims The global dimension of the current Tensor.
+  /// \param dist_attr The distributed attributes of the current tensor.
+  DistTensor(const DDim& dims, const TensorDistAttr& dist_attr);
+
+  /// \brief Destroy the tensor object and release exclusive resources.
+  virtual ~DistTensor() = default;
+
+  /// \brief Returns the name of the class for type traits.
+  /// \return The name of the class.
+  static const char* name() { return "DistTensor"; }
+
+  /// \brief Returns the global dims of the dist tensor.
+  /// \return The global dims of the dist tensor.
+  const DDim& dims() const override { return global_dims_; }
+
+  /// \brief Set the global dims of the dist tensor.
+  /// \return void
+  void unsafe_set_dims(const DDim& dims);
+
+  /// \brief Returns the dist attr of current dist tensor.
+  /// \return The TensorDistAttr's const reference
+  const TensorDistAttr& dist_attr() const { return dist_attr_; }
+
+  /// \brief Returns the process_mesh of current dist tensor.
+  /// \return The ProcessMesh's const reference
+  const ProcessMesh& process_mesh() const { return process_mesh_; }
+
+  /// \brief Returns the placements of current dist tensor.
+  /// \return The Placements's const reference
+  const Placements& placements() const { return placements_; }
+
+  /// \brief Returns the num_shard of current dist tensor.
+  /// \return int64_t
+  int64_t num_shard() const {
+    int64_t num_shard = 1;
+    const auto& mesh_shape = process_mesh_.shape();
+    for (size_t i = 0; i < placements_.size(); i++) {
+      if (placements_[i]->is_shard()) {
+        num_shard *= mesh_shape[i];
+      }
+    }
+    return num_shard;
+  }
+  /// \brief Set the dist attr of current dist tensor.
+  /// \return void
+  void unsafe_set_dist_attr(const TensorDistAttr& dist_attr);
+
+  /// \brief Returns the dense tensor value's const reference in dist tensor.
+  /// \return The DenseTensor value's const reference
+  const DenseTensor& value() const { return *value_; }
+
+  /// \brief Returns the shared_ptr of dense tensor value's in dist tensor.
+  /// \return The shared_ptr of dense tensor value
+  std::shared_ptr<DenseTensor> shared_value() { return value_; }
+
+  /// \brief Returns the mutable dense tensor value in dist tensor.
+  /// \note If DenseTensor value is modified externally, the corresponding
+  /// relationship between it and the current tensor's global dims and
+  /// dist attr may be destroyed, which may introduce some subtle bugs,
+  /// so you need to make sure to consider it thoroughly when using
+  /// this method.
+  /// \return The mutable pointer of DenseTensor value
+  DenseTensor* unsafe_mutable_value() { return value_.get(); }
+
+  /// \brief Returns the global dims of the dist tensor.
+  /// \return The global dims of the dist tensor.
+  const DDim& local_dims() const;
+
+  /// \brief Test whether the holder is created.
+  /// \return Whether the holder is created.
+  bool has_allocation() const override;
+
+  /// \brief Returns the global number of elements contained in tensor.
+  /// \return The number of elements contained in tensor.
+  int64_t numel() const override;
+
+  /// \brief Test whether the dense tensor value's storage is allocated.
+  /// \return Whether the dense tensor value's storage is allocated.
+  bool initialized() const override;
+
+  /// \brief Test whether the dense tensor value is defined.
+  /// \return Whether the dense tensor value is defined.
+  bool defined() const;
+
+  /// \brief Test whether the metadata is valid.
+  /// \return Whether the metadata is valid.
+  bool valid() const override;
+
+  /// \brief Returns the data type of the tensor.
+  /// \return The data type of the tensor.
+  DataType dtype() const override;
+
+  /// \brief Returns the data layout of the tensor.
+  /// \return The data layout of the tensor.
+  DataLayout layout() const override;
+
+  /// \brief Returns the data place of the tensor.
+  /// \return The data place of the tensor.
+  const Place& place() const override;
+
+  /// \brief Allocate memory with requested size from allocator.
+  /// \return The mutable data pointer value of type T.
+  void* AllocateFrom(Allocator* allocator,
+                     DataType dtype,
+                     size_t requested_size = 0,
+                     bool fake_alloc = false) override;
+
+  /// \brief Set the flag indicating whether to skip checking the process mesh.
+  /// \note Currently only used for the MoE apis,
+  /// it receives the inputs with different process meshes and outputs the dist
+  /// tensor with global process mesh.
+  /// \return void
+  void unsafe_set_skip_check_mesh(bool skip);
+
+  bool skip_check_mesh() const { return dist_attr_.skip_check_mesh(); }
+
+  void clear();
+
+ private:
+  friend class ReshardFunction;
+
+  // The global dimensions(shape), will move to DistTensorMeta
+  DDim global_dims_;
+  // The distributed attributes, will remove in the future
+  TensorDistAttr dist_attr_;
+  // The local DenseTensor value
+  std::shared_ptr<DenseTensor> value_;
+
+  ProcessMesh process_mesh_;
+
+  Placements placements_;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc

@@ -0,0 +1,206 @@
+/* Copyright (c) 2023 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. */
+
+#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
+
+namespace phi::distributed {
+
+InferSpmdContext::InferSpmdContext(
+    paddle::small_vector<DistMetaTensor, kInputSmallVectorSize> inputs,
+    paddle::small_vector<Attribute, kAttrSmallVectorSize> attrs) {
+  for (size_t i = 0; i < inputs.size(); i++) {
+    EmplaceBackInput(inputs[i]);
+  }
+  for (size_t i = 0; i < attrs.size(); i++) {
+    EmplaceBackAttr(attrs[i]);
+  }
+}
+
+void InferSpmdContext::EmplaceBackInput(DistMetaTensor input) {
+  int index = static_cast<int>(inputs_.size());
+  inputs_.emplace_back(std::move(input));
+  input_range_.emplace_back(std::pair<int, int>(index, index + 1));
+}
+
+void InferSpmdContext::EmplaceBackInputs(
+    paddle::small_vector<DistMetaTensor, kInputSmallVectorSize> inputs) {
+  int index = static_cast<int>(inputs_.size());
+  input_range_.emplace_back(std::pair<int, int>(index, index + inputs.size()));
+  inputs_.insert(inputs_.end(),
+                 std::make_move_iterator(inputs.begin()),
+                 std::make_move_iterator(inputs.end()));
+}
+
+void InferSpmdContext::EmplaceBackAttr(Attribute attr) {
+  attrs_.emplace_back(std::move(attr));
+}
+
+const DistMetaTensor& InferSpmdContext::InputAt(size_t idx) const {
+  return inputs_.at(idx);
+}
+
+template <typename AttrType>
+AttrType InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    return paddle::get<AttrType>(attrs_.at(idx));
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `%s`.",
+        attrs_.at(idx).type().name(),
+        std::type_index(typeid(AttrType)).name()));
+  }
+}
+
+template float InferSpmdContext::AttrAt(size_t idx) const;
+template int InferSpmdContext::AttrAt(size_t idx) const;
+template int64_t InferSpmdContext::AttrAt(size_t idx) const;
+template DataType InferSpmdContext::AttrAt(size_t idx) const;
+
+template <>
+bool InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    auto attr = attrs_.at(idx);
+    if (attr.type() == typeid(int)) {
+      return static_cast<bool>(paddle::get<int>(attr));
+    } else {
+      return paddle::get<bool>(attr);
+    }
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `bool`.",
+        attrs_.at(idx).type().name()));
+  }
+}
+
+template <>
+double InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    auto attr = attrs_.at(idx);
+    if (attr.type() == typeid(float)) {
+      return static_cast<double>(paddle::get<float>(attr));
+    } else {
+      return paddle::get<double>(attr);
+    }
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `double`.",
+        attrs_.at(idx).type().name()));
+  }
+}
+
+template <>
+std::vector<int> InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    auto attr = attrs_.at(idx);
+    if (attr.type() == typeid(std::vector<bool>)) {
+      std::vector<bool> val = PADDLE_GET_CONST(std::vector<bool>, attr);
+      return std::vector<int>(val.begin(), val.end());
+    } else if (attr.type() == typeid(std::vector<int64_t>) &&
+               paddle::get<std::vector<int64_t>>(attr).empty()) {
+      return std::vector<int>();
+    } else {
+      return paddle::get<std::vector<int>>(attr);
+    }
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `std::vector<int>`.",
+        attrs_.at(idx).type().name()));
+  }
+}
+
+template <>
+std::vector<int64_t> InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    auto attr = attrs_.at(idx);
+    if (attr.type() == typeid(std::vector<bool>)) {
+      std::vector<bool> val = PADDLE_GET_CONST(std::vector<bool>, attr);
+      return std::vector<int64_t>(val.begin(), val.end());
+    } else if (attr.type() == typeid(std::vector<int>)) {
+      std::vector<int> val = PADDLE_GET_CONST(std::vector<int>, attr);
+      return std::vector<int64_t>(val.begin(), val.end());
+    } else {
+      return PADDLE_GET_CONST(std::vector<int64_t>, attr);
+    }
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `std::vector<int64_t>`.",
+        attrs_.at(idx).type().name()));
+  }
+}
+
+template <>
+std::string InferSpmdContext::AttrAt(size_t idx) const {
+  try {
+    auto attr = attrs_.at(idx);
+    return PADDLE_GET_CONST(std::string, attr);
+  } catch (paddle::bad_variant_access const& e) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Attribute cast error in InferSpmd Context, the input attr type is "
+        "`%s`, but the expected attribute type is `std::string`.",
+        attrs_.at(idx).type().name()));
+  }
+}
+
+const Attribute& InferSpmdContext::AttrAt(size_t idx) const {
+  return attrs_.at(idx);
+}
+
+const std::pair<int, int>& InferSpmdContext::InputRangeAt(size_t idx) const {
+  return input_range_.at(idx);
+}
+
+const std::vector<const DistMetaTensor*> InferSpmdContext::InputsBetween(
+    size_t start, size_t end) const {
+  std::vector<const DistMetaTensor*> result;
+  result.reserve(end - start);
+  for (size_t i = start; i < end; ++i) {
+    auto& in = inputs_.at(i);
+    result.emplace_back(&in);
+    // result.emplace_back(in.initialized() ? &in : nullptr);
+  }
+
+  return result;
+}
+
+SpmdRuleFactory& SpmdRuleFactory::Instance() {
+  static SpmdRuleFactory g_spmd_rule_map;
+  return g_spmd_rule_map;
+}
+
+bool SpmdRuleFactory::ContainsSpmdRule(const std::string& kernel_name) const {
+  return spmd_rule_map_.count(kernel_name) > 0;
+}
+
+int SpmdRuleFactory::InsertSpmdRule(std::string kernel_name, SpmdRule rule) {
+  spmd_rule_map_.insert({std::move(kernel_name), rule});
+  return 0;
+}
+
+const SpmdRule& SpmdRuleFactory::GetSpmdRule(
+    const std::string& kernel_name) const {
+  auto it = spmd_rule_map_.find(kernel_name);
+  PADDLE_ENFORCE_NE(
+      it,
+      spmd_rule_map_.end(),
+      common::errors::NotFound("`%s` Kernel's Spmd rules is not registered.",
+                               kernel_name));
+  return it->second;
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h

@@ -0,0 +1,256 @@
+/* Copyright (c) 2023 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. */
+
+#pragma once
+
+#include <string>
+#include <typeindex>
+#include <typeinfo>
+#include <utility>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/attribute.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/type_defs.h"
+#include "paddle/utils/any.h"
+#include "paddle/utils/flat_hash_map.h"
+#include "paddle/utils/small_vector.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API InferSpmdContext {
+ public:
+  InferSpmdContext() = default;
+  InferSpmdContext(
+      paddle::small_vector<DistMetaTensor, kInputSmallVectorSize> inputs,
+      paddle::small_vector<Attribute, kAttrSmallVectorSize> attrs);
+
+  void EmplaceBackInput(DistMetaTensor input);
+  void EmplaceBackAttr(Attribute attr);
+  void EmplaceBackInputs(
+      paddle::small_vector<DistMetaTensor, kInputSmallVectorSize> inputs);
+
+  const DistMetaTensor& InputAt(size_t idx) const;
+
+  const std::pair<int, int>& InputRangeAt(size_t idx) const;
+  const std::vector<const DistMetaTensor*> InputsBetween(size_t start,
+                                                         size_t end) const;
+
+  template <typename AttrType>
+  AttrType AttrAt(size_t idx) const;
+
+  const Attribute& AttrAt(size_t idx) const;
+
+ private:
+  // Now we only need `inputs`, for backward, the `output` is passed as input
+  paddle::small_vector<DistMetaTensor, kInputSmallVectorSize> inputs_;
+  // Because the attribute arguments of dygraph do not have `attr name`,
+  // so we use vector instead of map
+  paddle::small_vector<Attribute, kAttrSmallVectorSize> attrs_;
+  // for vector arguments
+  paddle::small_vector<std::pair<int, int>, kInputSmallVectorSize> input_range_;
+};
+
+using InferSpmdFn = SpmdInfo (*)(const InferSpmdContext&);
+
+#define PD_INFER_SPMD(...)                                    \
+  ::phi::distributed::InferSpmdFnImpl<decltype(&__VA_ARGS__), \
+                                      &__VA_ARGS__>::Call
+
+template <typename T>
+struct InferSpmdTypeTag {};
+
+template <typename Fn, Fn fn>
+struct InferSpmdFnImpl;
+
+template <typename Return, typename... Args, Return (*infer_spmd_fn)(Args...)>
+struct InferSpmdFnImpl<Return (*)(Args...), infer_spmd_fn> {
+  static SpmdInfo Call(const InferSpmdContext& ctx) {
+    return InferSpmdFnCallHelper<Args..., InferSpmdTypeTag<int>>::
+        template Call<0, 0>(ctx);
+  }
+
+ private:
+  template <typename... RemainingArgs>
+  struct InferSpmdFnCallHelper;
+
+  // TODO(chenweihang): support other input type later as needed
+  template <typename... Tail>
+  struct InferSpmdFnCallHelper<const DistMetaTensor&, Tail...> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static SpmdInfo Call(const InferSpmdContext& ctx, PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0,
+                    "InferSpmd's Input should appear before Attributes.");
+      const std::pair<int, int> range = ctx.InputRangeAt(in_idx);
+      const DistMetaTensor& arg = ctx.InputAt(range.first);
+      return InferSpmdFnCallHelper<Tail...>::template Call<in_idx + 1,
+                                                           attr_idx>(
+          ctx, pargs..., arg);
+    }
+  };
+
+  // for vector slot
+  template <typename... Tail>
+  struct InferSpmdFnCallHelper<const std::vector<const DistMetaTensor*>&,
+                               Tail...> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static SpmdInfo Call(const InferSpmdContext& ctx, PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0,
+                    "InferSpmd's Input should appear before Attributes.");
+
+      const std::pair<int, int> range = ctx.InputRangeAt(in_idx);
+      std::vector<const DistMetaTensor*> arg =
+          ctx.InputsBetween(range.first, range.second);
+      return InferSpmdFnCallHelper<Tail...>::template Call<in_idx + 1,
+                                                           attr_idx>(
+          ctx, pargs..., arg);
+    }
+  };
+
+  // direct vector
+  template <typename... Tail>
+  struct InferSpmdFnCallHelper<const std::vector<DistMetaTensor>&, Tail...> {
+    template <int in_idx, int attr_idx, typename... PreviousArgs>
+    static SpmdInfo Call(const InferSpmdContext& ctx, PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0,
+                    "InferSpmd's Input should appear before Attributes.");
+      // TODO(liuzhenhai): parse input list as vector directly
+      const std::pair<int, int> range = ctx.InputRangeAt(in_idx);
+      std::vector<const DistMetaTensor*> tmp_arg =
+          ctx.InputsBetween(range.first, range.second);
+      std::vector<DistMetaTensor> arg;
+      std::transform(tmp_arg.begin(),
+                     tmp_arg.end(),
+                     std::back_inserter(arg),
+                     [](const DistMetaTensor* arg_ptr) { return *arg_ptr; });
+      return InferSpmdFnCallHelper<Tail...>::template Call<in_idx + 1,
+                                                           attr_idx>(
+          ctx, pargs..., arg);
+    }
+  };
+
+#define PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(attr_type)      \
+  template <typename... Tail>                                             \
+  struct InferSpmdFnCallHelper<attr_type, Tail...> {                      \
+    template <int in_idx, int attr_idx, typename... PreviousArgs>         \
+    static SpmdInfo Call(const InferSpmdContext& ctx,                     \
+                         PreviousArgs&... pargs) {                        \
+      attr_type arg = ctx.AttrAt<attr_type>(attr_idx);                    \
+      return InferSpmdFnCallHelper<Tail...>::template Call<in_idx,        \
+                                                           attr_idx + 1>( \
+          ctx, pargs..., arg);                                            \
+    }                                                                     \
+  }
+
+#define PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(attr_type) \
+  template <typename... Tail>                                                  \
+  struct InferSpmdFnCallHelper<const attr_type&, Tail...> {                    \
+    template <int in_idx, int attr_idx, typename... PreviousArgs>              \
+    static SpmdInfo Call(const InferSpmdContext& ctx,                          \
+                         PreviousArgs&... pargs) {                             \
+      attr_type arg = ctx.AttrAt<attr_type>(attr_idx);                         \
+      return InferSpmdFnCallHelper<Tail...>::template Call<in_idx,             \
+                                                           attr_idx + 1>(      \
+          ctx, pargs..., arg);                                                 \
+    }                                                                          \
+  }
+
+  // TODO(chenweihang): support other attr type later as needed
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(bool);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(int);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(float);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(double);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(int64_t);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_ATTRIBUTE(DataType);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(std::vector<int>);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(
+      std::vector<int64_t>);
+  PD_SPECIALIZE_InferSpmdFnCallHelper_FOR_CONST_ATTRIBUTE_REF(std::string);
+
+  /* End case */
+  template <typename T>
+  struct InferSpmdFnCallHelper<InferSpmdTypeTag<T>> {
+    template <int in_idx, int attr_idx>
+    static SpmdInfo Call(const InferSpmdContext& ctx UNUSED, Args&... args) {
+      return infer_spmd_fn(args...);
+    }
+  };
+};
+
+class SpmdRule {
+ public:
+  explicit SpmdRule(InferSpmdFn forward_fn)
+      : forward_fn_(forward_fn), backward_fn_(nullptr) {}
+
+  SpmdRule(InferSpmdFn forward_fn, InferSpmdFn backward_fn)
+      : forward_fn_(forward_fn), backward_fn_(backward_fn) {}
+
+  SpmdInfo InferForward(const InferSpmdContext& ctx) const {
+    PADDLE_ENFORCE_NE(forward_fn_,
+                      nullptr,
+                      common::errors::NotFound(
+                          "Current SpmdRule's forward function is not "
+                          "found, Please make sure "
+                          "that you have registered the rule correctly."));
+    return forward_fn_(ctx);
+  }
+
+  SpmdInfo InferBackward(const InferSpmdContext& ctx) const {
+    PADDLE_ENFORCE_NE(backward_fn_,
+                      nullptr,
+                      common::errors::NotFound(
+                          "Current SpmdRule's backward function is not "
+                          "found, Please make sure "
+                          "that you have registered the rule correctly."));
+    return backward_fn_(ctx);
+  }
+
+ private:
+  InferSpmdFn forward_fn_;
+  InferSpmdFn backward_fn_;
+};
+
+// SpmdRuleFactory manage the spmd rules and cache the propagate results
+// TODO(chenweihang): Add spmd caching impl later
+class PADDLE_API SpmdRuleFactory {
+ public:
+  static SpmdRuleFactory& Instance();
+
+  bool ContainsSpmdRule(const std::string& kernel_name) const;
+
+  int InsertSpmdRule(std::string kernel_name, SpmdRule rule);
+
+  const SpmdRule& GetSpmdRule(const std::string& kernel_name) const;
+
+ private:
+  SpmdRuleFactory() = default;
+
+  paddle::flat_hash_map<std::string, SpmdRule> spmd_rule_map_;
+
+  DISABLE_COPY_AND_ASSIGN(SpmdRuleFactory);
+};
+
+#define PD_REGISTER_SPMD_RULE(kernel_name, ...)                       \
+  UNUSED static int ___registrar_spmd_rule_for_##kernel_name =        \
+      ::phi::distributed::SpmdRuleFactory::Instance().InsertSpmdRule( \
+          #kernel_name, ::phi::distributed::SpmdRule(__VA_ARGS__));
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/placement_types.cc

@@ -0,0 +1,105 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/placement_types.h"
+
+namespace phi::distributed {
+
+int64_t DistTensorMeta::num_shard() const {
+  int64_t num_shard = 1;
+  const auto& mesh_shape = process_mesh_->shape();
+  for (size_t i = 0; i < placements_.size(); i++) {
+    if (placements_[i]->is_shard()) {
+      num_shard *= mesh_shape[i];
+    }
+  }
+  return num_shard;
+}
+
+std::vector<int64_t> DistTensorMeta::dim_mapping() const {
+  int64_t ndim = dims().size();
+  std::vector<int64_t> dim_map(ndim, -1);
+  for (size_t i = 0; i < placements_.size(); i++) {
+    auto& placement = placements_[i];
+    if (placement->is_shard()) {
+      auto shard_dim = dynamic_cast<const Shard&>(*placement).get_dim();
+      PADDLE_ENFORCE_EQ(
+          dim_map[shard_dim],
+          -1,
+          common::errors::InvalidArgument(
+              "Tensor dim %lld is already sharded on mesh dim %lld,"
+              " DistTensor operator implementation does not support things "
+              "like hybrid"
+              " sharding strategies yet (i.e. [Shard(0), Shard(0)])",
+              shard_dim,
+              dim_map[shard_dim]));
+      dim_map[shard_dim] = i;  // NOLINT
+    }
+  }
+  return dim_map;
+}
+
+bool DistTensorMeta::is_replicated() const {
+  return std::all_of(placements_.cbegin(),
+                     placements_.cend(),
+                     [](const auto& p) { return p->is_replicated(); });
+}
+
+bool equal_placements(const Placements& a, const Placements& b) {
+  if (a.size() != b.size()) {
+    return false;
+  }
+  for (size_t i = 0; i < a.size(); ++i) {
+    if (*a[i] != *b[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+phi::distributed::Placements cvt_dim_map_to_placements(
+    const ProcessMesh& process_mesh,
+    const std::vector<int64_t>& dim_mapping,
+    const paddle::flat_hash_map<int64_t, phi::ReduceType>& partial_status) {
+  phi::distributed::Placements placements;
+  placements.resize(process_mesh.ndim(),
+                    std::make_shared<phi::distributed::Replicate>());
+
+  for (const auto& pair : partial_status) {
+    placements[pair.first] =
+        std::make_shared<phi::distributed::Partial>(pair.second);
+  }
+
+  for (size_t i = 0; i < dim_mapping.size(); ++i) {
+    auto& mesh_id = dim_mapping[i];
+    if (mesh_id >= 0) {
+      auto& p = placements[mesh_id];
+      if (p->is_shard()) {
+        PADDLE_THROW(common::errors::PreconditionNotMet(
+            "ProcessMesh dimension can't be mapped to two  dimension of the "
+            "same tensor: {%d} and {%d}",
+            i,
+            dynamic_cast<phi::distributed::Shard&>(*p).get_dim()));
+      } else if (p->is_partial()) {
+        PADDLE_THROW(common::errors::PreconditionNotMet(
+            "ProcessMesh dimension {%d} cannot be both shard and partial!",
+            mesh_id));
+      }
+      placements[mesh_id] = std::make_shared<phi::distributed::Shard>(i);
+    }
+  }
+  return placements;
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/placement_types.h

@@ -0,0 +1,305 @@
+// Copyright (c) 2023 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.
+
+// This API design is inspired by:
+// https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tensor/placement_types.py
+// Git commit hash: 52e2b87d00ed527dc7f990d1a7a4c5498f99c513
+
+#pragma once
+
+#include <cstdint>
+#include <optional>
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/common/reduce_type.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/tensor_meta.h"
+#include "paddle/utils/flat_hash_map.h"
+
+namespace phi {
+namespace distributed {
+
+class Placement {
+ public:
+  virtual ~Placement() = default;
+
+  virtual bool is_shard(std::optional<int> dim = std::nullopt) const {
+    return false;
+  }
+
+  virtual bool is_replicated() const { return false; }
+
+  virtual bool is_partial() const { return false; }
+
+  virtual size_t hash() const { return 0; }
+
+  virtual std::string to_string() const { return ""; }
+
+  virtual bool operator==(const Placement& other) const {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Equal function is not implemented yet in Placement."));
+  }
+
+  virtual bool operator!=(const Placement& other) const {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Not Equal function is not implemented yet in Placement."));
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const Placement& p) {
+    os << p.to_string();
+    return os;
+  }
+};
+
+class Shard : public Placement {
+ public:
+  explicit Shard(int dim) : dim_(dim) {}
+
+  Shard(int dim, int split_factor) : dim_(dim), split_factor_(split_factor) {}
+
+  bool is_shard(std::optional<int> dim = std::nullopt) const override {
+    if (dim && *dim == this->dim_) {
+      return true;
+    } else {
+      return !dim.has_value();
+    }
+  }
+
+  bool operator==(const Placement& other) const override {
+    const Shard* other_shard = dynamic_cast<const Shard*>(&other);
+    if (!other_shard) return false;
+    if (other_shard->get_co_shard_order() != 0) return false;
+    return this->dim_ == other_shard->dim_ &&
+           this->split_factor_ == other_shard->split_factor_;
+  }
+
+  bool operator!=(const Placement& other) const override {
+    return !(*this == other);
+  }
+
+  std::size_t hash() const override {
+    return std::hash<std::string>{}(to_string());
+  }
+
+  int get_dim() const { return dim_; }
+
+  virtual int get_co_shard_order() const { return 0; }
+
+  void set_split_factor(int64_t sf) { split_factor_ = sf; }
+
+  int get_split_factor() const { return split_factor_; }
+
+  friend std::ostream& operator<<(std::ostream& os, const Shard& p) {
+    os << p.to_string();
+    return os;
+  }
+
+  std::string to_string() const override {
+    std::stringstream ss;
+    ss << "Shard(dim=" << std::to_string(dim_);
+    if (split_factor_ != 1) {
+      ss << ", split_factor=" << std::to_string(split_factor_);
+    }
+    ss << ")";
+
+    return ss.str();
+  }
+
+  virtual std::shared_ptr<Shard> copy() const {
+    return std::make_shared<Shard>(*this);
+  }
+
+  virtual std::shared_ptr<Shard> deepcopy() const {
+    return std::make_shared<Shard>(*this);
+  }
+
+ protected:
+  int dim_;
+  int split_factor_ = 1;
+};
+
+class CoShard : public Shard {
+ public:
+  CoShard(int64_t dim, int64_t co_shard_order)
+      : Shard(dim, 1), co_shard_order_(co_shard_order) {}
+
+  int get_co_shard_order() const override { return co_shard_order_; }
+
+  std::string to_string() const override {
+    std::stringstream ss;
+    ss << "Shard(dim=" << std::to_string(dim_);
+    ss << ", shard_order=" << std::to_string(co_shard_order_) << ")";
+
+    return ss.str();
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const CoShard& p) {
+    os << p.to_string();
+    return os;
+  }
+
+  std::shared_ptr<Shard> copy() const override {
+    return std::make_shared<CoShard>(*this);
+  }
+
+  std::shared_ptr<Shard> deepcopy() const override {
+    return std::make_shared<CoShard>(*this);
+  }
+
+  bool operator==(const Placement& other) const override {
+    if (const CoShard* other_coshard = dynamic_cast<const CoShard*>(&other)) {
+      return this->dim_ == other_coshard->dim_ &&
+             this->split_factor_ == other_coshard->split_factor_ &&
+             this->co_shard_order_ == other_coshard->co_shard_order_;
+    }
+    if (const Shard* other_shard = dynamic_cast<const Shard*>(&other)) {
+      return this->co_shard_order_ == 0 &&
+             this->dim_ == other_shard->get_dim() &&
+             this->split_factor_ == other_shard->get_split_factor();
+    }
+    return false;
+  }
+
+  bool operator!=(const Placement& other) const override {
+    return !(*this == other);
+  }
+
+  std::size_t hash() const override {
+    std::stringstream ss;
+    ss << "Shard(dim=" << std::to_string(dim_);
+    if (split_factor_ != 1) {
+      ss << ", split_factor=" << std::to_string(split_factor_);
+    }
+    if (co_shard_order_ != 0) {
+      ss << ", shard_order=" << std::to_string(co_shard_order_);
+    }
+    ss << ")";
+    return std::hash<std::string>{}(ss.str());
+  }
+
+ private:
+  int64_t co_shard_order_ = 0;
+};
+
+class Replicate : public Placement {
+ public:
+  bool is_replicated() const override { return true; }
+
+  bool operator==(const Placement& other) const override {
+    return dynamic_cast<const Replicate*>(&other) != nullptr;
+  }
+
+  bool operator!=(const Placement& other) const override {
+    return !(*this == other);
+  }
+
+  std::size_t hash() const override {
+    return std::hash<std::string>{}(to_string());
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const Replicate& p) {
+    os << p.to_string();
+    return os;
+  }
+
+  std::string to_string() const override { return "Replicate()"; }
+};
+
+class Partial : public Placement {
+ public:
+  explicit Partial(ReduceType reduce_type) : reduce_type_(reduce_type) {}
+  bool is_partial() const override { return true; }
+  ReduceType get_reduce_type() const { return reduce_type_; }
+
+  bool operator==(const Placement& other) const override {
+    const Partial* other_partial = dynamic_cast<const Partial*>(&other);
+    return other_partial && this->reduce_type_ == other_partial->reduce_type_;
+  }
+
+  bool operator!=(const Placement& other) const override {
+    return !(*this == other);
+  }
+
+  std::size_t hash() const override {
+    return std::hash<std::string>{}(to_string());
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const Partial& p) {
+    os << p.to_string();
+    return os;
+  }
+
+  std::string to_string() const override {
+    return "Partial(reduce_type=" +
+           std::string(ReduceTypeStrings[static_cast<int>(reduce_type_)]) + ")";
+  }
+
+ private:
+  ReduceType reduce_type_;
+};
+
+using Placements = std::vector<std::shared_ptr<Placement>>;
+class DistTensorMeta : public std::enable_shared_from_this<DistTensorMeta> {
+ public:
+  DistTensorMeta(const ProcessMesh& process_mesh,
+                 const Placements& placements,
+                 const DenseTensorMeta& tensor_meta)
+      : process_mesh_(std::make_shared<const ProcessMesh>(process_mesh)),
+        placements_(placements),
+        tensor_meta_(std::make_shared<const DenseTensorMeta>(tensor_meta)) {}
+
+  DistTensorMeta() = default;
+
+  const DDim& dims() const { return tensor_meta_->dims; }
+
+  const ProcessMesh& process_mesh() const { return *process_mesh_; }
+
+  const Placements& placements() const { return placements_; }
+
+  int64_t num_shard() const;
+
+  std::vector<int64_t> dim_mapping() const;
+
+  bool is_replicated() const;
+
+ private:
+  std::shared_ptr<const ProcessMesh> process_mesh_;
+  Placements placements_;
+  std::shared_ptr<const DenseTensorMeta> tensor_meta_;
+};
+
+PADDLE_API bool equal_placements(const Placements& a, const Placements& b);
+
+PADDLE_API phi::distributed::Placements cvt_dim_map_to_placements(
+    const ProcessMesh& process_mesh,
+    const std::vector<int64_t>& dim_mapping,
+    const paddle::flat_hash_map<int64_t, phi::ReduceType>& partial_status);
+
+}  // namespace distributed
+}  // namespace phi
+
+namespace std {
+template <>
+struct hash<phi::distributed::Placement> {
+  std::size_t operator()(const phi::distributed::Placement& p) const {
+    return p.hash();
+  }
+};
+
+}  // namespace std

paddle/phi/core/distributed/auto_parallel/process_mesh.cc

@@ -0,0 +1,241 @@
+/* Copyright (c) 2022 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. */
+
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+
+#include <algorithm>
+#include <iterator>
+#include <numeric>
+#include <set>
+
+#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+
+namespace phi::distributed {
+
+using phi::distributed::auto_parallel::has_duplicates;
+using phi::distributed::auto_parallel::ProcessMeshProto;
+using phi::distributed::auto_parallel::str_join;
+
+ProcessMesh::ProcessMesh(const std::vector<int64_t> &shape,
+                         const std::vector<int64_t> &process_ids,
+                         const std::vector<std::string> &dim_names) {
+  shape_ = shape;
+  int64_t size = this->size();
+  PADDLE_ENFORCE_EQ(
+      size,
+      process_ids.size(),
+      errors::InvalidArgument(
+          "The size of this process mesh must be equal to the size of its "
+          "process ids, but got mesh size %d and process id size %d.",
+          size,
+          process_ids.size()));
+  PADDLE_ENFORCE_EQ(
+      has_duplicates(process_ids),
+      false,
+      errors::InvalidArgument("The process ids [%s] must be unique.",
+                              str_join(process_ids_)));
+  process_ids_ = process_ids;
+
+  PADDLE_ENFORCE_EQ(
+      shape_.size(),
+      dim_names.size(),
+      errors::InvalidArgument(
+          "The size of mesh shape must be equal to the size of the dimension "
+          "names, but got shape size %d and dimension name size %d.",
+          shape_.size(),
+          dim_names.size()));
+  PADDLE_ENFORCE_EQ(has_duplicates(dim_names),
+                    false,
+                    errors::InvalidArgument(
+                        "The names [%s] of each dimension must be unique.",
+                        str_join(dim_names)));
+  dim_names_ = dim_names;
+}
+
+int64_t ProcessMesh::size() const {
+  if (shape_.empty()) return 0;
+  int64_t size = 1;
+  for (const int64_t dim_size : shape_) size *= dim_size;
+  return size;
+}
+
+bool ProcessMesh::contains(int64_t process_id) const {
+  auto result =
+      std::find(std::begin(process_ids_), std::end(process_ids_), process_id);
+  if (result != std::end(process_ids_)) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+std::string ProcessMesh::to_string() const {
+  std::string mesh_str = "{shape: [" + str_join(shape_) + "], ";
+  mesh_str += "process_ids: [" + str_join(process_ids_) + "], ";
+  mesh_str += "dim_names: [" + str_join(dim_names_) + "]}";
+  return mesh_str;
+}
+
+ProcessMesh ProcessMesh::from_proto(const ProcessMeshProto &proto) {
+  ProcessMesh mesh;
+
+  mesh.shape_.resize(proto.shape_size());
+  for (int i = 0; i < proto.shape_size(); ++i) {
+    mesh.shape_[i] = proto.shape(i);
+  }
+
+  mesh.process_ids_.resize(proto.process_ids_size());
+  for (int i = 0; i < proto.process_ids_size(); ++i) {
+    mesh.process_ids_[i] = proto.process_ids(i);
+  }
+
+  mesh.dim_names_.resize(proto.dim_names_size());
+  for (int i = 0; i < proto.dim_names_size(); ++i) {
+    mesh.dim_names_[i] = proto.dim_names(i);
+  }
+
+  return mesh;
+}
+
+void ProcessMesh::to_proto(ProcessMeshProto *proto) const {
+  for (const auto &i : shape_) {
+    proto->add_shape(i);
+  }
+
+  for (const auto &i : process_ids_) {
+    proto->add_process_ids(i);
+  }
+
+  for (const auto &i : dim_names_) {
+    proto->add_dim_names(i);
+  }
+}
+
+bool operator==(const ProcessMesh &lhs, const ProcessMesh &rhs) {
+  if (lhs.shape() != rhs.shape() || lhs.process_ids() != rhs.process_ids()) {
+    return false;
+  }
+  return true;
+}
+
+bool mesh_equal_ignore_shape1(const ProcessMesh &a,
+                              const ProcessMesh &b,
+                              int split_dim) {
+  if (a == b) {
+    return true;
+  }
+  if (a.process_ids() != b.process_ids()) {
+    return false;
+  }
+  std::vector<int64_t> a_shape = a.shape();
+  std::vector<int64_t> b_shape = b.shape();
+  if (a_shape[split_dim] != 1) {
+    return false;
+  }
+  a_shape.erase(a_shape.begin() + split_dim);
+  return a_shape == b_shape;
+}
+
+std::vector<ProcessMesh> SplitMesh(const ProcessMesh &mesh, int axis) {
+  std::vector<int64_t> mesh_shape = mesh.shape();
+  std::vector<int64_t> process_ids = mesh.process_ids();
+  std::vector<ProcessMesh> result;
+
+  int64_t total_elements = process_ids.size();
+
+  int64_t num_splits = mesh_shape[axis];
+  int64_t prod_before = std::accumulate(mesh_shape.begin(),
+                                        mesh_shape.begin() + axis,
+                                        1,
+                                        std::multiplies<int64_t>());
+  int64_t prod_after = std::accumulate(mesh_shape.begin() + axis + 1,
+                                       mesh_shape.end(),
+                                       1,
+                                       std::multiplies<int64_t>());
+
+  for (int i = 0; i < num_splits; ++i) {
+    std::vector<int64_t> new_shape = mesh_shape;
+    new_shape[axis] = 1;
+
+    std::vector<int64_t> new_process_ids;
+    for (int64_t j = 0; j < prod_before; ++j) {
+      for (int64_t k = 0; k < prod_after; ++k) {
+        int64_t index = j * mesh_shape[axis] * prod_after + i * prod_after + k;
+        if (index < total_elements) {
+          new_process_ids.push_back(process_ids[index]);
+        }
+      }
+    }
+
+    result.emplace_back(
+        ProcessMesh(new_shape, new_process_ids, mesh.dim_names()));
+  }
+
+  return result;
+}
+
+int SubMeshDim(const ProcessMesh &global_mesh, const ProcessMesh &sub_mesh) {
+  std::set<int64_t> global_ids(global_mesh.process_ids().begin(),
+                               global_mesh.process_ids().end());
+  std::set<int64_t> sub_ids(sub_mesh.process_ids().begin(),
+                            sub_mesh.process_ids().end());
+  if (!std::includes(
+          global_ids.begin(), global_ids.end(), sub_ids.begin(), sub_ids.end()))
+    return -1;
+
+  int sub_dim = -1;
+  std::vector<int64_t> global_shape = global_mesh.shape();
+  std::vector<int64_t> sub_shape = sub_mesh.shape();
+
+  if (global_mesh.ndim() == sub_mesh.ndim() + 1) {
+    // for the case that the `1` is not explicitly specified in the shape
+    // e.g.
+    //  global_mesh: shape = [2,3], process_ids = [0,1,2,3,4,5]
+    //  sub_mesh: shape = [3], process_ids = [0,1,2]
+    int global_ndim = global_mesh.ndim();
+    for (int i = 0; i < global_ndim - 1; ++i) {
+      std::vector<ProcessMesh> sub_meshes = SplitMesh(global_mesh, i);
+      for (const ProcessMesh &mesh : sub_meshes) {
+        if (mesh_equal_ignore_shape1(mesh, sub_mesh, i)) {
+          return i;
+        }
+      }
+    }
+    return sub_dim;
+  } else if (global_mesh.ndim() != sub_mesh.ndim()) {
+    return -1;
+  }
+
+  // e.g.
+  //  global_mesh: shape = [1,2], process_ids = [0,1]; sub_mesh: shape = [1, 1],
+  //  process_ids = [0] global_mesh: shape = [2,2], process_ids = [0,1,2,3];
+  //  sub_mesh: shape = [2, 1], process_ids = [0, 2]
+  auto it =
+      std::mismatch(sub_shape.begin(), sub_shape.end(), global_shape.begin());
+  if (it.first == sub_shape.end()) {
+    return -1;
+  }
+
+  sub_dim = it.first - sub_shape.begin();
+  std::vector<ProcessMesh> sub_meshes = SplitMesh(global_mesh, sub_dim);
+  if (std::find(sub_meshes.begin(), sub_meshes.end(), sub_mesh) !=
+      sub_meshes.end()) {
+    return sub_dim;
+  }
+
+  return -1;
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/process_mesh.h

@@ -0,0 +1,114 @@
+/* Copyright (c) 2022 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. */
+
+#pragma once
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace distributed {
+
+namespace auto_parallel {
+class ProcessMeshProto;
+}
+
+class PADDLE_API ProcessMesh {
+ public:
+  ProcessMesh() = default;
+
+  ProcessMesh(const std::vector<int64_t>& shape,
+              const std::vector<int64_t>& process_ids,
+              const std::vector<std::string>& dim_names);
+
+  const std::vector<int64_t>& shape() const { return shape_; }
+
+  const std::vector<int64_t>& process_ids() const { return process_ids_; }
+
+  const std::vector<std::string>& dim_names() const { return dim_names_; }
+
+  int64_t size() const;
+
+  int64_t ndim() const { return shape_.size(); }
+
+  int64_t dim_size(int64_t dim) const {
+    int64_t cdim = auto_parallel::canonical_dim(dim, shape_.size());
+    return shape_[cdim];
+  }
+
+  int64_t dim_size(const std::string& dim_name) const {
+    for (std::size_t i = 0; i < dim_names_.size(); ++i) {
+      if (dim_names_[i] == dim_name) {
+        return shape_[i];
+      }
+    }
+    PADDLE_THROW(errors::InvalidArgument(
+        "Cannot find the dimension of %s in this process mesh.", dim_name));
+  }
+
+  bool empty() const { return (shape_.empty() || process_ids_.empty()); }
+  bool contains(int64_t process_id) const;
+
+  size_t hash() const { return std::hash<std::string>{}(to_string()); }
+
+  // ProcessMesh from_string(const std::string& mesh_str);
+  std::string to_string() const;
+
+  static ProcessMesh from_proto(const auto_parallel::ProcessMeshProto& proto);
+  void to_proto(auto_parallel::ProcessMeshProto* proto) const;
+
+ private:
+  std::vector<int64_t> shape_;
+  std::vector<int64_t> process_ids_;
+  std::vector<std::string> dim_names_;
+};
+
+inline std::ostream& operator<<(std::ostream& os, const ProcessMesh& obj) {
+  os << obj.to_string();
+  return os;
+}
+
+PADDLE_API bool operator==(const ProcessMesh& lhs, const ProcessMesh& rhs);
+
+inline bool operator!=(const ProcessMesh& lhs, const ProcessMesh& rhs) {
+  return !operator==(lhs, rhs);
+}
+
+// split the mesh into sub-meshes at the given axis
+PADDLE_API std::vector<ProcessMesh> SplitMesh(const ProcessMesh& mesh,
+                                              int axis);
+
+// return which dimension that the sub_mesh is split from the global_mesh,
+// if sub_mesh is not a subset of global_mesh, return -1
+PADDLE_API int SubMeshDim(const ProcessMesh& global_mesh,
+                          const ProcessMesh& sub_mesh);
+
+// when the shapes of two meshes are different and their process_ids
+// are the same, check whether the only difference is that mesh 'a'
+// has an additional '1' on the split dim of its shape.
+// e.g. a.shape = [2], b.shape = [2, 1], and the process_ids are the
+// same, then they are equal.
+PADDLE_API bool mesh_equal_ignore_shape1(const ProcessMesh& a,
+                                         const ProcessMesh& b,
+                                         int split_dim);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/proto_helper.cc

@@ -0,0 +1,63 @@
+// Copyright (c) 2024 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/proto_helper.h"
+#include "paddle/phi/core/distributed/auto_parallel/auto_parallel.pb.h"
+#include "paddle/phi/core/distributed/auto_parallel/device_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_mapper.h"
+
+#define TO_PROTO_HELPER(object, proto_type) \
+  proto_type proto;                         \
+  object.to_proto(&proto);                  \
+  return proto
+
+namespace phi::distributed {
+
+auto_parallel::TensorDistAttrProto to_proto(const TensorDistAttr& dist_attr) {
+  TO_PROTO_HELPER(dist_attr, auto_parallel::TensorDistAttrProto);
+}
+
+auto_parallel::ProcessMeshProto to_proto(const ProcessMesh& process_mesh) {
+  TO_PROTO_HELPER(process_mesh, auto_parallel::ProcessMeshProto);
+}
+
+auto_parallel::DeviceCapabilityProto to_proto(
+    const auto_parallel::DeviceCapability& device_capability) {
+  TO_PROTO_HELPER(device_capability, auto_parallel::DeviceCapabilityProto);
+}
+
+auto_parallel::DeviceProto to_proto(const auto_parallel::Device& device) {
+  TO_PROTO_HELPER(device, auto_parallel::DeviceProto);
+}
+
+auto_parallel::LinkCapabilityProto to_proto(
+    const auto_parallel::LinkCapability& link_capability) {
+  TO_PROTO_HELPER(link_capability, auto_parallel::LinkCapabilityProto);
+}
+
+auto_parallel::LinkProto to_proto(const auto_parallel::Link& link) {
+  TO_PROTO_HELPER(link, auto_parallel::LinkProto);
+}
+
+auto_parallel::DeviceMeshProto to_proto(
+    const auto_parallel::DeviceMesh& device_mesh) {
+  TO_PROTO_HELPER(device_mesh, auto_parallel::DeviceMeshProto);
+}
+
+auto_parallel::DistributedMapperProto to_proto(
+    const auto_parallel::DistributedMapper& dist_mapper) {
+  TO_PROTO_HELPER(dist_mapper, auto_parallel::DistributedMapperProto);
+}
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/proto_helper.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2024 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.
+
+#pragma once
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/distributed/auto_parallel/auto_parallel.pb.h"
+namespace phi {
+namespace distributed {
+class TensorDistAttr;
+class ProcessMesh;
+namespace auto_parallel {
+struct DeviceCapability;
+class Device;
+struct LinkCapability;
+class Link;
+class DeviceMesh;
+class DistributedMapper;
+}  // namespace auto_parallel
+PADDLE_API auto_parallel::TensorDistAttrProto to_proto(
+    const TensorDistAttr& dist_attr);
+PADDLE_API auto_parallel::ProcessMeshProto to_proto(
+    const ProcessMesh& dist_attr);
+
+PADDLE_API auto_parallel::DeviceCapabilityProto to_proto(
+    const auto_parallel::DeviceCapability& device_capability);
+PADDLE_API auto_parallel::DeviceProto to_proto(
+    const auto_parallel::Device& device);
+PADDLE_API auto_parallel::LinkCapabilityProto to_proto(
+    const auto_parallel::LinkCapability& link_capability);
+PADDLE_API auto_parallel::LinkProto to_proto(const auto_parallel::Link& link);
+PADDLE_API auto_parallel::DeviceMeshProto to_proto(
+    const auto_parallel::DeviceMesh& link);
+PADDLE_API auto_parallel::DistributedMapperProto to_proto(
+    const auto_parallel::DistributedMapper& dist_mapper);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/CMakeLists.txt

@@ -0,0 +1,18 @@
+collect_srcs(
+  core_srcs
+  SRCS
+  reshard_utils.cc
+  reshard_function.cc
+  r_to_s_reshard_function.cc
+  s_to_r_reshard_function.cc
+  r_to_p_reshard_function.cc
+  p_to_r_reshard_function.cc
+  s_to_s_reshard_function.cc
+  p_to_s_reshard_function.cc
+  s_to_p_reshard_function.cc
+  x_to_r_reshard_function.cc
+  r_to_x_reshard_function.cc
+  nd_mesh_reshard_function.cc
+  same_status_reshard_function.cc
+  global_and_sub_mesh_reshard_function.cc
+  reshard_function_registry.cc)

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.cc

@@ -0,0 +1,150 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/p_recv_kernel.h"
+#include "paddle/phi/kernels/p_send_kernel.h"
+
+namespace phi::distributed {
+
+bool GlobalToSubMeshReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const TensorDistAttr& in_dist_attr = in.dist_attr();
+
+  const ProcessMesh& in_process_mesh = in_dist_attr.process_mesh();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+
+  int sub_mesh_dim = SubMeshDim(in_process_mesh, out_process_mesh);
+  RESHARD_SHORTCUT_IF_FALSE(sub_mesh_dim != -1);
+  // 1. the split dimension must be replicated
+  // 2. out mesh is the value of a certain dimension of global mesh
+  // e.g. global_mesh = [[1, 2], [3, 4]], out_mesh = [1, 2] or [3, 4]
+  //      global_mesh = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
+  //      out_mesh = [[1, 2], [3, 4]] or [[5, 6], [7, 8]]
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated(sub_mesh_dim));
+
+  return true;
+}
+
+void GlobalToSubMeshReshardFunction::Eval(DeviceContext* dev_ctx,
+                                          const DistTensor& in,
+                                          const TensorDistAttr& out_dist_attr,
+                                          DistTensor* out) {
+  VLOG(3) << "Call GlobalToSubMeshReshardFunction Eval";
+  const DenseTensor& in_dense_value = in.value();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+  if (IsCurRankInMesh(out_process_mesh)) {
+    SetValue(out, in_dense_value);
+  } else {
+    *(out->unsafe_mutable_value()) =
+        DenseTensor(std::make_shared<phi::Allocation>(
+                        nullptr, 0, phi::distributed::GetDefaultPlace()),
+                    in.value().meta());
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool SubMeshToGlobalReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const TensorDistAttr& in_dist_attr = in.dist_attr();
+  const ProcessMesh& in_process_mesh = in_dist_attr.process_mesh();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+
+  int sub_mesh_dim = SubMeshDim(out_process_mesh, in_process_mesh);
+  RESHARD_SHORTCUT_IF_FALSE(sub_mesh_dim != -1);
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated(sub_mesh_dim));
+
+  return true;
+}
+
+void SubMeshToGlobalReshardFunction::Eval(DeviceContext* dev_ctx,
+                                          const DistTensor& in,
+                                          const TensorDistAttr& out_dist_attr,
+                                          DistTensor* out) {
+  VLOG(3) << "Call SubMeshToGlobalReshardFunction Eval";
+  const TensorDistAttr& in_dist_attr = in.dist_attr();
+  const ProcessMesh& in_process_mesh = in_dist_attr.process_mesh();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+
+  int sub_mesh_dim = SubMeshDim(out_process_mesh, in_process_mesh);
+  std::vector<ProcessMesh> sub_process_meshes =
+      SplitMesh(out_process_mesh, sub_mesh_dim);
+  const std::vector<int64_t>& in_process_ids = in_process_mesh.process_ids();
+  const std::vector<int64_t>& out_process_ids = out_process_mesh.process_ids();
+  std::unordered_map<int64_t, std::vector<int64_t>> send2recv_map;
+  std::unordered_map<int64_t, int64_t> recv2send_map;
+
+  for (const ProcessMesh& sub_mesh : sub_process_meshes) {
+    if (mesh_equal_ignore_shape1(sub_mesh, in_process_mesh, sub_mesh_dim)) {
+      continue;
+    }
+    const std::vector<int64_t>& sub_process_ids = sub_mesh.process_ids();
+    for (size_t i = 0; i < sub_process_ids.size(); ++i) {
+      int64_t send_id = in_process_ids[i];
+      send2recv_map[send_id].push_back(sub_process_ids[i]);
+      recv2send_map[sub_process_ids[i]] = send_id;
+    }
+  }
+
+  std::vector<int64_t> all_process_ids =
+      GetUnionProcessIds(in_process_ids, out_process_ids);
+  int64_t cur_global_rank = GetCurGlobalRank();
+  DataType dtype = in.dtype();
+  if (IsCurRankInMesh(in_process_mesh)) {
+    const DenseTensor& in_dense_value = in.value();
+    std::vector<int64_t>& recv_vec = send2recv_map[cur_global_rank];
+    for (int64_t recv_id : recv_vec) {
+      auto relative_recv_rank = recv_id;
+      for (size_t i = 0; i < all_process_ids.size(); ++i) {
+        if (all_process_ids[i] == recv_id) {
+          relative_recv_rank = i;
+        }
+      }
+      RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                PSendKernel,
+                                dtype,
+                                all_process_ids,
+                                in_dense_value,
+                                relative_recv_rank, /*peer*/
+                                true /*dynamic_shape*/);
+    }
+    SetValue(out, in_dense_value);
+  } else {
+    int64_t send_id = recv2send_map[cur_global_rank];
+    auto relative_send_rank = send_id;
+    for (size_t i = 0; i < all_process_ids.size(); ++i) {
+      if (all_process_ids[i] == send_id) {
+        relative_send_rank = i;
+      }
+    }
+
+    RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                              PRecv,
+                              dtype,
+                              all_process_ids,
+                              relative_send_rank, /*peer*/
+                              {} /*out_shape*/,
+                              true /*dynamic_shape*/,
+                              GetMutableTensor(out));
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.h

@@ -0,0 +1,49 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API GlobalToSubMeshReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "GlobalToSubMeshReshardFunction"; }
+};
+
+class PADDLE_API SubMeshToGlobalReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SubMeshToGlobalReshardFunction"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.cc

@@ -0,0 +1,550 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+
+namespace phi::distributed {
+
+namespace {
+ProcessMesh GetSubProcessMesh(const ProcessMesh& mesh, int64_t axis) {
+  int64_t shape_of_axis = mesh.dim_size(axis);
+  std::vector<int64_t> shape = {shape_of_axis};
+  std::vector<std::string> dim_names = {mesh.dim_names()[axis]};
+  std::vector<int64_t> coord = GetCurRankCoordInMesh(mesh);
+
+  std::vector<int64_t> process_ids;
+  for (int64_t i = 0; i < shape_of_axis; ++i) {
+    coord[axis] = i;
+    int64_t rank = 0;
+    int64_t degree = 1;
+    for (int64_t j = static_cast<int64_t>(coord.size() - 1); j >= 0; --j) {
+      rank += coord[j] * degree;
+      degree *= mesh.dim_size(j);
+    }
+    process_ids.emplace_back(mesh.process_ids()[rank]);
+  }
+
+  ProcessMesh out_mesh(shape, process_ids, dim_names);
+  return out_mesh;
+}
+
+// Given the input two dist_attr, traversing from high-dimension axis to
+// low-dimension. Find and return the first different axis which is shard status
+// between these two. For example, the input two dims_mapping are [-1, 0, -1,
+// -1] and [-1, -1, 0, -1], the first diff shard axis is 2.
+int64_t FindFirstDiffShardAxis(const TensorDistAttr& in_dist_attr,
+                               const TensorDistAttr& out_dist_attr) {
+  const auto& in_dims_mapping = in_dist_attr.multi_dims_mapping();
+  const auto& out_dims_mapping = out_dist_attr.multi_dims_mapping();
+
+  VLOG(3) << "In find diff axis, in dim mapping "
+          << auto_parallel::str_join(in_dims_mapping) << ", out dim mapping "
+          << auto_parallel::str_join(out_dims_mapping);
+  int64_t axis = -1;
+
+  for (int64_t i = static_cast<int64_t>(in_dims_mapping.size() - 1); i >= 0;
+       --i) {
+    if (in_dims_mapping[i] != out_dims_mapping[i]) {
+      axis = i;
+      break;
+    }
+    auto predicate = [&in_dist_attr, &out_dist_attr](int64_t mesh_dim) {
+      if (in_dist_attr.get_split_factor(mesh_dim) !=
+          out_dist_attr.get_split_factor(mesh_dim)) {
+        return true;
+      }
+      return false;
+    };
+    if (std::any_of(
+            in_dims_mapping[i].begin(), in_dims_mapping[i].end(), predicate)) {
+      axis = i;
+      break;
+    }
+  }
+
+  return axis;
+}
+}  // namespace
+
+class ReshardContext final {
+ public:
+  ReshardContext(DeviceContext* dev_ctx,
+                 const DistTensor& in,
+                 const TensorDistAttr& out_dist_attr,
+                 DistTensor* out)
+      : dev_ctx(dev_ctx), in(in), out_dist_attr(out_dist_attr), out(out) {}
+
+  TensorDistAttr CreateOneDimDistAttr(
+      const ProcessMesh& sub_mesh,
+      bool is_partial = false,
+      std::optional<ReduceType> reduce_type = std::nullopt,
+      std::optional<int64_t> cur_tensor_dim = std::nullopt,
+      std::optional<int64_t> cur_mesh_split_factor = std::nullopt) const {
+    TensorDistAttr dist_attr(common::vectorize(in.dims()));
+    dist_attr.set_process_mesh(sub_mesh);
+
+    if (is_partial) {
+      if (reduce_type) {
+        dist_attr.set_partial_status(std::vector<int64_t>{0},
+                                     reduce_type.value());
+      } else {
+        dist_attr.set_partial_status(std::vector<int64_t>{0});
+      }
+    }
+
+    if (cur_tensor_dim) {
+      auto dims_mapping = dist_attr.multi_dims_mapping();
+      PADDLE_ENFORCE_GE(
+          cur_tensor_dim.value(),
+          0,
+          common::errors::InvalidArgument(
+              "tensor dim should be greater than or equal to 0, but got %d.",
+              cur_tensor_dim.value()));
+      dims_mapping[cur_tensor_dim.value()] = {0};
+      dist_attr.set_dims_mapping(dims_mapping);
+      if (cur_mesh_split_factor) {
+        dist_attr.set_split_factor(0, cur_mesh_split_factor.value());
+      }
+    }
+    return dist_attr;
+  }
+
+  ProcessMesh GetSubProcessMesh(int64_t axis) const {
+    return phi::distributed::GetSubProcessMesh(out_dist_attr.process_mesh(),
+                                               axis);
+  }
+
+  DeviceContext* dev_ctx;
+  const DistTensor& in;
+  const TensorDistAttr& out_dist_attr;
+  DistTensor* out;
+  DistTensor tmp_result;
+};
+
+template <typename ReshardFunc>
+class SingleDimReshardStrategy
+    : public SameNdMeshReshardFunction::ReshardStrategy {
+ public:
+  SingleDimReshardStrategy(int64_t cur_tensor_dim,
+                           int64_t cur_mesh_dim,
+                           ReshardContext ctx)
+      : cur_tensor_dim_(cur_tensor_dim),
+        cur_mesh_dim_(cur_mesh_dim),
+        ctx_(ctx) {}
+  void Eval() override {
+    auto cur_dist_attr = CalculateNewDistAttr();
+    VLOG(3) << "New Dist Attr " << cur_dist_attr;
+    auto sub_mesh = ctx_.GetSubProcessMesh(cur_mesh_dim_);
+    VLOG(3) << "Get Sub Mesh " << sub_mesh;
+    auto in_one_dim = CreateOneDimInDistAttr(sub_mesh);
+    VLOG(3) << "One dim In Attr " << in_one_dim;
+    auto out_one_dim = CreateOneDimOutDistAttr(sub_mesh);
+    VLOG(3) << "One dim Out Attr " << out_one_dim;
+
+    SetDistProps(ctx_.out, in_one_dim);
+    VLOG(3) << "Set One dim In Attr";
+    ReshardFunc func;
+    func.Eval(ctx_.dev_ctx, *ctx_.out, out_one_dim, &ctx_.tmp_result);
+    VLOG(3) << "Finish reshard func.";
+    SetValue(ctx_.out, ctx_.tmp_result.value());
+    VLOG(3) << "Set local value";
+    SetDistProps(ctx_.out, cur_dist_attr);
+    VLOG(3) << "Set Cur Dist Attr";
+  }
+
+ private:
+  virtual TensorDistAttr CalculateNewDistAttr() const = 0;
+  virtual TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const = 0;
+  virtual TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const = 0;
+
+ protected:
+  int64_t cur_tensor_dim_;
+  int64_t cur_mesh_dim_;
+  ReshardContext ctx_;
+};
+
+class PartialToReplicate final
+    : public SingleDimReshardStrategy<PToRReshardFunction> {
+ private:
+  using SingleDimReshardStrategy<PToRReshardFunction>::SingleDimReshardStrategy;
+
+  TensorDistAttr CalculateNewDistAttr() const override {
+    auto real_out_attr = ctx_.out->dist_attr();
+    real_out_attr.clean_partial_dims({cur_mesh_dim_});
+    return real_out_attr;
+  }
+
+  TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    auto input_reduce_type =
+        ctx_.out->dist_attr().partial_status().at(cur_mesh_dim_);
+    return ctx_.CreateOneDimDistAttr(sub_mesh, true, input_reduce_type);
+  }
+
+  TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    return ctx_.CreateOneDimDistAttr(sub_mesh);
+  }
+};
+
+class ShardToReplicate final
+    : public SingleDimReshardStrategy<SToRReshardFunction> {
+ private:
+  using SingleDimReshardStrategy<SToRReshardFunction>::SingleDimReshardStrategy;
+
+  TensorDistAttr CalculateNewDistAttr() const override {
+    auto real_out_attr = ctx_.out->dist_attr();
+    std::vector<std::vector<int64_t>> real_dims_mapping =
+        real_out_attr.multi_dims_mapping();
+    real_dims_mapping[cur_tensor_dim_] = {};
+    real_out_attr.set_dims_mapping(real_dims_mapping);
+    real_out_attr.clear_split_factor(cur_mesh_dim_);
+    return real_out_attr;
+  }
+
+  TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    auto split_factor = ctx_.out->dist_attr().get_split_factor(cur_mesh_dim_);
+    VLOG(3) << "In S To R, cur mesh dim is " << cur_mesh_dim_
+            << ", split factor is " << split_factor;
+    return ctx_.CreateOneDimDistAttr(
+        sub_mesh, false, std::nullopt, cur_tensor_dim_, split_factor);
+  }
+
+  TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    return ctx_.CreateOneDimDistAttr(sub_mesh);
+  }
+};
+
+class ReplicateToPartial final
+    : public SingleDimReshardStrategy<RToPReshardFunction> {
+ private:
+  using SingleDimReshardStrategy<RToPReshardFunction>::SingleDimReshardStrategy;
+  TensorDistAttr CalculateNewDistAttr() const override {
+    TensorDistAttr real_out_dist_attr = ctx_.out->dist_attr();
+    real_out_dist_attr.set_partial_status(std::vector<int64_t>{cur_mesh_dim_});
+    return real_out_dist_attr;
+  }
+
+  TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    return ctx_.CreateOneDimDistAttr(sub_mesh);
+  }
+
+  TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    return ctx_.CreateOneDimDistAttr(sub_mesh, true);
+  }
+};
+
+class ReplicateToShard final
+    : public SingleDimReshardStrategy<RToSReshardFunction> {
+ private:
+  using SingleDimReshardStrategy<RToSReshardFunction>::SingleDimReshardStrategy;
+
+  TensorDistAttr CalculateNewDistAttr() const override {
+    TensorDistAttr real_out_dist_attr(ctx_.out->dist_attr());
+    std::vector<std::vector<int64_t>> real_dims_mapping =
+        real_out_dist_attr.multi_dims_mapping();
+    real_dims_mapping[cur_tensor_dim_].push_back(cur_mesh_dim_);
+    real_out_dist_attr.set_dims_mapping(real_dims_mapping);
+
+    auto split_factor = ctx_.out_dist_attr.get_split_factor(cur_mesh_dim_);
+    real_out_dist_attr.set_split_factor(cur_mesh_dim_, split_factor);
+    return real_out_dist_attr;
+  }
+
+  TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    return ctx_.CreateOneDimDistAttr(sub_mesh);
+  }
+
+  TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    auto split_factor = ctx_.out_dist_attr.get_split_factor(cur_mesh_dim_);
+    VLOG(3) << "In R to S mesh dim is " << cur_mesh_dim_ << ", split factor is "
+            << split_factor;
+    return ctx_.CreateOneDimDistAttr(
+        sub_mesh, false, std::nullopt, cur_tensor_dim_, split_factor);
+  }
+};
+
+class PartialToShard final
+    : public SingleDimReshardStrategy<PToSReshardFunction> {
+ private:
+  using SingleDimReshardStrategy<PToSReshardFunction>::SingleDimReshardStrategy;
+
+  TensorDistAttr CalculateNewDistAttr() const override {
+    TensorDistAttr real_out_dist_attr(ctx_.out->dist_attr());
+    std::vector<std::vector<int64_t>> real_dims_mapping =
+        real_out_dist_attr.multi_dims_mapping();
+    real_dims_mapping[cur_tensor_dim_].push_back(cur_mesh_dim_);
+    real_out_dist_attr.set_dims_mapping(real_dims_mapping);
+    if (real_out_dist_attr.is_partial(cur_mesh_dim_)) {
+      real_out_dist_attr.clean_partial_dims({cur_mesh_dim_});
+    }
+
+    auto split_factor = ctx_.out_dist_attr.get_split_factor(cur_mesh_dim_);
+    real_out_dist_attr.set_split_factor(cur_mesh_dim_, split_factor);
+    return real_out_dist_attr;
+  }
+
+  TensorDistAttr CreateOneDimInDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    auto input_reduce_type =
+        ctx_.out->dist_attr().partial_status().at(cur_mesh_dim_);
+    return ctx_.CreateOneDimDistAttr(sub_mesh, true, input_reduce_type);
+  }
+
+  TensorDistAttr CreateOneDimOutDistAttr(
+      const ProcessMesh& sub_mesh) const override {
+    auto split_factor = ctx_.out_dist_attr.get_split_factor(cur_mesh_dim_);
+    VLOG(3) << "In P to S mesh dim is " << cur_mesh_dim_ << ", split factor is "
+            << split_factor;
+    return ctx_.CreateOneDimDistAttr(
+        sub_mesh, false, std::nullopt, cur_tensor_dim_, split_factor);
+  }
+};
+
+void ProcessPartialToReplicated(DeviceContext* dev_ctx,
+                                const DistTensor& in,
+                                TensorDistAttr out_dist_attr,
+                                DistTensor* out) {
+  if (in.dist_attr().is_partial()) {
+    auto partial_status = in.dist_attr().partial_status();
+    auto out_partial_status = out_dist_attr.partial_status();
+    ReshardContext ctx(dev_ctx, in, out_dist_attr, out);
+    for (const auto& [k, v] : partial_status) {
+      VLOG(3) << "Step1: partial axis " << k;
+      if (out_partial_status.count(k) != 0 || out_dist_attr.is_shard(k)) {
+        continue;
+      }
+      auto strategy = std::make_unique<PartialToReplicate>(-1, k, ctx);
+      strategy->Eval();
+    }
+    VLOG(3) << "After P to R, dist attr is " << out->dist_attr();
+  }
+}
+
+void ProcessShardToReplicated(DeviceContext* dev_ctx,
+                              const DistTensor& in,
+                              TensorDistAttr out_dist_attr,
+                              DistTensor* out) {
+  auto is_same_shard = [&out_dist_attr, &out](
+                           std::vector<int64_t> in_mesh_axis,
+                           std::vector<int64_t> out_mesh_axis) {
+    if (in_mesh_axis != out_mesh_axis) {
+      return false;
+    }
+    for (auto dim : in_mesh_axis) {
+      if (out_dist_attr.get_split_factor(dim) !=
+          out->dist_attr().get_split_factor(dim)) {
+        return false;
+      }
+    }
+    return true;
+  };
+  int64_t first_diff_axis =
+      FindFirstDiffShardAxis(out->dist_attr(), out_dist_attr);
+  VLOG(3) << "In S to R, first diff axis is " << first_diff_axis;
+  for (int cur_tensor_dim = first_diff_axis; cur_tensor_dim >= 0;
+       --cur_tensor_dim) {
+    auto in_mesh_axis = out->dist_attr().multi_dims_mapping()[cur_tensor_dim];
+    auto out_mesh_axis = out_dist_attr.multi_dims_mapping()[cur_tensor_dim];
+    if (in_mesh_axis.size() == 0 ||
+        is_same_shard(in_mesh_axis, out_mesh_axis)) {
+      continue;
+    }
+    VLOG(3) << "Step2: in_mesh axis " << auto_parallel::str_join(in_mesh_axis);
+    ReshardContext ctx(dev_ctx, in, out_dist_attr, out);
+    for (int64_t idx = in_mesh_axis.size() - 1; idx >= 0; idx--) {
+      int64_t cur_mesh_dim = in_mesh_axis.at(idx);
+      auto strategy =
+          std::make_unique<ShardToReplicate>(cur_tensor_dim, cur_mesh_dim, ctx);
+      strategy->Eval();
+    }
+  }
+  if (first_diff_axis >= 0) {
+    VLOG(3) << "After S to R, dist attr is " << out->dist_attr();
+  }
+}
+
+void ProcessReplicatedToPartial(DeviceContext* dev_ctx,
+                                const DistTensor& in,
+                                TensorDistAttr out_dist_attr,
+                                DistTensor* out) {
+  if (out_dist_attr.is_partial()) {
+    const auto& in_partial_status = out->dist_attr().partial_status();
+    const auto& out_partial_status = out_dist_attr.partial_status();
+    for (const auto& [k, v] : out_partial_status) {
+      if (in_partial_status.count(k) != 0) {
+        continue;
+      }
+      VLOG(3) << "Step3: Partial status mesh axis " << k;
+      ReshardContext ctx(dev_ctx, in, out_dist_attr, out);
+      auto strategy = std::make_unique<ReplicateToPartial>(-1, k, ctx);
+      strategy->Eval();
+    }
+    VLOG(3) << "After R to P, dist attr is " << out->dist_attr();
+  }
+}
+
+void ProcessReplicateOrPartialToShard(DeviceContext* dev_ctx,
+                                      const DistTensor& in,
+                                      TensorDistAttr out_dist_attr,
+                                      DistTensor* out) {
+  int64_t first_diff_axis =
+      FindFirstDiffShardAxis(out->dist_attr(), out_dist_attr);
+  VLOG(3) << "In P or R to S, first diff axis is " << first_diff_axis;
+  for (int64_t cur_tensor_dim = first_diff_axis; cur_tensor_dim >= 0;
+       --cur_tensor_dim) {
+    const auto& in_mesh_axis =
+        out->dist_attr().multi_dims_mapping()[cur_tensor_dim];
+    const auto& out_mesh_axis =
+        out_dist_attr.multi_dims_mapping()[cur_tensor_dim];
+    if (in_mesh_axis == out_mesh_axis) {
+      continue;
+    }
+
+    const auto& in_partial_status = out->dist_attr().partial_status();
+    ReshardContext ctx(dev_ctx, in, out_dist_attr, out);
+
+    for (auto cur_mesh_dim : out_mesh_axis) {
+      bool is_partial = in_partial_status.count(cur_mesh_dim) != 0;
+      VLOG(3) << "Step4: out_mesh axis : " << cur_mesh_dim
+              << "; partial state :" << is_partial;
+      std::shared_ptr<SameNdMeshReshardFunction::ReshardStrategy> strategy;
+      if (is_partial) {
+        strategy =
+            std::make_unique<PartialToShard>(cur_tensor_dim, cur_mesh_dim, ctx);
+      } else {
+        strategy = std::make_unique<ReplicateToShard>(
+            cur_tensor_dim, cur_mesh_dim, ctx);
+      }
+      strategy->Eval();
+    }
+  }
+  if (first_diff_axis >= 0) {
+    VLOG(3) << "After P or R to S, dist attr is " << out->dist_attr();
+  }
+}
+
+bool SameNdMeshReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  RESHARD_SHORTCUT_IF_FALSE(in.dist_attr().process_mesh() ==
+                            out_dist_attr.process_mesh());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.process_mesh().ndim() > 1);
+
+  // check the input and output dims_mapping is not equal
+  RESHARD_SHORTCUT_IF_FALSE(in.dist_attr() != out_dist_attr);
+
+  return true;
+}
+
+void SameNdMeshReshardFunction::Eval(DeviceContext* dev_ctx,
+                                     const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr,
+                                     DistTensor* out) {
+  auto out_dist_attr_orig = out_dist_attr;
+  SetValue(out, in.value());
+  SetDistProps(out, in.dims(), in.dist_attr());
+  // 1. change all the partial status to replicated status if needed
+  ProcessPartialToReplicated(dev_ctx, in, out_dist_attr_orig, out);
+  // 2. change all the shard status to replicated status
+  ProcessShardToReplicated(dev_ctx, in, out_dist_attr_orig, out);
+  // 3. Change replicated to partial
+  ProcessReplicatedToPartial(dev_ctx, in, out_dist_attr_orig, out);
+  // 4. Change replicated/partial to shard
+  ProcessReplicateOrPartialToShard(dev_ctx, in, out_dist_attr_orig, out);
+
+  // TODO(lfw): refine this, now is reports wrong info.
+  // Final attr check
+  // PADDLE_ENFORCE_EQ(out->dist_attr() == out_dist_attr_orig,
+  //                   true,
+  //                   ::common::errors::InvalidArgument("Expected that out of
+  //                   reshard has to be target dist, " "out dist att is " +
+  //                   out->dist_attr().to_string() + ", but target is " +
+  //                   out_dist_attr_orig.to_string()));
+}
+
+bool CrossNdMeshReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const ProcessMesh& in_process_mesh = in.dist_attr().process_mesh();
+  const ProcessMesh& out_process_mesh = out_dist_attr.process_mesh();
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.process_mesh().ndim() > 1);
+
+  // check the input and output dims_mapping is not equal
+  RESHARD_SHORTCUT_IF_FALSE(in.dist_attr() != out_dist_attr);
+
+  return true;
+}
+
+void CrossNdMeshReshardFunction::Eval(DeviceContext* dev_ctx,
+                                      const DistTensor& in,
+                                      const TensorDistAttr& out_dist_attr,
+                                      DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+
+  // Construct a `DistTensor` by `dtype` of `in` tensor to avoid using default
+  // dtype `float32`. The default dtype `float32` may cause error in amp.
+  DistTensor tmp_result(in.dtype());
+  TensorDistAttr in_dist_attr_shard = in_dist_attr;
+  in_dist_attr_shard.set_partial_status(out_dist_attr.partial_status());
+  in_dist_attr_shard.set_dims_mapping(out_dist_attr.dims_mapping());
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_dist_attr.process_mesh().contains(cur_global_rank)) {
+    SameNdMeshReshardFunction same_nd_reshard_func;
+    PADDLE_ENFORCE(
+        same_nd_reshard_func.IsSuitable(in, in_dist_attr_shard),
+        common::errors::InvalidArgument(
+            "Invoke the same nd reshard function is not valid from %s to %s.",
+            in_dist_attr,
+            in_dist_attr_shard));
+    same_nd_reshard_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
+  } else {
+    SetDistProps(&tmp_result, in.dims(), in_dist_attr_shard);
+  }
+
+  SameStatusReshardFunction same_status_func;
+  PADDLE_ENFORCE(
+      same_status_func.IsSuitable(tmp_result, out_dist_attr),
+      common::errors::InvalidArgument("Invoke the same status reshard function "
+                                      "is not valid from %s to %s.",
+                                      tmp_result.dist_attr(),
+                                      out_dist_attr));
+  same_status_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.h

@@ -0,0 +1,62 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API SameNdMeshReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SameNdMeshReshard"; }
+
+  class PADDLE_API ReshardStrategy {
+   public:
+    virtual ~ReshardStrategy() = default;
+    virtual void Eval() = 0;
+    void SetValue(DistTensor* tensor, const DenseTensor& value) {
+      ReshardFunction::SetValue(tensor, value);
+    }
+
+    void SetDistProps(DistTensor* tensor, const TensorDistAttr& dist_attr) {
+      ReshardFunction::SetDistProps(tensor, dist_attr);
+    }
+  };
+};
+
+class PADDLE_API CrossNdMeshReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "CrossNdMeshReshard"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.cc

@@ -0,0 +1,149 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/all_reduce_kernel.h"
+#include "paddle/phi/kernels/elementwise_divide_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+
+namespace phi::distributed {
+
+bool PToRReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  RESHARD_SHORTCUT_IF_FALSE(in.dist_attr().is_partial());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in.dist_attr().process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void PToRReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& in_process_ids = in_process_mesh.process_ids();
+  if (in_process_ids.size() == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+  const auto& in_partial_status = in_dist_attr.partial_status();
+  auto in_reduce_type = in_partial_status.at(0);
+  bool reduce_mean = false;
+  auto dtype = in.dtype();
+
+  if (in_reduce_type == ReduceType::kRedAvg) {
+    in_reduce_type = ReduceType::kRedSum;
+    reduce_mean = true;
+  }
+  int64_t reduce_type = static_cast<int64_t>(in_reduce_type);
+  VLOG(3) << "Transfer from partial to replicated status with reduce type "
+          << reduce_type;
+
+  RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                            AllReduce,
+                            dtype,
+                            in_process_ids,
+                            in.value(),
+                            reduce_type,
+                            GetMutableTensor(out));
+
+  if (reduce_mean) {
+    VLOG(3) << "Do reduce mean after all reduce sum";
+    DenseTensor tensor_of_num_process;
+    IntArray shape({1});
+    RESHARD_FUNCTOR(dev_ctx,
+                    Full,
+                    in.dtype(),
+                    shape,
+                    static_cast<int64_t>(in_process_ids.size()),
+                    &tensor_of_num_process);
+    RESHARD_FUNCTOR(dev_ctx,
+                    Divide,
+                    dtype,
+                    out->value(),
+                    tensor_of_num_process,
+                    GetMutableTensor(out));
+  }
+
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool PToRReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_partial());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void PToRReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  DistTensor tmp_result;
+
+  SameStatusReshardFunction same_status_func;
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    PToRReshardFunction p_to_r_func;
+    PADDLE_ENFORCE(
+        p_to_r_func.IsSuitable(tmp_result, out_dist_attr),
+        common::errors::InvalidArgument(
+            "Invoke the p to r reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    p_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h

@@ -0,0 +1,52 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API PToRReshardFunction final : public ReshardFunction {
+ public:
+  PToRReshardFunction() = default;
+  ~PToRReshardFunction() = default;
+
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "PToRReshard"; }
+};
+
+class PADDLE_API PToRReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "PToRReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.cc

@@ -0,0 +1,238 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/reduce_scatter_kernel.h"
+#include "paddle/phi/kernels/split_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi::distributed {
+
+bool PToSReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_partial());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void ReshardPToSWithPadding(DeviceContext* dev_ctx,
+                            int64_t split_axis,
+                            const std::vector<int64_t>& process_ids,
+                            const DenseTensor& in,
+                            int64_t padding_nums,
+                            DenseTensor* out) {
+  DenseTensor in_reduce_scatter;
+  std::vector<int> axis;
+  const auto& logical_ddim = in.dims();
+  auto dtype = in.dtype();
+
+  if (split_axis != 0) {
+    for (size_t i = 0; i < common::vectorize(logical_ddim).size(); ++i) {
+      axis.emplace_back(i);
+    }
+    std::swap(axis[0], axis[split_axis]);
+    RESHARD_FUNCTOR(dev_ctx, Transpose, dtype, in, axis, &in_reduce_scatter);
+  } else {
+    in_reduce_scatter.ShareDataNoCheckWith(in);
+  }
+
+  DenseTensor out_reduce_scatter;
+  RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                            ReduceScatter,
+                            dtype,
+                            process_ids,
+                            in_reduce_scatter,
+                            static_cast<int64_t>(process_ids.size()),
+                            &out_reduce_scatter);
+
+  DenseTensor out_result;
+  if (split_axis != 0) {
+    RESHARD_FUNCTOR(
+        dev_ctx, Transpose, dtype, out_reduce_scatter, axis, &out_result);
+  } else {
+    out_result.ShareDataNoCheckWith(out_reduce_scatter);
+  }
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (cur_global_rank == process_ids.back() && padding_nums != 0) {
+    std::vector<DenseTensor> tmp_out_vec;
+    IntArray tmp_sections(std::vector<int64_t>{
+        out_result.dims()[split_axis] - padding_nums, padding_nums});
+    RESHARD_FUNCTOR(dev_ctx,
+                    Split,
+                    dtype,
+                    out_result,
+                    tmp_sections,
+                    split_axis,
+                    &tmp_out_vec);
+    // TODO(liyurui): Since we can not separate local tensor with [0, 10] shape
+    // and uninitialized tensor, here we use a tricky solution.
+    // Give local tensor which has, for example [0, 10] shape, a little
+    // allocation, to make it difference from uninitialized tensor in pipeline
+    // strategy.
+    if (tmp_out_vec[0].dims()[split_axis] == 0) {
+      tmp_out_vec[0].mutable_data(tmp_out_vec[0].place(), 4);
+    }
+    out->ShareDataNoCheckWith(tmp_out_vec[0]);
+  } else {
+    out->ShareDataNoCheckWith(out_result);
+  }
+}
+
+void PToSReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& in_process_ids = in_process_mesh.process_ids();
+
+  int out_split_axis =
+      GetSplitAxisWithDimsMapping(out_dist_attr.dims_mapping()).begin()->first;
+  int64_t num_of_process = in_process_mesh.size();
+  if (num_of_process == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+  int64_t num_of_padding = in.dims()[out_split_axis] % num_of_process;
+  bool is_balanced_split = (num_of_padding == 0);
+
+  if (is_balanced_split) {
+    VLOG(3) << "Balanced reshard from partial to shard";
+    ReshardPToSWithPadding(dev_ctx,
+                           out_split_axis,
+                           in_process_ids,
+                           in.value(),
+                           /*padding_nums*/ 0,
+                           GetMutableTensor(out));
+  } else {
+    VLOG(3) << "Unbalanced reshard from partial to shard";
+    int64_t avg_size_on_split_axis =
+        (in.dims()[out_split_axis] + num_of_process - 1) / num_of_process;
+    int64_t padding_nums =
+        avg_size_on_split_axis * num_of_process - in.dims()[out_split_axis];
+
+    DDim concat_local_shape = in.local_dims();
+    concat_local_shape[out_split_axis] = padding_nums;
+    IntArray concat_local_shape_int_array(concat_local_shape.Get(),
+                                          concat_local_shape.size());
+    auto dtype = in.dtype();
+
+    DenseTensor concat_local_tensor;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Full,
+                    dtype,
+                    concat_local_shape_int_array,
+                    0,
+                    &concat_local_tensor);
+
+    DenseTensor in_local_tensor = in.value();
+    std::vector<const DenseTensor*> concat_input_vec = {&in_local_tensor,
+                                                        &concat_local_tensor};
+
+    DenseTensor concat_result;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Concat,
+                    dtype,
+                    concat_input_vec,
+                    out_split_axis,
+                    &concat_result);
+
+    ReshardPToSWithPadding(dev_ctx,
+                           out_split_axis,
+                           in_process_ids,
+                           concat_result,
+                           padding_nums,
+                           GetMutableTensor(out));
+  }
+
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool PToSReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_partial());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void PToSReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+
+  DistTensor tmp_result;
+  TensorDistAttr in_dist_attr_shard = in_dist_attr;
+  in_dist_attr_shard.clean_partial_status();
+  in_dist_attr_shard.set_dims_mapping(out_dist_attr.dims_mapping());
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_dist_attr.process_mesh().contains(cur_global_rank)) {
+    PToSReshardFunction p_to_s_func;
+    PADDLE_ENFORCE(
+        p_to_s_func.IsSuitable(in, in_dist_attr_shard),
+        common::errors::InvalidArgument(
+            "Invoke the p to s reshard function is not valid from %s to %s.",
+            in_dist_attr,
+            in_dist_attr_shard));
+    p_to_s_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
+  } else {
+    SetDistProps(&tmp_result, in.dims(), in_dist_attr_shard);
+    SetValue(&tmp_result, in.value());
+  }
+
+  SameStatusReshardFunction same_status_func;
+  PADDLE_ENFORCE(
+      same_status_func.IsSuitable(tmp_result, out_dist_attr),
+      common::errors::InvalidArgument("Invoke the same status reshard function "
+                                      "is not valid from %s to %s.",
+                                      tmp_result.dist_attr(),
+                                      out_dist_attr));
+  same_status_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API PToSReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "PToSReshard"; }
+};
+
+class PADDLE_API PToSReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "PToSReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.cc

@@ -0,0 +1,129 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/assign_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+
+namespace phi::distributed {
+
+bool RToPReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void RToPReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_process_ids = in.dist_attr().process_mesh().process_ids();
+  if (in_process_ids.size() == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+  int64_t local_rank = GetCurRankCoordInMesh(out_process_mesh)[0];
+  const auto& in_reduce_type = out_dist_attr.partial_status().at(0);
+
+  if (local_rank != 0) {
+    if (in_reduce_type == ReduceType::kRedAvg) {
+      // assign the input value to output
+      RESHARD_FUNCTOR_WITHOUT_DTYPE(
+          dev_ctx, Assign, in.value(), GetMutableTensor(out));
+    } else {
+      // reset the physical tensor to zero
+      IntArray shape(in.local_dims().Get(), in.local_dims().size());
+      RESHARD_FUNCTOR(
+          dev_ctx, Full, in.dtype(), shape, 0, GetMutableTensor(out));
+    }
+  } else {
+    // assign the input value to output
+    RESHARD_FUNCTOR_WITHOUT_DTYPE(
+        dev_ctx, Assign, in.value(), GetMutableTensor(out));
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool RToPReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void RToPReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  DistTensor tmp_result;
+
+  SameStatusReshardFunction same_status_func;
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    RToPReshardFunction r_to_p_func;
+    PADDLE_ENFORCE(
+        r_to_p_func.IsSuitable(tmp_result, out_dist_attr),
+        common::errors::InvalidArgument(
+            "Invoke the r to p reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    r_to_p_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h

@@ -0,0 +1,49 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API RToPReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "RToPReshard"; }
+};
+
+class PADDLE_API RToPReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "RToPReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.cc

@@ -0,0 +1,196 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/slice_kernel.h"
+
+namespace phi::distributed {
+
+bool RToSReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+namespace {
+std::map<int, int64_t> GetSplitAxisWithDimsMapping(
+    const std::vector<std::vector<int64_t>>& dims_mapping) {
+  std::map<int, int64_t> split_axis_to_mesh_axis;
+  for (size_t i = 0; i < dims_mapping.size(); ++i) {
+    if (dims_mapping[i].size() > 0) {
+      split_axis_to_mesh_axis.emplace(i, dims_mapping[i][0]);
+    }
+  }
+  return split_axis_to_mesh_axis;
+}
+}  // namespace
+
+void RToSReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_dims_mapping = out_dist_attr.multi_dims_mapping();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+  const DenseTensor& in_physical_tensor_cur_rank = in.value();
+
+  std::map<int, int64_t> split_axis_to_mesh_axis =
+      GetSplitAxisWithDimsMapping(out_dims_mapping);
+  std::vector<int64_t> coord_in_mesh = GetCurRankCoordInMesh(out_process_mesh);
+
+  int split_axis = split_axis_to_mesh_axis.begin()->first;
+  int64_t mesh_axis = split_axis_to_mesh_axis.begin()->second;
+
+  VLOG(3) << "split axis is " << split_axis << ", mesh axis is " << mesh_axis;
+  VLOG(3) << "shape size is " << out_process_mesh.shape().size();
+
+  int64_t num_of_process = out_process_mesh.shape()[mesh_axis];
+  if (num_of_process == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+  VLOG(3) << "RToSReshard: Tensor will be split on axis " << split_axis
+          << ". Slice will use axis " << mesh_axis << " of process_mesh."
+          << " There will have " << num_of_process
+          << " process participate in.";
+
+  int64_t num_group = out_dist_attr.get_split_factor(mesh_axis);
+  VLOG(3) << "num group = " << num_group;
+
+  std::vector<int64_t> split_num_vec =
+      BalancedSplit(in.value().dims()[split_axis], num_of_process * num_group);
+
+  auto dtype = in_physical_tensor_cur_rank.dtype();
+
+  int64_t slice_stride = num_of_process;
+  std::vector<DenseTensor> dense_out_vec(num_group);
+  for (int64_t i = 0; i < num_group; i++) {
+    int64_t start =
+        split_num_vec[0] * (coord_in_mesh[mesh_axis] + i * slice_stride);
+    int64_t end = std::min(start + split_num_vec[0],
+                           in_physical_tensor_cur_rank.dims()[split_axis]);
+    VLOG(4) << "start is " << start << ", end is " << end;
+
+    PADDLE_ENFORCE_LE(start,
+                      end,
+                      common::errors::InvalidArgument(
+                          "Slice Args 'start' should be less or qual to 'end', "
+                          "but got 'start' is %d, 'end' is %d.",
+                          start,
+                          end));
+    RESHARD_FUNCTOR(dev_ctx,
+                    Slice,
+                    dtype,
+                    in_physical_tensor_cur_rank,
+                    {split_axis},
+                    {start},
+                    {end},
+                    &(dense_out_vec[i]));
+  }
+
+  if (num_group == 1) {
+    SetValue(out, dense_out_vec[0]);
+  } else {
+    std::vector<const DenseTensor*> d_tensor_ptr_vec;
+    for (auto& d_tensor : dense_out_vec) {
+      d_tensor_ptr_vec.push_back(&d_tensor);
+    }
+
+    DenseTensor dense_out;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Concat,
+                    dtype,
+                    d_tensor_ptr_vec,
+                    Scalar(split_axis),
+                    &dense_out);
+    SetValue(out, dense_out);
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool RToSReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void RToSReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+
+  DistTensor tmp_result;
+  TensorDistAttr in_dist_attr_shard = in_dist_attr;
+  in_dist_attr_shard.set_dims_mapping(out_dist_attr.dims_mapping());
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_dist_attr.process_mesh().contains(cur_global_rank)) {
+    RToSReshardFunction r_to_s_func;
+    PADDLE_ENFORCE(
+        r_to_s_func.IsSuitable(in, in_dist_attr_shard),
+        common::errors::InvalidArgument(
+            "Invoke the r to s reshard function is not valid from %s to %s.",
+            in_dist_attr,
+            in_dist_attr_shard));
+    r_to_s_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
+  } else {
+    SetDistProps(&tmp_result, in.dims(), in_dist_attr_shard);
+    SetValue(&tmp_result, in.value());
+  }
+  SameStatusReshardFunction same_status_func;
+  PADDLE_ENFORCE(
+      same_status_func.IsSuitable(tmp_result, out_dist_attr),
+      common::errors::InvalidArgument("Invoke the same status reshard function "
+                                      "is not valid from %s to %s.",
+                                      tmp_result.dist_attr(),
+                                      out_dist_attr));
+  same_status_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h

@@ -0,0 +1,49 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API RToSReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "RToSReshard"; }
+};
+
+class PADDLE_API RToSReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "RToSReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.cc

@@ -0,0 +1,144 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/add_n_kernel.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/p_recv_kernel.h"
+#include "paddle/phi/kernels/p_send_kernel.h"
+#include "paddle/phi/kernels/split_kernel.h"
+
+namespace phi::distributed {
+
+bool RToXExpandReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.process_ids().size() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.process_ids().size() != 1);
+
+  return true;
+}
+
+void RToXExpandReshardFunction::Eval(DeviceContext* dev_ctx,
+                                     const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr,
+                                     DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& out_dims_mapping = out_dist_attr.dims_mapping();
+  const auto& in_mesh = in_dist_attr.process_mesh();
+  const auto& out_mesh = out_dist_attr.process_mesh();
+  const auto& in_process_ids = in_mesh.process_ids();
+  const auto& out_process_ids = out_mesh.process_ids();
+  int64_t cur_global_rank = GetCurGlobalRank();
+  int64_t root_rank = in_process_ids[0];
+  auto all_process_ids = GetUnionProcessIds(in_process_ids, out_process_ids);
+  auto dtype = in.dtype();
+  const auto& out_partial_status = out_dist_attr.partial_status();
+  bool cur_rank_in_out_mesh =
+      (std::find(out_process_ids.begin(),
+                 out_process_ids.end(),
+                 cur_global_rank) != out_process_ids.end());
+  DenseTensor result_value;
+
+  if (root_rank == cur_global_rank) {
+    for (const auto& out_process_id : out_process_ids) {
+      if (out_process_id != root_rank) {
+        // Should acculate the relative rank in communication.
+        auto relative_rank = out_process_id;
+        for (size_t i = 0; i < all_process_ids.size(); ++i) {
+          if (all_process_ids[i] == out_process_id) {
+            relative_rank = i;
+          }
+        }
+
+        RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                  PSendKernel,
+                                  dtype,
+                                  all_process_ids,
+                                  in.value(),
+                                  /*peer*/ relative_rank,
+                                  /*dynamic_shape=*/true);
+      }
+    }
+    if (cur_rank_in_out_mesh) {
+      result_value = in.value();
+    }
+  } else {
+    RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                              PRecv,
+                              dtype,
+                              all_process_ids,
+                              root_rank,
+                              {} /*out_shape*/,
+                              /*dynamic_shape=*/true,
+                              &result_value);
+  }
+
+  if (cur_rank_in_out_mesh) {
+    if (out_dist_attr.is_partial()) {
+      auto out_reduce_type = out_partial_status.at(0);
+      if (out_reduce_type == ReduceType::kRedSum &&
+          cur_global_rank != out_process_ids[0]) {
+        IntArray shape(result_value.dims().Get(), result_value.dims().size());
+        RESHARD_FUNCTOR(dev_ctx, Full, dtype, shape, 0, &result_value);
+      }
+      SetValue(out, result_value);
+    } else if (out_dist_attr.is_shard()) {
+      std::map<int, int64_t> split_axis_to_mesh_axis =
+          GetSplitAxisWithDimsMapping(out_dims_mapping);
+      std::vector<int64_t> coord_in_mesh = GetCurRankCoordInMesh(out_mesh);
+
+      int split_axis = split_axis_to_mesh_axis.begin()->first;
+      int64_t mesh_axis = split_axis_to_mesh_axis.begin()->second;
+      int64_t num_of_process = out_mesh.shape()[mesh_axis];
+
+      std::vector<int64_t> split_num_vec =
+          BalancedSplit(in.dims()[split_axis], num_of_process);
+      IntArray sections(split_num_vec);
+
+      std::vector<DenseTensor> split_out_vec;
+      RESHARD_FUNCTOR(dev_ctx,
+                      Split,
+                      dtype,
+                      result_value,
+                      sections,
+                      split_axis,
+                      &split_out_vec);
+
+      SetValue(out, split_out_vec[coord_in_mesh[mesh_axis]]);
+    } else {
+      SetValue(out, result_value);
+    }
+    SetDistProps(out, in.dims(), out_dist_attr);
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.h

@@ -0,0 +1,36 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class RToXExpandReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "RToXExpandReshard"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc

@@ -0,0 +1,76 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+
+namespace phi::distributed {
+
+using phi::distributed::auto_parallel::str_join;
+
+std::shared_ptr<DistTensor> ReshardFunction::Eval(
+    DeviceContext* dev_ctx,
+    const DistTensor& in,
+    const TensorDistAttr& out_dist_attr) {
+  phi::RecordEvent reshard_record_event(
+      Name(), phi::TracerEventType::OperatorInner, 1);
+  std::shared_ptr<DistTensor> out = std::make_shared<DistTensor>();
+  Eval(dev_ctx, in, out_dist_attr, out.get());
+  return out;
+}
+
+void ReshardFunction::SetValue(DistTensor* tensor, const DenseTensor& value) {
+  tensor->value_ = std::make_shared<DenseTensor>(value);
+}
+
+void ReshardFunction::SetDistProps(DistTensor* tensor,
+                                   const DDim& dims,
+                                   const TensorDistAttr& dist_attr) {
+  PADDLE_ENFORCE_EQ(dist_attr.verify_dynamic(common::vectorize(dims)),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input dist_attr [%s] and dims [%s] are improper.",
+                        dist_attr.to_string(),
+                        str_join(vectorize(dims))));
+
+  tensor->global_dims_ = dims;
+  tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
+}
+
+void ReshardFunction::SetDistProps(DistTensor* tensor,
+                                   const TensorDistAttr& dist_attr) {
+  PADDLE_ENFORCE_EQ(dist_attr.verify_dynamic(common::vectorize(tensor->dims())),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input dist_attr [%s] and dims [%s] are improper.",
+                        dist_attr.to_string(),
+                        str_join(vectorize(tensor->dims()))));
+
+  tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
+}
+
+DenseTensor* ReshardFunction::GetMutableTensor(DistTensor* tensor) {
+  return tensor->value_.get();
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h

@@ -0,0 +1,58 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+#include <memory>
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+class DeviceContext;
+
+namespace distributed {
+
+class DistTensor;
+class TensorDistAttr;
+
+class PADDLE_API ReshardFunction {
+ public:
+  ReshardFunction() = default;
+  virtual ~ReshardFunction() = default;
+
+  virtual bool IsSuitable(const DistTensor& in,
+                          const TensorDistAttr& out_dist_attr) = 0;
+
+  std::shared_ptr<DistTensor> Eval(DeviceContext* dev_ctx,
+                                   const DistTensor& in,
+                                   const TensorDistAttr& out_dist_attr);
+
+  virtual void Eval(DeviceContext* dev_ctx,
+                    const DistTensor& in,
+                    const TensorDistAttr& out_dist_attr,
+                    DistTensor* out) = 0;
+
+  virtual std::string Name() { return "ReshardBase"; }
+
+ protected:
+  static void SetValue(DistTensor* tensor, const DenseTensor& value);
+  static void SetDistProps(DistTensor* tensor,
+                           const DDim& dims,
+                           const TensorDistAttr& dist_attr);
+  static void SetDistProps(DistTensor* tensor, const TensorDistAttr& dist_attr);
+  DenseTensor* GetMutableTensor(DistTensor* tensor);
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.cc

@@ -0,0 +1,81 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/global_and_sub_mesh_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_x_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/x_to_r_reshard_function.h"
+
+namespace phi::distributed {
+
+ReshardFunction* ChooseProperReshardFunction(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  for (const auto& func : GetReshardFunctionList()) {
+    if (func->IsSuitable(in, out_dist_attr)) {
+      VLOG(4) << "Choose ReshardFunction: " << func->Name();
+      return func.get();
+    }
+  }
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Can not reshard from in_dist_attr=%s to out_dist_attr=%s.",
+      in.dist_attr().to_string(),
+      out_dist_attr.to_string()));
+}
+
+std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList() {
+  static std::vector<std::unique_ptr<ReshardFunction>> func_list;
+  return func_list;
+}
+
+// NOTE: be aware of the registration order of the reshard function.
+// Higher priority will be granted to the reshard function
+// which was registered earlier.
+// Reshard function with higher priority will be evoked
+// when more than one reshard function satisfy the request.
+REGISTER_RESHARD_FUNC(SToRReshardFunction);
+REGISTER_RESHARD_FUNC(SToRReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(SToPReshardFunction);
+REGISTER_RESHARD_FUNC(SToPReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(RToSReshardFunction);
+REGISTER_RESHARD_FUNC(RToSReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(RToPReshardFunction);
+REGISTER_RESHARD_FUNC(RToPReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(PToRReshardFunction);
+REGISTER_RESHARD_FUNC(PToRReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(PToSReshardFunction);
+REGISTER_RESHARD_FUNC(PToSReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(SToSReshardFunction);
+REGISTER_RESHARD_FUNC(SToSReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(XToRShrinkReshardFunction);
+REGISTER_RESHARD_FUNC(RToXExpandReshardFunction);
+REGISTER_RESHARD_FUNC(SameStatusReshardFunction);
+REGISTER_RESHARD_FUNC(SameNdMeshReshardFunction);
+REGISTER_RESHARD_FUNC(CrossNdMeshReshardFunction);
+REGISTER_RESHARD_FUNC(GlobalToSubMeshReshardFunction);
+REGISTER_RESHARD_FUNC(SubMeshToGlobalReshardFunction);
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.h

@@ -0,0 +1,40 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList();
+
+#define REGISTER_RESHARD_FUNC(func_type)                                    \
+  class __RegisterReshard_##func_type {                                     \
+   public:                                                                  \
+    __RegisterReshard_##func_type() {                                       \
+      GetReshardFunctionList().emplace_back(std::make_unique<func_type>()); \
+    }                                                                       \
+  };                                                                        \
+  static __RegisterReshard_##func_type local_reshard_func_##func_type
+
+ReshardFunction* ChooseProperReshardFunction(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.cc

@@ -0,0 +1,285 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi::distributed {
+
+namespace {
+std::string GenUniqueCommKey(const std::vector<int64_t>& process_ids) {
+  std::string unique_comm_key = "ReshardGroup";
+  for (const auto& id : process_ids) {
+    unique_comm_key += "/" + std::to_string(id);
+  }
+  return unique_comm_key;
+}
+}  // namespace
+
+std::vector<int64_t> GetUnionProcessIds(std::vector<int64_t> in_process_ids,
+                                        std::vector<int64_t> out_process_ids) {
+  std::vector<int64_t> result;
+  std::sort(in_process_ids.begin(), in_process_ids.end());
+  std::sort(out_process_ids.begin(), out_process_ids.end());
+  std::set_union(in_process_ids.begin(),
+                 in_process_ids.end(),
+                 out_process_ids.begin(),
+                 out_process_ids.end(),
+                 std::back_inserter(result));
+  return result;
+}
+
+int64_t GetLocalRankInParticipate(const std::vector<int64_t>& process_ids,
+                                  int64_t global_rank) {
+  if (global_rank == -1) {
+    global_rank = GetCurGlobalRank();
+  }
+  auto iter = std::find(process_ids.begin(), process_ids.end(), global_rank);
+  PADDLE_ENFORCE_NE(
+      iter,
+      process_ids.end(),
+      common::errors::NotFound(
+          "Global rank %lld cannot be found in process_mesh", global_rank));
+  return iter - process_ids.begin();
+}
+
+std::vector<int64_t> GetCurRankCoordInMesh(const ProcessMesh& process_mesh) {
+  const auto& process_shape = process_mesh.shape();
+  const auto& process_ids = process_mesh.process_ids();
+  int64_t ndims_mesh = static_cast<int64_t>(process_shape.size());
+  int64_t cur_global_rank = GetCurGlobalRank();
+
+  VLOG(3) << "Searching current global rank " << cur_global_rank
+          << " in process_mesh " << process_mesh;
+
+  auto iter =
+      std::find(process_ids.begin(), process_ids.end(), cur_global_rank);
+  PADDLE_ENFORCE_NE(
+      iter,
+      process_ids.end(),
+      common::errors::NotFound("Rank %lld cannot be found in process_mesh",
+                               cur_global_rank));
+
+  int64_t flat_idx_in_mesh = iter - process_ids.begin();
+
+  std::vector<int64_t> coord(ndims_mesh, -1);
+  for (int64_t i = ndims_mesh - 1; i >= 0; --i) {
+    coord[i] = flat_idx_in_mesh % process_shape[i];
+    flat_idx_in_mesh /= process_shape[i];
+  }
+  return coord;
+}
+
+CommContext* CreateOrGetCommContext(const DeviceContext& dev_ctx,
+                                    const std::vector<int64_t>& process_ids) {
+  std::string unique_comm_key = GenUniqueCommKey(process_ids);
+
+  if (!CommContextManager::GetInstance().Has(unique_comm_key)) {
+    int64_t world_size = static_cast<int64_t>(process_ids.size());
+    int64_t rank = GetLocalRankInParticipate(process_ids);
+    VLOG(3) << "local world size: " << world_size << " local rank: " << rank;
+    auto store = CreateOrGetGlobalTCPStore();
+    if (phi::CPUContext::classof(&dev_ctx)) {
+#if defined(PADDLE_WITH_GLOO)
+      CommContextManager::CreateGlooCommContext(store,
+                                                unique_comm_key,
+                                                static_cast<int>(rank),
+                                                static_cast<int>(world_size));
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Cannot use gloo on CPU, please turn PADDLE_WITH_GLOO flag on."));
+#endif
+    }
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    else if (phi::GPUContext::classof(&dev_ctx)) {  // NOLINT
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+      CommContextManager::CreateNCCLCommContext(store,
+                                                unique_comm_key,
+                                                static_cast<int>(rank),
+                                                static_cast<int>(world_size));
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Cannot use nccl on GPU, please turn WITH_NCCL flag on."));
+#endif
+    }
+#elif defined(PADDLE_WITH_XPU)
+    else if (phi::XPUContext::classof(&dev_ctx)) {  // NOLINT
+#if defined(PADDLE_WITH_XPU_BKCL)
+      CommContextManager::CreateBKCLCommContext(store,
+                                                unique_comm_key,
+                                                static_cast<int>(rank),
+                                                static_cast<int>(world_size));
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Cannot use xpu on GPU, please turn WITH_XPU_BKCL flag on."));
+#endif
+    }
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+    else if (phi::CustomContext::classof(&dev_ctx)) {  // NOLINT
+      CommContextManager::CreateXCCLCommContext(
+          store, unique_comm_key, dev_ctx.GetPlace(), rank, world_size);
+    }
+#endif
+    else {  // NOLINT
+      PADDLE_THROW(common::errors::Unimplemented(
+          "CommContext is only supported CPU, GPU, XPU, and CustomDevice."));
+    }
+  }
+
+  auto* comm_context = CommContextManager::GetInstance().Get(unique_comm_key);
+  return comm_context;
+}
+
+std::map<int, int64_t> GetSplitAxisWithDimsMapping(
+    const std::vector<int64_t>& dims_mapping) {
+  std::map<int, int64_t> split_axis_to_mesh_axis;
+  for (size_t i = 0; i < dims_mapping.size(); ++i) {
+    if (dims_mapping[i] != -1) {
+      split_axis_to_mesh_axis.emplace(i, dims_mapping[i]);
+    }
+  }
+  return split_axis_to_mesh_axis;
+}
+
+std::vector<int64_t> BalancedSplit(int64_t total_nums, int64_t num_of_pieces) {
+  bool has_remainder = (total_nums % num_of_pieces != 0);
+  std::vector<int64_t> result(num_of_pieces,
+                              (total_nums + num_of_pieces - 1) / num_of_pieces);
+  if (has_remainder) {
+    int64_t& last_value = result.back();
+    last_value = last_value - (last_value * num_of_pieces - total_nums);
+  }
+  return result;
+}
+
+bool IsCurRankInMesh(const ProcessMesh& process_mesh) {
+  int64_t cur_global_rank = GetCurGlobalRank();
+  const auto& process_ids = process_mesh.process_ids();
+  return (std::find(process_ids.begin(), process_ids.end(), cur_global_rank) !=
+          process_ids.end());
+}
+
+// Only Input is DistTensor and current device id isn't in DistTensor's mesh
+// will return true.
+bool NeedComputationClipForPP(
+    const std::shared_ptr<phi::TensorBase>& tensor_impl) {
+  PADDLE_ENFORCE_EQ(
+      DistTensor::classof(tensor_impl.get()),
+      true,
+      common::errors::InvalidArgument(
+          "The input tensor of NeedComputationClipForPP should be "
+          "``DistTensor``. "
+          "However it's %s",
+          typeid(tensor_impl.get()).name()));
+  return !IsCurRankInMesh(std::static_pointer_cast<DistTensor>(tensor_impl)
+                              ->dist_attr()
+                              .process_mesh());
+}
+
+Place GetDefaultPlace() {
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+  auto dev_types = phi::DeviceManager::GetAllCustomDeviceTypes();
+  if (phi::DeviceManager::GetDeviceCount(dev_types[0]) > 0) {
+    return paddle::DefaultCustomPlace();
+  }
+#elif defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (phi::backends::gpu::GetGPUDeviceCount() >= 0) {
+    return paddle::DefaultGPUPlace();
+  }
+#elif defined(PADDLE_WITH_XPU)
+  if (phi::backends::xpu::GetXPUDeviceCount() >= 0) {
+    return paddle::DefaultXPUPlace();
+  }
+#endif
+  return paddle::CPUPlace();
+}
+
+DeviceContext* GetDistTensorDeviceContext(DistTensor* input) {
+  // TODO(GhostScreaming): pipeline parallel may create an undefined middle grad
+  // tensor. In such case, we need to get default place.
+  auto place =
+      input && input->initialized() ? input->place() : GetDefaultPlace();
+  return phi::DeviceContextPool::Instance().Get(place);
+}
+
+DDim InferShapeForReshardFromReplicate(
+    const std::shared_ptr<DenseTensor>& global_value,
+    const TensorDistAttr& dist_attr) {
+  DDim out_dim = global_value->dims();
+  auto coord_id = GetCurRankCoordInMesh(dist_attr.process_mesh());
+  for (int tensor_axis = 0; tensor_axis < global_value->dims().size();
+       ++tensor_axis) {
+    if (dist_attr.is_shard(-1, tensor_axis)) {
+      for (int mesh_axis = 0; mesh_axis < dist_attr.process_mesh().ndim();
+           ++mesh_axis) {
+        if (dist_attr.is_shard(mesh_axis, tensor_axis)) {
+          // handle the shard axis
+          int64_t global_shape = out_dim[tensor_axis];
+          int64_t mesh_size = dist_attr.process_mesh().dim_size(mesh_axis);
+          auto balance_shard = BalancedSplit(global_shape, mesh_size);
+          out_dim[tensor_axis] = balance_shard[coord_id[mesh_axis]];
+        }
+      }
+    }
+  }
+  return out_dim;
+}
+
+// 1. Get all the sub meshes of global_mesh
+// e.g. global_mesh = [[1, 2], [3, 4]], out_mesh = [1, 2] and [3, 4]
+//      global_mesh = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
+//      out_mesh = [[1, 2], [3, 4]] and [[5, 6], [7, 8]]
+std::vector<ProcessMesh> GetSubMeshes(const ProcessMesh& process_mesh) {
+  const std::vector<int64_t>& shape = process_mesh.shape();
+  const std::vector<int64_t>& process_ids = process_mesh.process_ids();
+  const std::vector<std::string>& dim_names = process_mesh.dim_names();
+  int64_t total_process_num = process_ids.size();
+  int64_t sub_process_num = total_process_num / shape[0];
+  std::vector<int64_t> sub_process_mesh_shape(shape.begin() + 1, shape.end());
+  std::vector<std::string> sub_process_mesh_dim_names(dim_names.begin() + 1,
+                                                      dim_names.end());
+
+  std::vector<ProcessMesh> sub_process_meshes;
+  for (int i = 0; i < shape[0]; ++i) {
+    int64_t start_position = i * sub_process_num;
+    int64_t end_position = start_position + sub_process_num;
+    std::vector<int64_t> sub_process_ids(process_ids.begin() + start_position,
+                                         process_ids.begin() + end_position);
+
+    sub_process_meshes.emplace_back(
+        sub_process_mesh_shape, sub_process_ids, sub_process_mesh_dim_names);
+  }
+  return sub_process_meshes;
+}
+
+bool IsSubMesh(const ProcessMesh& global_mesh, const ProcessMesh& sub_mesh) {
+  std::vector<ProcessMesh> sub_process_meshes = GetSubMeshes(global_mesh);
+  for (const ProcessMesh& mesh : sub_process_meshes) {
+    if (mesh == sub_mesh) {
+      return true;
+    }
+  }
+  return false;
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h

@@ -0,0 +1,231 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/tensor_base.h"
+#include "paddle/phi/core/visit_type.h"
+
+#if defined(PADDLE_WITH_XPU)
+#include <xpu/runtime.h>
+#endif
+
+namespace phi {
+class DeviceContext;
+
+namespace distributed {
+class ProcessMesh;
+
+PADDLE_API std::vector<int64_t> GetUnionProcessIds(
+    std::vector<int64_t> in_process_ids, std::vector<int64_t> out_process_ids);
+
+PADDLE_API bool IsCurRankInMesh(const ProcessMesh& process_mesh);
+
+PADDLE_API bool NeedComputationClipForPP(
+    const std::shared_ptr<phi::TensorBase>& tensor_impl);
+
+PADDLE_API Place GetDefaultPlace();
+
+PADDLE_API DeviceContext* GetDistTensorDeviceContext(DistTensor* input);
+
+PADDLE_API int64_t GetLocalRankInParticipate(
+    const std::vector<int64_t>& process_ids, int64_t global_rank = -1);
+
+// Get the coordinate of cur rank in process mesh. For example, the process mesh
+// is [[0, 1], [2, 3], [4, 5], [6, 7]], if the current rank is 4, then will
+// return [2, 0]; if the current rank is 3, then will return [1, 1].
+PADDLE_API std::vector<int64_t> GetCurRankCoordInMesh(
+    const ProcessMesh& process_mesh);
+
+// If the index i's value in dims_mapping is x ( x != -1), means the ith axis of
+// tensor need be split by xth axis of process_mesh. The function analyze the
+// input vector, return a key-value map of tensor_split_axis and
+// process_mesh_split_axis.
+// For example, if dims_mapping is [-1, 1, -1, 0], will return {1: 1, 3: 0}.
+PADDLE_API std::map<int, int64_t> GetSplitAxisWithDimsMapping(
+    const std::vector<int64_t>& dims_mapping);
+
+// If given a number, balance split it to multiple pieces.
+// For example, the input value is 12, split it to 5 pieces, then return
+// {3, 3, 2, 2, 2}.
+PADDLE_API std::vector<int64_t> BalancedSplit(int64_t total_nums,
+                                              int64_t num_of_pieces);
+
+// Create a comm context of the input process_ids. Once the newly comm context
+// created, it will be cached in the global instance, and get from the global
+// cache later. If the input dev_ctx is GPU, then nccl comm context will be
+// created. If the input dev_ctx is CPU, then gloo comm context will be created.
+PADDLE_API CommContext* CreateOrGetCommContext(
+    const DeviceContext& dev_ctx, const std::vector<int64_t>& process_ids);
+
+PADDLE_API DDim InferShapeForReshardFromReplicate(
+    const std::shared_ptr<DenseTensor>& global_value,
+    const TensorDistAttr& dist_attr);
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#define DEVICE_CONTEXT GPUContext
+#elif defined(PADDLE_WITH_XPU)
+#define DEVICE_CONTEXT XPUContext
+#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
+#define DEVICE_CONTEXT CustomContext
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+#define DEVICE_WAIT(dev_ctx) \
+  do {                       \
+    xpu_wait();              \
+    (dev_ctx)->Wait();       \
+  } while (0)
+#else
+#define DEVICE_WAIT(dev_ctx) \
+  do {                       \
+  } while (0)  // no need to wait on other devices.
+#endif
+
+// Some reshard function supports fewer data types on xpu than on gpu. For
+// example, `Transpose`, `Split`, and `Divide` do not support double type.
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#define PD_VISIT_RESHARD_TYPES PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES
+#else
+#define PD_VISIT_RESHARD_TYPES(TYPE, NAME, ...)                                \
+  [&] {                                                                        \
+    const auto& __dtype__ = TYPE;                                              \
+    switch (__dtype__) {                                                       \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT32, int, __VA_ARGS__)     \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT64, int64_t, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::FLOAT32, float, __VA_ARGS__) \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::FLOAT16, ::phi::float16, __VA_ARGS__)         \
+      PD_PRIVATE_CASE_TYPE_BFLOAT16(NAME, __VA_ARGS__)                         \
+      default:                                                                 \
+        PD_THROW("Reshard function " #NAME                                     \
+                 " is not implemented"                                         \
+                 " for data type `",                                           \
+                 __dtype__,                                                    \
+                 "`");                                                         \
+    }                                                                          \
+  }()
+#endif
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_XPU)
+#define RESHARD_FUNCTOR_IMPL(dev_ctx, fn_name, dtype, ...)                  \
+  do {                                                                      \
+    if (phi::CPUContext::classof(dev_ctx)) {                                \
+      VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU.";         \
+      PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES_CPU(                    \
+          dtype, #fn_name, ([&] {                                           \
+            fn_name<data_t>(static_cast<const CPUContext&>(*dev_ctx),       \
+                            __VA_ARGS__);                                   \
+          }));                                                              \
+    } else if (DEVICE_CONTEXT::classof(dev_ctx)) {                          \
+      DEVICE_WAIT(dev_ctx);                                                 \
+      VLOG(4) << "Call `" << #fn_name << "` in Resharding on device.";      \
+      PD_VISIT_RESHARD_TYPES(                                               \
+          dtype, #fn_name, ([&] {                                           \
+            fn_name<data_t>(static_cast<const DEVICE_CONTEXT&>(*dev_ctx),   \
+                            __VA_ARGS__);                                   \
+          }));                                                              \
+      DEVICE_WAIT(dev_ctx);                                                 \
+    } else {                                                                \
+      PADDLE_THROW(common::errors::Unimplemented(                           \
+          "The %s in reshard only supported on CPU, GPU, and XPU for now.", \
+          #fn_name));                                                       \
+    }                                                                       \
+  } while (0)
+#else
+#define RESHARD_FUNCTOR_IMPL(dev_ctx, fn_name, dtype, ...)                \
+  do {                                                                    \
+    if (phi::CPUContext::classof(dev_ctx)) {                              \
+      VLOG(4) << "Call `" << #fn_name << "` in Resharding on CPU.";       \
+      PD_VISIT_BOOL_AND_FLOATING_AND_INTEGRAL_TYPES_CPU(                  \
+          dtype, #fn_name, ([&] {                                         \
+            fn_name<data_t>(static_cast<const CPUContext&>(*dev_ctx),     \
+                            __VA_ARGS__);                                 \
+          }));                                                            \
+    } else {                                                              \
+      PADDLE_THROW(common::errors::Unimplemented(                         \
+          "The %s in reshard only supported on CPU for now.", #fn_name)); \
+    }                                                                     \
+  } while (0)
+#endif
+
+#define RESHARD_FUNCTOR_WITH_COMM(dev_ctx, fn_name, dtype, process_ids, ...) \
+  do {                                                                       \
+    auto* comm_context = CreateOrGetCommContext(*dev_ctx, process_ids);      \
+    dev_ctx->SetCommContext(comm_context);                                   \
+    RESHARD_FUNCTOR_IMPL(dev_ctx, fn_name, dtype, __VA_ARGS__);              \
+  } while (0)
+
+#define RESHARD_FUNCTOR(dev_ctx, fn_name, dtype, ...)           \
+  do {                                                          \
+    RESHARD_FUNCTOR_IMPL(dev_ctx, fn_name, dtype, __VA_ARGS__); \
+  } while (0)
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_XPU)
+#define RESHARD_FUNCTOR_WITHOUT_DTYPE(dev_ctx, fn_name, ...)              \
+  do {                                                                    \
+    if (phi::CPUContext::classof(dev_ctx)) {                              \
+      VLOG(4) << "Call `" << #fn_name                                     \
+              << "`without DType in  Resharding on CPU.";                 \
+      fn_name(static_cast<const CPUContext&>(*dev_ctx), __VA_ARGS__);     \
+    } else if (DEVICE_CONTEXT::classof(dev_ctx)) {                        \
+      VLOG(4) << "Call `" << #fn_name                                     \
+              << "`without DType in  Resharding on device.";              \
+      fn_name(static_cast<const DEVICE_CONTEXT&>(*dev_ctx), __VA_ARGS__); \
+    } else {                                                              \
+      PADDLE_THROW(common::errors::Unimplemented(                         \
+          "The %s in reshard only supported CPU, GPU, and XPU Device",    \
+          #fn_name));                                                     \
+    }                                                                     \
+  } while (0)
+#else
+#define RESHARD_FUNCTOR_WITHOUT_DTYPE(dev_ctx, fn_name, ...)            \
+  do {                                                                  \
+    if (phi::CPUContext::classof(dev_ctx)) {                            \
+      VLOG(4) << "Call `" << #fn_name                                   \
+              << "`without DType in Resharding on CPU.";                \
+      fn_name(static_cast<const CPUContext&>(*dev_ctx), __VA_ARGS__);   \
+    } else {                                                            \
+      PADDLE_THROW(common::errors::Unimplemented(                       \
+          "The %s in reshard only supported CPU, GPU, and XPU Device.", \
+          #fn_name));                                                   \
+    }                                                                   \
+  } while (0)
+#endif
+
+#define RESHARD_SHORTCUT_IF_FALSE(expr) \
+  do {                                  \
+    if (!(expr)) {                      \
+      return false;                     \
+    }                                   \
+  } while (0)
+
+std::vector<ProcessMesh> GetSubMeshes(const ProcessMesh& process_mesh);
+PADDLE_API bool IsSubMesh(const ProcessMesh& global_mesh,
+                          const ProcessMesh& sub_mesh);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.cc

@@ -0,0 +1,113 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/reduce_scatter_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi::distributed {
+
+bool SToPReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void SToPReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+
+  // step 1, create tmp dist attr and tmp dist tensor
+  TensorDistAttr tmp_attr(out_dist_attr);
+  DistTensor tmp_tensor;
+  tmp_attr.clean_partial_status();
+
+  // step 2, do s to r reshard on `in` to `tmp`
+  SToRReshardFunction s_to_r;
+  s_to_r.Eval(dev_ctx, in, tmp_attr, &tmp_tensor);
+
+  // step 3, do r to p reshard on `tmp` to `out`
+  RToPReshardFunction r_to_p;
+  r_to_p.Eval(dev_ctx, tmp_tensor, out_dist_attr, out);
+}
+
+bool SToPReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void SToPReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  DistTensor tmp_result;
+
+  SameStatusReshardFunction same_status_func;
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    SToPReshardFunction s_to_p_func;
+    PADDLE_ENFORCE(
+        s_to_p_func.IsSuitable(tmp_result, out_dist_attr),
+        common::errors::InvalidArgument(
+            "Invoke the s to p reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    s_to_p_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API SToPReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToPReshard"; }
+};
+
+class PADDLE_API SToPReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToPReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.cc

@@ -0,0 +1,257 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/all_gather_kernel.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/split_kernel.h"
+
+namespace phi::distributed {
+
+namespace {
+
+void ReshardSToRWithPadding(DeviceContext* dev_ctx,
+                            int64_t split_axis,
+                            const std::vector<int64_t>& process_ids,
+                            const DenseTensor& in,
+                            int64_t padding_nums,
+                            DenseTensor* out,
+                            int64_t split_factor = 1) {
+  VLOG(4) << "split factor is " << split_factor;
+  int64_t num_of_process = process_ids.size();
+  auto dtype = in.dtype();
+
+  // For balanced split to replicate, we need to do all gather first.
+  // If the input value doesn't split on axis 0, we need to split
+  // and concat on specific axis.
+  RESHARD_FUNCTOR_WITH_COMM(
+      dev_ctx, AllGather, dtype, process_ids, in, num_of_process, out);
+
+  if (split_axis != 0 || padding_nums != 0 || split_factor > 1) {
+    PADDLE_ENFORCE_EQ((in.dims()[0] % split_factor),
+                      0,
+                      common::errors::InvalidArgument(
+                          "The specified axis of tensor should be evenly "
+                          "splited, but got %d and %d.",
+                          in.dims()[0],
+                          split_factor));
+    IntArray sections(std::vector<int64_t>(num_of_process * split_factor,
+                                           in.dims()[0] / split_factor));
+
+    std::vector<DenseTensor> split_out_vec;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Split,
+                    dtype,
+                    *out,
+                    sections,
+                    /*split_axis*/ 0,
+                    &split_out_vec);
+
+    if (padding_nums != 0) {
+      std::vector<DenseTensor> tmp_out_vec;
+      IntArray tmp_sections(std::vector<int64_t>{
+          in.dims()[split_axis] - padding_nums, padding_nums});
+      RESHARD_FUNCTOR(dev_ctx,
+                      Split,
+                      dtype,
+                      split_out_vec[num_of_process - 1],
+                      tmp_sections,
+                      split_axis,
+                      &tmp_out_vec);
+      split_out_vec[num_of_process - 1] = tmp_out_vec[0];
+    }
+
+    // Concat the result after split on correct axis.
+    std::vector<const DenseTensor*> concat_input_vec;
+    concat_input_vec.reserve(split_out_vec.size());
+
+    for (size_t idx = 0; idx < split_out_vec.size(); idx++) {
+      int64_t new_idx =
+          idx % num_of_process * split_factor + idx / num_of_process;
+      concat_input_vec.emplace_back(&(split_out_vec.at(new_idx)));
+    }
+
+    RESHARD_FUNCTOR(dev_ctx, Concat, dtype, concat_input_vec, split_axis, out);
+  }
+}
+
+std::map<int, int64_t> GetSplitAxisWithDimsMapping(
+    const std::vector<std::vector<int64_t>>& dims_mapping) {
+  std::map<int, int64_t> split_axis_to_mesh_axis;
+  for (size_t i = 0; i < dims_mapping.size(); ++i) {
+    if (dims_mapping[i].size() > 0) {
+      split_axis_to_mesh_axis.emplace(i, dims_mapping[i][0]);
+    }
+  }
+  return split_axis_to_mesh_axis;
+}
+}  // namespace
+
+bool SToRReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void SToRReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& in_dims_mapping = in_dist_attr.multi_dims_mapping();
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& in_process_ids = in_process_mesh.process_ids();
+  if (in_process_ids.size() == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+
+  int split_axis = GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
+  int64_t mesh_axis =
+      GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->second;
+  int64_t num_of_process = in_process_mesh.size();
+  int64_t num_of_padding = in.dims()[split_axis] % num_of_process;
+  VLOG(4) << "dist attr " << in_dist_attr.to_string();
+  int64_t split_factor = in_dist_attr.get_split_factor(mesh_axis);
+  bool is_balanced_split = (num_of_padding == 0);
+
+  if (is_balanced_split) {
+    VLOG(3) << "Balanced reshard from shard to replicated";
+    ReshardSToRWithPadding(dev_ctx,
+                           split_axis,
+                           in_process_ids,
+                           in.value(),
+                           /*padding_nums*/ 0,
+                           GetMutableTensor(out),
+                           split_factor);
+  } else {
+    VLOG(3) << "Unbalanced reshard from shard to replicated";
+    int64_t avg_size_on_split_axis =
+        (in.dims()[split_axis] + num_of_process - 1) / num_of_process;
+    int64_t padding_nums =
+        avg_size_on_split_axis * num_of_process - in.dims()[split_axis];
+    bool need_padding = (in.local_dims()[split_axis] != avg_size_on_split_axis);
+
+    if (need_padding) {
+      DDim concat_local_shape = in.local_dims();
+      concat_local_shape[split_axis] = padding_nums;
+      IntArray concat_local_shape_int_array(concat_local_shape.Get(),
+                                            concat_local_shape.size());
+      auto dtype = in.dtype();
+
+      DenseTensor concat_local_tensor;
+      RESHARD_FUNCTOR(dev_ctx,
+                      Full,
+                      dtype,
+                      concat_local_shape_int_array,
+                      0,
+                      &concat_local_tensor);
+
+      DenseTensor in_local_tensor = in.value();
+      std::vector<const DenseTensor*> concat_input_vec = {&in_local_tensor,
+                                                          &concat_local_tensor};
+
+      DenseTensor concat_result;
+      RESHARD_FUNCTOR(
+          dev_ctx, Concat, dtype, concat_input_vec, split_axis, &concat_result);
+      ReshardSToRWithPadding(dev_ctx,
+                             split_axis,
+                             in_process_ids,
+                             concat_result,
+                             padding_nums,
+                             GetMutableTensor(out));
+    } else {
+      ReshardSToRWithPadding(dev_ctx,
+                             split_axis,
+                             in_process_ids,
+                             in.value(),
+                             padding_nums,
+                             GetMutableTensor(out));
+    }
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool SToRReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void SToRReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  SameStatusReshardFunction same_status_func;
+  DistTensor tmp_result;
+
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    SToRReshardFunction s_to_r_func;
+    PADDLE_ENFORCE(
+        s_to_r_func.IsSuitable(tmp_result, out_dist_attr),
+        common::errors::InvalidArgument(
+            "Invoke the s to r reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    s_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h

@@ -0,0 +1,52 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API SToRReshardFunction final : public ReshardFunction {
+ public:
+  SToRReshardFunction() = default;
+  ~SToRReshardFunction() = default;
+
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToRReshard"; }
+};
+
+class PADDLE_API SToRReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToRReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.cc

@@ -0,0 +1,198 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/all_to_all_kernel.h"
+#include "paddle/phi/kernels/reshape_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi::distributed {
+
+bool SToSReshardFunction::IsSuitable(const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.multi_dims_mapping() !=
+                            out_dist_attr.multi_dims_mapping());
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh == out_process_mesh);
+
+  return true;
+}
+
+void SToSReshardFunction::Eval(DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_process_mesh = in.dist_attr().process_mesh();
+  const auto& in_process_ids = in_process_mesh.process_ids();
+  if (in_process_ids.size() == 1) {
+    SetValue(out, in.value());
+    SetDistProps(out, in.dims(), out_dist_attr);
+    return;
+  }
+  auto dtype = in.dtype();
+  const auto& logical_ddim = in.dims();
+  int64_t nranks = static_cast<int64_t>(in_process_ids.size());
+  int in_split_axis =
+      GetSplitAxisWithDimsMapping(in.dist_attr().dims_mapping()).begin()->first;
+  int out_split_axis =
+      GetSplitAxisWithDimsMapping(out_dist_attr.dims_mapping()).begin()->first;
+
+  DenseTensor in_all_to_all;
+  // 1. preprocess, reshape and transpose the input tensor
+  if (out_split_axis != 0) {
+    // 1.1 calc the shape and reshape
+    std::vector<int64_t> pre_shape_vec = common::vectorize(logical_ddim);
+    pre_shape_vec[in_split_axis] /= nranks;
+    pre_shape_vec[out_split_axis] /= nranks;
+    pre_shape_vec.insert(pre_shape_vec.begin() + out_split_axis, nranks);
+
+    DenseTensor out_reshape1;
+    RESHARD_FUNCTOR(
+        dev_ctx, Reshape, dtype, in.value(), pre_shape_vec, &out_reshape1);
+
+    // 1.2 calc the the desire axis and transpose
+    std::vector<int> axis;
+    axis.emplace_back(out_split_axis);
+    for (size_t i = 0; i < pre_shape_vec.size(); ++i) {
+      if (static_cast<int>(i) != out_split_axis) {
+        axis.emplace_back(i);
+      }
+    }
+    DenseTensor out_transpose;
+    RESHARD_FUNCTOR(
+        dev_ctx, Transpose, dtype, out_reshape1, axis, &out_transpose);
+
+    // 1.3 calc the final shape and reshape
+    pre_shape_vec.erase(pre_shape_vec.begin() + out_split_axis);
+    pre_shape_vec[in_split_axis] *= nranks;
+    RESHARD_FUNCTOR(
+        dev_ctx, Reshape, dtype, out_transpose, pre_shape_vec, &in_all_to_all);
+  } else {
+    in_all_to_all.ShareDataNoCheckWith(in.value());
+  }
+
+  // 2. use all to all to switch data to other ranks
+  RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                            AllToAll,
+                            dtype,
+                            in_process_ids,
+                            in_all_to_all,
+                            GetMutableTensor(out));
+
+  // 3. postprocess, reshape and transpose the output tensor
+  if (in_split_axis != 0) {
+    // 3.1 calc the shape and reshape
+    std::vector<int64_t> post_shape_vec = common::vectorize(logical_ddim);
+    post_shape_vec[in_split_axis] /= nranks;
+    post_shape_vec[out_split_axis] /= nranks;
+    post_shape_vec.insert(post_shape_vec.begin(), nranks);
+
+    DenseTensor out_reshape1;
+    RESHARD_FUNCTOR(
+        dev_ctx, Reshape, dtype, out->value(), post_shape_vec, &out_reshape1);
+
+    // 3.2 calc the the desire axis and transpose
+    std::vector<int> axis;
+    for (size_t i = 1; i < post_shape_vec.size(); ++i) {
+      axis.emplace_back(i);
+    }
+    axis.insert(axis.begin() + in_split_axis, 0);
+    DenseTensor out_transpose;
+    RESHARD_FUNCTOR(
+        dev_ctx, Transpose, dtype, out_reshape1, axis, &out_transpose);
+
+    // 3.3 calc the final shape and reshape
+    post_shape_vec.erase(post_shape_vec.begin());
+    post_shape_vec[in_split_axis] *= nranks;
+    RESHARD_FUNCTOR(dev_ctx,
+                    Reshape,
+                    dtype,
+                    out_transpose,
+                    post_shape_vec,
+                    GetMutableTensor(out));
+  }
+
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+bool SToSReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.multi_dims_mapping() !=
+                            out_dist_attr.multi_dims_mapping());
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void SToSReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  SameStatusReshardFunction same_status_func;
+  DistTensor tmp_result;
+
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    SToSReshardFunction s_to_s_func;
+    PADDLE_ENFORCE(
+        s_to_s_func.IsSuitable(tmp_result, out_dist_attr),
+        common::errors::InvalidArgument(
+            "Invoke the s to s reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    s_to_s_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h

@@ -0,0 +1,55 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API SToSReshardFunction final : public ReshardFunction {
+ public:
+  SToSReshardFunction() = default;
+  ~SToSReshardFunction() = default;
+
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToSReshard"; }
+};
+
+class PADDLE_API SToSReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  SToSReshardFunctionCrossMesh() = default;
+  ~SToSReshardFunctionCrossMesh() = default;
+
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToSReshardCrossMesh"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.cc

@@ -0,0 +1,114 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+
+#include <algorithm>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/p_recv_kernel.h"
+#include "paddle/phi/kernels/p_send_kernel.h"
+
+namespace phi::distributed {
+
+bool SameStatusReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.multi_dims_mapping() ==
+                            out_dist_attr.multi_dims_mapping());
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.partial_dims() ==
+                            out_dist_attr.partial_dims());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.shape() ==
+                            out_process_mesh.shape());
+
+  return true;
+}
+
+void SameStatusReshardFunction::Eval(DeviceContext* dev_ctx,
+                                     const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr,
+                                     DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& in_process_ids = in_process_mesh.process_ids();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+  const auto& out_process_ids = out_process_mesh.process_ids();
+  auto all_process_ids = GetUnionProcessIds(in_process_ids, out_process_ids);
+
+  // TODO(GhostScreaming): After cross-mesh reshard, current device may
+  // needs to execute next layer. When it construct next layer's backward
+  // graph, out->place() will be called such as in SetGradOutMeta method. As
+  // a result, out can't be undefined. Try to allocate a zero-memory value
+  // for out. Following send/recv will cover this empty DenseTensor
+  // construction.
+  VLOG(3) << "Same_status_reshard_function create an empty DenseTensor for "
+             "cross-mesh DistTensor.";
+  *(out->unsafe_mutable_value()) =
+      DenseTensor(std::make_shared<phi::Allocation>(
+                      nullptr, 0, phi::distributed::GetDefaultPlace()),
+                  in.value().meta());
+
+  std::vector<std::pair<int64_t, int64_t>> p2p_pair;
+  for (size_t i = 0; i < out_process_ids.size(); ++i) {
+    p2p_pair.emplace_back(in_process_ids[i], out_process_ids[i]);
+  }
+  int64_t cur_global_rank = GetCurGlobalRank();
+  for (const auto& iter : p2p_pair) {
+    int64_t src = iter.first;
+    int64_t dst = iter.second;
+    if (src == cur_global_rank) {
+      VLOG(3) << "Send from src " << src << " to dst " << dst;
+      int64_t dst_local_rank = GetLocalRankInParticipate(all_process_ids, dst);
+      // Since send kernel only has input, so we don't need to infermeta
+      // actually. According to this reason, just use the kernel directly.
+      RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                PSendKernel,
+                                in.dtype(),
+                                all_process_ids,
+                                in.value(),
+                                dst_local_rank,
+                                /*dynamic_shape=*/true);
+      // TODO(liyurui): Use dynamic shape will lead to poor performance, but we
+      // don't have any other good idea now. For the following reasons:
+      // 1. We can not ensure the meta being right deduce by the infermeta.
+      // 2. The meta of some kernels can't decide in compile time.
+      // 3. DenseTensor with empty value only need infermeta and skip the real
+      // kernel execution.
+    } else if (dst == cur_global_rank) {
+      VLOG(3) << "Recv from src " << src << " to dst " << dst;
+      int64_t src_local_rank = GetLocalRankInParticipate(all_process_ids, src);
+      RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                PRecv,
+                                in.dtype(),
+                                all_process_ids,
+                                src_local_rank,
+                                {} /*out_shape*/,
+                                /*dynamic_shape=*/true,
+                                GetMutableTensor(out));
+    }
+  }
+  SetDistProps(out, in.dims(), out_dist_attr);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h

@@ -0,0 +1,36 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API SameStatusReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SameStatusReshard"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/x_to_r_reshard_function.cc

@@ -0,0 +1,132 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/x_to_r_reshard_function.h"
+
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
+#include "paddle/phi/kernels/add_n_kernel.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/p_recv_kernel.h"
+#include "paddle/phi/kernels/p_send_kernel.h"
+
+namespace phi::distributed {
+
+bool XToRShrinkReshardFunction::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_replicated());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.process_ids().size() != 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.process_ids().size() == 1);
+
+  return true;
+}
+
+void XToRShrinkReshardFunction::Eval(DeviceContext* dev_ctx,
+                                     const DistTensor& in,
+                                     const TensorDistAttr& out_dist_attr,
+                                     DistTensor* out) {
+  VLOG(3) << "Call " << Name();
+  const auto& in_dist_attr = in.dist_attr();
+  const auto& in_dims_mapping = in_dist_attr.dims_mapping();
+  const auto& in_mesh = in_dist_attr.process_mesh();
+  const auto& out_mesh = out_dist_attr.process_mesh();
+  const auto& in_process_ids = in_mesh.process_ids();
+  const auto& out_process_ids = out_mesh.process_ids();
+  int64_t cur_global_rank = GetCurGlobalRank();
+  int64_t root_rank = out_process_ids[0];
+  auto dtype = in.dtype();
+  const auto& in_partial_status = in_dist_attr.partial_status();
+  auto all_process_ids = GetUnionProcessIds(in_process_ids, out_process_ids);
+  std::unordered_map<int64_t, DenseTensor> rank_to_result;
+
+  // Step 1: other ranks need to send value to the root
+  if (!in_dist_attr.is_replicated()) {
+    if (cur_global_rank != root_rank) {
+      // send dense tensor to root
+      RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                PSendKernel,
+                                dtype,
+                                all_process_ids,
+                                in.value(),
+                                root_rank,
+                                /*dynamic_shape=*/true);
+    } else {
+      for (size_t i = 0; i < all_process_ids.size(); ++i) {
+        if (all_process_ids[i] != root_rank) {
+          rank_to_result.emplace(all_process_ids[i], DenseTensor());
+          RESHARD_FUNCTOR_WITH_COMM(dev_ctx,
+                                    PRecv,
+                                    dtype,
+                                    all_process_ids,
+                                    all_process_ids[i],
+                                    {} /*out_shape*/,
+                                    /*dynamic_shape=*/true,
+                                    &rank_to_result[all_process_ids[i]]);
+        }
+      }
+    }
+  }
+
+  // Step 2: concat or reduce based on dist attr
+  if (cur_global_rank == root_rank) {
+    std::vector<const DenseTensor*> input_vec;
+    for (const auto& in_process_id : in_process_ids) {
+      if (in_process_id == cur_global_rank) {
+        input_vec.emplace_back(&(in.value()));
+      } else {
+        input_vec.emplace_back(&(rank_to_result[in_process_id]));
+      }
+    }
+    if (in_dist_attr.is_shard()) {
+      int split_axis =
+          GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
+      RESHARD_FUNCTOR(
+          dev_ctx, Concat, dtype, input_vec, split_axis, GetMutableTensor(out));
+    } else if (in_dist_attr.is_partial()) {
+      auto in_reduce_type = in_partial_status.at(0);
+      if (in_reduce_type == ReduceType::kRedSum) {
+        DenseTensor result_add_out = *input_vec[0];
+        for (size_t i = 1; i < input_vec.size(); ++i) {
+          RESHARD_FUNCTOR(dev_ctx,
+                          Add,
+                          dtype,
+                          *input_vec[i],
+                          result_add_out,
+                          &result_add_out);
+        }
+        SetValue(out, result_add_out);
+      } else {
+        PADDLE_THROW(common::errors::Unavailable(
+            "The reduce type is not supported, will be supported soon."));
+      }
+    } else {
+      SetValue(out, in.value());
+    }
+    SetDistProps(out, in.dims(), out_dist_attr);
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/auto_parallel/reshard/x_to_r_reshard_function.h

@@ -0,0 +1,36 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API XToRShrinkReshardFunction final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "XToRShrinkReshard"; }
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/auto_parallel/utils.h

@@ -0,0 +1,191 @@
+/* Copyright (c) 2022 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. */
+#pragma once
+
+#include <algorithm>
+#include <iterator>
+#include <map>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace distributed {
+namespace auto_parallel {
+
+// struct Indent {
+//   Indent(int &level) : level(level) { ++level; }
+//   ~Indent() { --level; }
+//   int &level;
+// };
+
+// inline std::string str_indent(std::string& str, cur_indent) {
+//   string spaces(cur_indent, " ");
+//   return str + std::string(cur_indent, " ");
+// }
+
+template <class T>
+bool has_duplicates(const std::vector<T>& vec) {
+  std::unordered_map<T, int> map;
+  for (const auto& i : vec) {
+    ++map[i];
+    if (map[i] > 1) return true;
+  }
+  return false;
+}
+
+inline int64_t canonical_dim(int dim, int ndim) {
+  PADDLE_ENFORCE_EQ(
+      dim >= -ndim && dim < ndim,
+      true,
+      errors::InvalidArgument(
+          "Dimension %d is outside of [-%d, %d).", dim, ndim, ndim));
+  if (dim < 0) {
+    return dim + ndim;
+  }
+  return dim;
+}
+
+// Refer to https://stackoverflow.com/a/5289170
+template <typename Range, typename Value = typename Range::value_type>
+std::string str_join(Range const& elements,
+                     const std::string& delimiter = ",") {
+  std::ostringstream os;
+  auto b = std::begin(elements), e = std::end(elements);
+
+  if (b != e) {
+    std::copy(b, prev(e), std::ostream_iterator<Value>(os, delimiter.c_str()));
+    b = prev(e);
+  }
+  if (b != e) {
+    os << *b;
+  }
+
+  return os.str();
+}
+
+inline std::string str_join(std::map<std::string, bool> const& elements,
+                            const std::string& delimiter = ",") {
+  std::string str;
+  for (const auto& item : elements) {
+    str += item.first + ": " + std::to_string(item.second) + delimiter;
+  }
+  return str.substr(0, str.size() - 1);
+}
+
+inline std::string str_join(const std::vector<std::vector<int64_t>>& elements) {
+  std::stringstream ss;
+  for (const auto& e : elements) {
+    ss << "[" << str_join(e) << "] ";
+  }
+  return ss.str();
+}
+
+inline std::string str_join(const std::unordered_map<int64_t, int64_t>& map) {
+  std::stringstream ss;
+  for (const auto& [k, v] : map) {
+    ss << "mesh dim: " << std::to_string(k)
+       << ", split factor: " << std::to_string(v);
+  }
+  return ss.str();
+}
+
+// Refer to https://stackoverflow.com/a/46931770
+inline std::vector<std::string> str_split(std::string const& input,
+                                          const std::string& delimiter = ",") {
+  size_t pos_start = 0, pos_end, delim_len = delimiter.length();
+  std::string token;
+  std::vector<std::string> output;
+  while ((pos_end = input.find(delimiter, pos_start)) != std::string::npos) {
+    token = input.substr(pos_start, pos_end - pos_start);
+    pos_start = pos_end + delim_len;
+    output.push_back(token);
+  }
+  output.push_back(input.substr(pos_start));
+  return output;
+}
+
+// Refer to https://stackoverflow.com/a/29200671/2358969
+template <typename T>
+std::string to_string_with_precision(const T a_value, const int n = 2) {
+  std::ostringstream out;
+  out.precision(n);
+  out << std::fixed << a_value;
+  return out.str();
+}
+
+class SplitFactor final {
+ public:
+  SplitFactor() {}
+  SplitFactor(std::unordered_map<int64_t, int64_t> split_factor_map)
+      : split_factor_map_(split_factor_map) {}
+
+  void set_split_factor(int64_t mesh_dim, int64_t split_factor) {
+    // default value is 1
+    if (split_factor > 1) {
+      split_factor_map_[mesh_dim] = split_factor;
+    }
+    PADDLE_ENFORCE_LE(split_factor_map_.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "At now only support to rearrange at one mesh dim."));
+  }
+
+  int64_t get_split_factor(int64_t mesh_dim) const {
+    return split_factor_map_.count(mesh_dim) ? split_factor_map_.at(mesh_dim)
+                                             : 1;
+  }
+
+  void clear_split_factor(int64_t mesh_dim) {
+    if (split_factor_map_.count(mesh_dim)) {
+      split_factor_map_.erase(mesh_dim);
+    }
+  }
+  bool operator==(const SplitFactor& other) const {
+    return split_factor_map_ == other.split_factor_map_;
+  }
+
+  bool operator!=(const SplitFactor& other) const {
+    return !(this->operator==(other));
+  }
+
+  std::string to_string() const {
+    std::stringstream ss;
+    for (const auto& [k, v] : split_factor_map_) {
+      ss << "mesh dim: " << std::to_string(k)
+         << ", split factor: " << std::to_string(v);
+    }
+    return ss.str();
+  }
+
+ private:
+  std::unordered_map<int64_t, int64_t> split_factor_map_;
+};
+
+}  // namespace auto_parallel
+}  // namespace distributed
+}  // namespace phi
+namespace std {
+template <>
+struct hash<phi::distributed::auto_parallel::SplitFactor> {
+  size_t operator()(
+      const phi::distributed::auto_parallel::SplitFactor& split_factor) const {
+    string str = split_factor.to_string();
+    return hash<string>()(str);
+  }
+};
+}  // namespace std

paddle/phi/core/distributed/bkcl_comm_context.cc

@@ -0,0 +1,396 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/bkcl_comm_context.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/distributed/check/static_check.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi {
+namespace distributed {
+
+BKCLCommContext::BKCLCommContext(int rank, int size, BKCLUniqueId bkcl_id)
+    : CommContext(rank, size) {
+  PADDLE_ENFORCE_BKCL_SUCCESS(
+      bkcl_init_rank(&bkcl_comm_, rank_, size_, &bkcl_id));
+}
+
+#if defined(PADDLE_WITH_FLAGCX)
+BKCLCommContext::BKCLCommContext(int rank,
+                                 int size,
+                                 flagcxHandlerGroup_t flagcx_handler)
+    : CommContext(rank, size), flagcx_handler_(flagcx_handler) {
+  phi::dynload::flagcxCommInitRank(
+      &flagcx_handler_->comm, size_, flagcx_handler_->uniqueId, rank_);
+}
+#endif
+
+BKCLContext_t BKCLCommContext::GetBKCLComm() { return bkcl_comm_; }
+
+XPUStream BKCLCommContext::GetStream() { return dev_ctx_->stream(); }
+
+phi::XPUContext* BKCLCommContext::GetDevContext() { return dev_ctx_.get(); }
+
+void BKCLCommContext::SetDevContext(
+    std::unique_ptr<phi::XPUContext>&& dev_ctx) {
+  dev_ctx_ = std::move(dev_ctx);
+}
+
+XPUEvent BKCLCommContext::GetComputeEvent() { return compute_event_.get(); }
+
+void BKCLCommContext::SetComputeEvent(
+    std::shared_ptr<std::remove_pointer<XPUEvent>::type>&& compute_event) {
+  compute_event_ = std::move(compute_event);
+}
+
+XPUEvent BKCLCommContext::GetCommEvent() { return comm_event_.get(); }
+
+void BKCLCommContext::SetCommEvent(
+    std::shared_ptr<std::remove_pointer<XPUEvent>::type>&& comm_event) {
+  comm_event_ = std::move(comm_event);
+}
+
+void BKCLCommContext::Broadcast(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                int root,
+                                XPUStream stream) {
+  CommStaticCheck::SameShape(*out_tensor,
+                             in_tensor,
+                             /*dst_rank*/ rank_,
+                             /*cur_rank*/ rank_,
+                             size_,
+                             AllocationType::XPU);
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxBroadcast(in_tensor.data(),
+                                    out_tensor->data(),
+                                    in_tensor.numel(),
+                                    ToFlagcxDataType(in_tensor.type()),
+                                    root,
+                                    flagcx_handler_->comm,
+                                    reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_broadcast(bkcl_comm_,
+                                             in_tensor.data(),
+                                             out_tensor->data(),
+                                             in_tensor.numel(),
+                                             ToBKCLDataType(in_tensor.type()),
+                                             root,
+                                             stream));
+#endif
+}
+
+void BKCLCommContext::AllGather(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                XPUStream stream) {
+  phi::distributed::CommStaticCheck::GatherLikeShape(*out_tensor,
+                                                     in_tensor,
+                                                     /*dst_rank*/ rank_,
+                                                     /*cur_rank*/ rank_,
+                                                     size_,
+                                                     AllocationType::XPU);
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxAllGather(in_tensor.data(),
+                                    out_tensor->data(),
+                                    in_tensor.numel(),
+                                    ToFlagcxDataType(in_tensor.type()),
+                                    flagcx_handler_->comm,
+                                    reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_all_gather(bkcl_comm_,
+                                              in_tensor.data(),
+                                              in_tensor.numel(),
+                                              out_tensor->data(),
+                                              ToBKCLDataType(in_tensor.type()),
+                                              stream));
+#endif
+}
+
+void BKCLCommContext::ReduceScatter(DenseTensor* out_tensor,
+                                    const DenseTensor& in_tensor,
+                                    BKCLOp reduce_type,
+                                    XPUStream stream) {
+  phi::distributed::CommStaticCheck::ScatterLikeShape(*out_tensor,
+                                                      in_tensor,
+                                                      /*dst_rank*/ rank_,
+                                                      /*cur_rank*/ rank_,
+                                                      size_,
+                                                      AllocationType::XPU);
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(phi::dynload::flagcxReduceScatter(
+      in_tensor.data(),
+      out_tensor->data(),
+      out_tensor->numel(),
+      ToFlagcxDataType(in_tensor.type()),
+      BkclToFlagcxRedType(reduce_type),
+      flagcx_handler_->comm,
+      reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(
+      bkcl_reduce_scatter(bkcl_comm_,
+                          in_tensor.data(),
+                          out_tensor->data(),
+                          out_tensor->numel(),
+                          ToBKCLDataType(in_tensor.type()),
+                          reduce_type,
+                          stream));
+#endif
+}
+
+#if defined(PADDLE_WITH_FLAGCX)
+void BKCLCommContext::Scatter(DenseTensor* out_tensor,
+                              const DenseTensor& in_tensor,
+                              int root,
+                              XPUStream stream) {
+  phi::distributed::CommStaticCheck::ScatterLikeShape(*out_tensor,
+                                                      in_tensor,
+                                                      /*dst_rank*/ rank_,
+                                                      /*cur_rank*/ rank_,
+                                                      size_,
+                                                      AllocationType::XPU);
+
+  FLAGCX_CHECK(
+      phi::dynload::flagcxScatter(in_tensor.data(),
+                                  out_tensor->data(),
+                                  out_tensor->numel(),
+                                  ToFlagcxDataType(in_tensor.type()),
+                                  root,
+                                  flagcx_handler_->comm,
+                                  reinterpret_cast<flagcxStream_t>(&stream)));
+}
+#endif
+
+void BKCLCommContext::Send(const DenseTensor& in_tensor,
+                           const int64_t& count,
+                           const int& peer,
+                           XPUStream stream) {
+  phi::distributed::CommStaticCheck::CheckShape(
+      in_tensor, rank_, size_, AllocationType::XPU);
+
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxSend(in_tensor.data(),
+                               count,
+                               ToFlagcxDataType(in_tensor.dtype()),
+                               peer,
+                               flagcx_handler_->comm,
+                               reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_send(bkcl_comm_,
+                                        in_tensor.data(),
+                                        count,
+                                        peer,
+                                        ToBKCLDataType(in_tensor.dtype()),
+                                        stream));
+#endif
+  VLOG(3) << "rank " << GetRank() << " send " << phi::product(in_tensor.dims())
+          << " to " << peer;
+}
+
+void BKCLCommContext::Recv(DenseTensor* out_tensor,
+                           const int64_t& count,
+                           const int& peer,
+                           XPUStream stream) {
+  phi::distributed::CommStaticCheck::CheckShape(
+      *out_tensor, rank_, size_, AllocationType::XPU);
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxRecv(out_tensor->data(),
+                               count,
+                               ToFlagcxDataType(out_tensor->dtype()),
+                               peer,
+                               flagcx_handler_->comm,
+                               reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_recv(bkcl_comm_,
+                                        out_tensor->data(),
+                                        count,
+                                        peer,
+                                        ToBKCLDataType(out_tensor->dtype()),
+                                        stream));
+#endif
+  VLOG(3) << "rank " << GetRank() << " recv "
+          << common::product(out_tensor->dims()) << " from " << peer;
+}
+
+void BKCLCommContext::AllReduce(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                BKCLOp reduce_type,
+                                XPUStream stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ rank_,
+                                               /*cur_rank*/ rank_,
+                                               size_,
+                                               AllocationType::XPU);
+
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxAllReduce(in_tensor.data(),
+                                    out_tensor->data(),
+                                    in_tensor.numel(),
+                                    ToFlagcxDataType(in_tensor.type()),
+                                    BkclToFlagcxRedType(reduce_type),
+                                    flagcx_handler_->comm,
+                                    reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_all_reduce(bkcl_comm_,
+                                              in_tensor.data(),
+                                              out_tensor->data(),
+                                              in_tensor.numel(),
+                                              ToBKCLDataType(in_tensor.type()),
+                                              reduce_type,
+                                              stream));
+#endif
+}
+
+void BKCLCommContext::AllToAll(DenseTensor* out_tensor,
+                               const DenseTensor& in_tensor,
+                               XPUStream stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ rank_,
+                                               /*cur_rank*/ rank_,
+                                               size_,
+                                               AllocationType::XPU);
+
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxAlltoAll(in_tensor.data(),
+                                   out_tensor->data(),
+                                   in_tensor.numel() / size_,
+                                   ToFlagcxDataType(in_tensor.type()),
+                                   flagcx_handler_->comm,
+                                   reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_all_to_all(bkcl_comm_,
+                                              in_tensor.data(),
+                                              in_tensor.numel() / size_,
+                                              out_tensor->data(),
+                                              ToBKCLDataType(in_tensor.type()),
+                                              stream));
+#endif
+}
+
+void BKCLCommContext::AllToAllUnequalSplit(DenseTensor* out_tensor,
+                                           const DenseTensor& in_tensor,
+                                           const DenseTensor& out_size_tensor,
+                                           const DenseTensor& out_offset_tensor,
+                                           const DenseTensor& in_size_tensor,
+                                           const DenseTensor& in_offset_tensor,
+                                           XPUStream stream) {
+  auto in_size_ptr = reinterpret_cast<const size_t*>(in_size_tensor.data());
+  auto in_offset_ptr = reinterpret_cast<const size_t*>(in_offset_tensor.data());
+  auto out_size_ptr = reinterpret_cast<const size_t*>(out_size_tensor.data());
+  auto out_offset_ptr =
+      reinterpret_cast<const size_t*>(out_offset_tensor.data());
+
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxAlltoAllv(in_tensor.data(),
+                                    const_cast<size_t*>(in_size_ptr),
+                                    const_cast<size_t*>(in_offset_ptr),
+                                    out_tensor->data(),
+                                    const_cast<size_t*>(out_size_ptr),
+                                    const_cast<size_t*>(out_offset_ptr),
+                                    ToFlagcxDataType(in_tensor.type()),
+                                    flagcx_handler_->comm,
+                                    reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+
+  PADDLE_ENFORCE_BKCL_SUCCESS(
+      bkcl_all_to_all_v(bkcl_comm_,
+                        in_tensor.data(),
+                        in_size_ptr,
+                        in_offset_ptr,
+                        ToBKCLDataType(in_tensor.type()),
+                        out_tensor->data(),
+                        out_size_ptr,
+                        out_offset_ptr,
+                        ToBKCLDataType(out_tensor->type()),
+                        stream));
+#endif
+}
+
+void BKCLCommContext::Reduce(DenseTensor* out_tensor,
+                             const DenseTensor& in_tensor,
+                             BKCLOp reduce_type,
+                             int root,
+                             XPUStream stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ root,
+                                               /*cur_rank*/ rank_,
+                                               size_,
+                                               AllocationType::XPU);
+
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(
+      phi::dynload::flagcxReduce(in_tensor.data(),
+                                 out_tensor->data(),
+                                 in_tensor.numel(),
+                                 ToFlagcxDataType(in_tensor.type()),
+                                 BkclToFlagcxRedType(reduce_type),
+                                 root,
+                                 flagcx_handler_->comm,
+                                 reinterpret_cast<flagcxStream_t>(&stream)));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_reduce(bkcl_comm_,
+                                          in_tensor.data(),
+                                          out_tensor->data(),
+                                          in_tensor.numel(),
+                                          ToBKCLDataType(in_tensor.type()),
+                                          reduce_type,
+                                          root,
+                                          stream));
+#endif
+}
+
+void BKCLCommContext::GroupStart() {
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(phi::dynload::flagcxGroupStart(flagcx_handler_->comm));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_group_start());
+#endif
+}
+void BKCLCommContext::GroupEnd() {
+#if defined(PADDLE_WITH_FLAGCX)
+  FLAGCX_CHECK(phi::dynload::flagcxGroupEnd(flagcx_handler_->comm));
+#else
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_group_end());
+#endif
+}
+
+#if defined(PADDLE_WITH_FLAGCX)
+flagcxRedOp_t BKCLCommContext::BkclToFlagcxRedType(BKCLOp redOp) {
+  switch (redOp) {
+    case BKCL_MIN:
+      return flagcxMin;
+    case BKCL_MAX:
+      return flagcxMax;
+    case BKCL_ADD:
+      return flagcxSum;
+  }
+}
+#endif
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/bkcl_comm_context.h

@@ -0,0 +1,138 @@
+// Copyright (c) 2023 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.
+#pragma once
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+
+#if defined(PADDLE_WITH_FLAGCX)
+#include "paddle/phi/backends/dynload/flagcx.h"
+#include "paddle/phi/core/distributed/flagcx_tools.h"
+#endif
+
+namespace phi {
+class DenseTensor;
+namespace distributed {
+
+class BKCLCommContext final : public CommContext {
+ public:
+  BKCLCommContext(int rank, int size, BKCLUniqueId BKCL_id);
+#if defined(PADDLE_WITH_FLAGCX)
+  BKCLCommContext(int rank, int size, flagcxHandlerGroup_t flagcx_handler);
+#endif
+  ~BKCLCommContext() override = default;
+
+  BKCLContext_t GetBKCLComm();
+
+  XPUStream GetStream();
+
+  XPUEvent GetComputeEvent();
+
+  void SetComputeEvent(
+      std::shared_ptr<std::remove_pointer<XPUEvent>::type>&& compute_event);
+
+  XPUEvent GetCommEvent();
+
+  void SetCommEvent(
+      std::shared_ptr<std::remove_pointer<XPUEvent>::type>&& comm_event);
+
+  phi::XPUContext* GetDevContext();
+
+  void SetDevContext(std::unique_ptr<phi::XPUContext>&& dev_ctx);
+
+  void Broadcast(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 int root,
+                 XPUStream stream);
+
+  void Send(const DenseTensor& in_tensor,
+            const int64_t& count,
+            const int& peer,
+            XPUStream stream);
+
+  void Recv(DenseTensor* out_tensor,
+            const int64_t& count,
+            const int& peer,
+            XPUStream stream);
+
+  void ReduceScatter(DenseTensor* out_tensor,
+                     const DenseTensor& in_tensor,
+                     BKCLOp reduce_type,
+                     XPUStream stream);
+
+#if defined(PADDLE_WITH_FLAGCX)
+  void Scatter(DenseTensor* out_tensor,
+               const DenseTensor& in_tensor,
+               int root,
+               XPUStream stream);
+#endif
+
+  void AllGather(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 XPUStream stream);
+
+  void AllReduce(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 BKCLOp reduce_type,
+                 XPUStream stream);
+
+  void AllToAll(DenseTensor* out_tensor,
+                const DenseTensor& in_tensor,
+                XPUStream stream);
+
+  void AllToAllUnequalSplit(DenseTensor* out_tensor,
+                            const DenseTensor& in_tensor,
+                            const DenseTensor& out_size_tensor,
+                            const DenseTensor& out_offset_tensor,
+                            const DenseTensor& in_size_tensor,
+                            const DenseTensor& in_offset_tensor,
+                            XPUStream stream);
+
+  void Reduce(DenseTensor* out_tensor,
+              const DenseTensor& in_tensor,
+              BKCLOp reduce_type,
+              int root,
+              XPUStream stream);
+
+  void GroupStart();
+
+  void GroupEnd();
+
+#if defined(PADDLE_WITH_FLAGCX)
+  flagcxRedOp_t BkclToFlagcxRedType(BKCLOp redOp);
+#endif
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(BKCLCommContext);
+
+  BKCLContext_t bkcl_comm_;
+
+  std::unique_ptr<phi::XPUContext> dev_ctx_;
+
+  // used for comm wait compute, compute_stream-->event-->comm_stream
+  std::shared_ptr<std::remove_pointer<XPUEvent>::type> compute_event_;
+
+  // used for compute wait comm, comm_stream-->event-->compute_stream
+  std::shared_ptr<std::remove_pointer<XPUEvent>::type> comm_event_;
+
+#if defined(PADDLE_WITH_FLAGCX)
+
+ public:
+  flagcxHandlerGroup_t flagcx_handler_;
+#endif
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/check/CMakeLists.txt

@@ -0,0 +1,11 @@
+set(CHECK_COMMON_SRCS static_check.cc)
+
+if(WITH_NCCL OR WITH_RCCL)
+  list(APPEND CHECK_COMMON_SRCS nccl_dynamic_check.cc)
+endif()
+
+if(WITH_XPU_BKCL)
+  list(APPEND CHECK_COMMON_SRCS bkcl_dynamic_check.cc)
+endif()
+
+collect_srcs(core_srcs SRCS ${CHECK_COMMON_SRCS})

paddle/phi/core/distributed/check/bkcl_dynamic_check.cc

@@ -0,0 +1,156 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/check/bkcl_dynamic_check.h"
+
+#include "glog/logging.h"
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/utils/string/string_helper.h"
+
+namespace phi::distributed {
+
+void BKCLDynamicCheck::CheckDataType(const DenseTensor& tensor, int64_t dtype) {
+  PADDLE_ENFORCE_EQ(
+      static_cast<int64_t>(tensor.dtype()),
+      dtype,
+      common::errors::InvalidArgument(
+          "Tensors in communication are expected to have the same data type."));
+}
+
+void BKCLDynamicCheck::CheckDataType(const DenseTensor& tensor,
+                                     int root_rank,
+                                     int cur_rank,
+                                     BKCLContext_t comm) {
+  constexpr int kSize = sizeof(int64_t);
+  int64_t dtype_host = static_cast<int64_t>(tensor.dtype());
+  int64_t* dtype_device;
+  PADDLE_ENFORCE_XPU_SUCCESS(
+      xpu_malloc(reinterpret_cast<void**>(&dtype_device), kSize));
+  PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+  PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(
+      dtype_device, &dtype_host, kSize, XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+
+  PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_broadcast(
+      comm, dtype_device, dtype_device, 1, BKCL_INT64, root_rank, 0));
+
+  if (root_rank == cur_rank) {
+    VLOG(3) << "Dynamic check broadcast metadata, dtype: " << dtype_host;
+  } else {
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(
+        &dtype_host, dtype_device, kSize, XPUMemcpyKind::XPU_DEVICE_TO_HOST));
+    VLOG(3) << "Dynamic check recv metadata, dtype: " << dtype_host;
+    CheckDataType(tensor, dtype_host);
+  }
+  PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+  PADDLE_ENFORCE_XPU_SUCCESS(xpu_free(dtype_device));
+}
+
+void BKCLDynamicCheck::CheckShape(const DenseTensor& tensor, int64_t shape) {
+  PADDLE_ENFORCE_EQ(
+      tensor.numel(),
+      shape,
+      common::errors::InvalidArgument(
+          "Tensors in communication are expected to have matching sizes."));
+}
+
+void BKCLDynamicCheck::CheckShape(const DenseTensor& out_tensor,
+                                  const DenseTensor& in_tensor,
+                                  const std::vector<int64_t>& in_size_each_rank,
+                                  int cur_rank,
+                                  int world_size,
+                                  BKCLContext_t comm) {
+  CheckDataType(out_tensor, /*root_rank*/ 0, cur_rank, comm);
+  CheckDataType(in_tensor, /*root_rank*/ 0, cur_rank, comm);
+
+  constexpr int kSize = sizeof(int64_t);
+  int64_t in_row_size =
+      in_tensor.dims()[0] == 0 ? 0 : in_tensor.numel() / in_tensor.dims()[0];
+
+  for (int rank = 0; rank < world_size; ++rank) {
+    int64_t in_shape_host = in_size_each_rank[rank] * in_row_size;
+    int64_t* in_shape_device;
+    PADDLE_ENFORCE_XPU_SUCCESS(
+        xpu_malloc(reinterpret_cast<void**>(&in_shape_device), kSize));
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(in_shape_device,
+                                          &in_shape_host,
+                                          kSize,
+                                          XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+    PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_reduce(comm,
+                                            in_shape_device,
+                                            in_shape_device,
+                                            1,
+                                            BKCL_INT64,
+                                            BKCL_ADD,
+                                            rank,
+                                            0));
+    if (rank == cur_rank) {
+      PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+      PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(&in_shape_host,
+                                            in_shape_device,
+                                            kSize,
+                                            XPUMemcpyKind::XPU_DEVICE_TO_HOST));
+      VLOG(3) << "Dynamic check recv metadata, shape: " << in_shape_host;
+      CheckShape(out_tensor, in_shape_host);
+    }
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_free(in_shape_device));
+  }
+}
+
+void BKCLDynamicCheck::CheckAlltoAllShape(
+    const std::vector<DenseTensor>& out_tensor,
+    const std::vector<DenseTensor>& in_tensor,
+    int cur_rank,
+    int world_size,
+    BKCLContext_t comm) {
+  int64_t first_dtype = static_cast<int64_t>(in_tensor[0].dtype());
+  constexpr int kSize = sizeof(int64_t);
+  CheckDataType(in_tensor[0], /*root_rank*/ 0, cur_rank, comm);
+  for (int rank = 0; rank < world_size; ++rank) {
+    CheckDataType(in_tensor[rank], first_dtype);
+    CheckDataType(out_tensor[rank], first_dtype);
+
+    int64_t in_shape_host = in_tensor[rank].numel();
+    int64_t* in_shape_device;
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_malloc(
+        reinterpret_cast<void**>(&in_shape_device), kSize * world_size));
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(in_shape_device + cur_rank,
+                                          &in_shape_host,
+                                          kSize,
+                                          XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+    PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_all_gather(
+        comm, in_shape_device + cur_rank, 1, in_shape_device, BKCL_INT64, 0));
+    if (rank == cur_rank) {
+      std::vector<int64_t> in_shapes_recv_host(world_size);
+      PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait());
+      PADDLE_ENFORCE_XPU_SUCCESS(xpu_memcpy(in_shapes_recv_host.data(),
+                                            in_shape_device,
+                                            kSize * world_size,
+                                            XPUMemcpyKind::XPU_DEVICE_TO_HOST));
+      VLOG(3) << "Dynamic check recv metadata, shape: "
+              << paddle::string::join_strings(in_shapes_recv_host, ',');
+      for (int out_rank = 0; out_rank < world_size; ++out_rank) {
+        CheckShape(out_tensor[out_rank], in_shapes_recv_host[out_rank]);
+      }
+    }
+    PADDLE_ENFORCE_XPU_SUCCESS(xpu_free(in_shape_device));
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/check/bkcl_dynamic_check.h

@@ -0,0 +1,51 @@
+// Copyright (c) 2025 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.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/distributed/bkcl_comm_context.h"
+
+namespace phi {
+// forward declaration
+class DenseTensor;
+
+namespace distributed {
+struct BKCLDynamicCheck {
+  static void CheckDataType(const DenseTensor& tensor, int64_t dtype);
+
+  static void CheckDataType(const DenseTensor& tensor,
+                            int root_rank,
+                            int cur_rank,
+                            BKCLContext_t comm);
+
+  static void CheckShape(const DenseTensor& tensor, int64_t shape);
+  static void CheckShape(const DenseTensor& out_tensor,
+                         const DenseTensor& in_tensor,
+                         const std::vector<int64_t>& in_size_each_rank,
+                         int cur_rank,
+                         int world_size,
+                         BKCLContext_t comm);
+  static void CheckAlltoAllShape(const std::vector<DenseTensor>& out_tensor,
+                                 const std::vector<DenseTensor>& in_tensor,
+                                 int cur_rank,
+                                 int world_size,
+                                 BKCLContext_t comm);
+};
+
+}  //  namespace distributed
+}  //  namespace phi

paddle/phi/core/distributed/check/nccl_dynamic_check.cc

@@ -0,0 +1,242 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/check/nccl_dynamic_check.h"
+
+#include "glog/logging.h"
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/utils/string/string_helper.h"
+
+#if defined(PADDLE_WITH_RCCL)
+#include <hip/hip_runtime.h>
+
+#include "paddle/phi/backends/dynload/rccl.h"
+
+#define gpuMalloc hipMalloc
+#define gpuMemcpy hipMemcpy
+#define gpuMemcpyDeviceToHost hipMemcpyDeviceToHost
+#define gpuMemcpyHostToDevice hipMemcpyHostToDevice
+#define gpuFree hipFree
+#else
+#include <cuda_runtime.h>
+
+#include "paddle/phi/backends/dynload/nccl.h"
+
+#define gpuMalloc cudaMalloc
+#define gpuMemcpy cudaMemcpy
+#define gpuMemcpyDeviceToHost cudaMemcpyDeviceToHost
+#define gpuMemcpyHostToDevice cudaMemcpyHostToDevice
+#define gpuFree cudaFree
+#endif
+
+namespace phi::distributed {
+void NCCLDynamicCheck::CheckDataType(const DenseTensor& tensor, int64_t dtype) {
+  PADDLE_ENFORCE_EQ(
+      static_cast<int64_t>(tensor.dtype()),
+      dtype,
+      common::errors::InvalidArgument(
+          "Tensors in communication are expected to have the same data type."));
+}
+
+void NCCLDynamicCheck::CheckDataType(const DenseTensor& tensor,
+                                     int root_rank,
+                                     int cur_rank,
+                                     ncclComm_t comm) {
+  constexpr int kSize = sizeof(int64_t);
+  int64_t dtype_host = static_cast<int64_t>(tensor.dtype());
+  int64_t* dtype_device;
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuMalloc(&dtype_device, kSize));
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      gpuMemcpy(dtype_device, &dtype_host, kSize, gpuMemcpyHostToDevice));
+
+  PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclBroadcast(dtype_device,
+                                                         dtype_device,
+                                                         1,
+                                                         ncclInt64,
+                                                         root_rank,
+                                                         comm,
+                                                         kDefaultStream));
+
+  if (root_rank == cur_rank) {
+    VLOG(3) << "Dynamic check broadcast metadata, dtype: " << dtype_host;
+  } else {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        gpuMemcpy(&dtype_host, dtype_device, kSize, gpuMemcpyDeviceToHost));
+    VLOG(3) << "Dynamic check recv metadata, dtype: " << dtype_host;
+    CheckDataType(tensor, dtype_host);
+  }
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuFree(dtype_device));
+}
+
+void NCCLDynamicCheck::CheckShape(const DenseTensor& tensor, int64_t shape) {
+  PADDLE_ENFORCE_EQ(
+      tensor.numel(),
+      shape,
+      common::errors::InvalidArgument(
+          "Tensors in communication are expected to have matching sizes."));
+}
+
+void NCCLDynamicCheck::CheckShape(const DenseTensor& tensor,
+                                  int root_rank,
+                                  int cur_rank,
+                                  ncclComm_t comm) {
+  CheckDataType(tensor, root_rank, cur_rank, comm);
+
+  constexpr int kSize = sizeof(int64_t);
+  int64_t shape_host = tensor.numel();
+  int64_t* shape_device;
+
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuMalloc(&shape_device, kSize));
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      gpuMemcpy(shape_device, &shape_host, kSize, gpuMemcpyHostToDevice));
+
+  PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclBroadcast(shape_device,
+                                                         shape_device,
+                                                         1,
+                                                         ncclInt64,
+                                                         root_rank,
+                                                         comm,
+                                                         kDefaultStream));
+
+  if (root_rank == cur_rank) {
+    VLOG(3) << "Dynamic check broadcast metadata, shape: " << shape_host;
+  } else {
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        gpuMemcpy(&shape_host, shape_device, kSize, gpuMemcpyDeviceToHost));
+    VLOG(3) << "Dynamic check recv metadata, shape: " << shape_host;
+    CheckShape(tensor, shape_host);
+  }
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuFree(shape_device));
+}
+
+void NCCLDynamicCheck::CheckShape(const DenseTensor& out_tensor,
+                                  const DenseTensor& in_tensor,
+                                  const std::vector<int64_t>& in_size_each_rank,
+                                  int cur_rank,
+                                  int world_size,
+                                  ncclComm_t comm) {
+  CheckDataType(out_tensor, /*root_rank*/ 0, cur_rank, comm);
+  CheckDataType(in_tensor, /*root_rank*/ 0, cur_rank, comm);
+
+  constexpr int kSize = sizeof(int64_t);
+  int64_t in_row_size =
+      in_tensor.dims()[0] == 0 ? 0 : in_tensor.numel() / in_tensor.dims()[0];
+
+  for (int rank = 0; rank < world_size; ++rank) {
+    int64_t in_shape_host = in_size_each_rank[rank] * in_row_size;
+    int64_t* in_shape_device;
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuMalloc(&in_shape_device, kSize));
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(
+        in_shape_device, &in_shape_host, kSize, gpuMemcpyHostToDevice));
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclReduce(in_shape_device,
+                                                        in_shape_device,
+                                                        1,
+                                                        ncclInt64,
+                                                        ncclSum,
+                                                        rank,
+                                                        comm,
+                                                        kDefaultStream));
+    if (rank == cur_rank) {
+      PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(
+          &in_shape_host, in_shape_device, kSize, gpuMemcpyDeviceToHost));
+      VLOG(3) << "Dynamic check recv metadata, shape: " << in_shape_host;
+      CheckShape(out_tensor, in_shape_host);
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuFree(in_shape_device));
+  }
+}
+
+void NCCLDynamicCheck::CheckAlltoAllShape(
+    const std::vector<DenseTensor>& out_tensor,
+    const std::vector<DenseTensor>& in_tensor,
+    int cur_rank,
+    int world_size,
+    ncclComm_t comm) {
+  int64_t first_dtype = static_cast<int64_t>(in_tensor[0].dtype());
+  constexpr int kSize = sizeof(int64_t);
+  CheckDataType(in_tensor[0], /*root_rank*/ 0, cur_rank, comm);
+  for (int rank = 0; rank < world_size; ++rank) {
+    CheckDataType(in_tensor[rank], first_dtype);
+    CheckDataType(out_tensor[rank], first_dtype);
+
+    int64_t in_shape_host = in_tensor[rank].numel();
+    int64_t* in_shape_device;
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuMalloc(&in_shape_device, kSize * world_size));
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(in_shape_device + cur_rank,
+                                         &in_shape_host,
+                                         kSize,
+                                         gpuMemcpyHostToDevice));
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        phi::dynload::ncclAllGather(in_shape_device + cur_rank,
+                                    in_shape_device,
+                                    1,
+                                    ncclInt64,
+                                    comm,
+                                    kDefaultStream));
+    if (rank == cur_rank) {
+      std::vector<int64_t> in_shapes_recv_host(world_size);
+      PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(in_shapes_recv_host.data(),
+                                           in_shape_device,
+                                           kSize * world_size,
+                                           gpuMemcpyDeviceToHost));
+      VLOG(3) << "Dynamic check recv metadata, shape: "
+              << paddle::string::join_strings(in_shapes_recv_host, ',');
+      for (int out_rank = 0; out_rank < world_size; ++out_rank) {
+        CheckShape(out_tensor[out_rank], in_shapes_recv_host[out_rank]);
+      }
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuFree(in_shape_device));
+  }
+}
+
+void NCCLDynamicCheck::CheckGatherShape(
+    const DenseTensor& in_tensor,
+    const std::vector<DenseTensor>& out_tensors,
+    int root_rank,
+    int cur_rank,
+    int world_size,
+    ncclComm_t comm) {
+  std::vector<int64_t> shapes(world_size, 0);
+  shapes[cur_rank] = in_tensor.numel();
+  int64_t* in_shape_device;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      gpuMalloc(&in_shape_device, world_size * sizeof(int64_t)));
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(in_shape_device,
+                                       shapes.data(),
+                                       world_size * sizeof(int64_t),
+                                       gpuMemcpyHostToDevice));
+
+  PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclAllReduce(in_shape_device,
+                                                         in_shape_device,
+                                                         world_size,
+                                                         ncclInt64,
+                                                         ncclSum,
+                                                         comm,
+                                                         kDefaultStream));
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuMemcpy(shapes.data(),
+                                       in_shape_device,
+                                       world_size * sizeof(int64_t),
+                                       gpuMemcpyDeviceToHost));
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuFree(in_shape_device));
+
+  if (cur_rank == root_rank) {
+    for (int i = 0; i < world_size; i++) {
+      CheckShape(out_tensors[i], shapes[i]);
+    }
+  }
+}
+}  // namespace phi::distributed

paddle/phi/core/distributed/check/nccl_dynamic_check.h

@@ -0,0 +1,75 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#include "paddle/phi/backends/gpu/forwards.h"
+
+#if defined(PADDLE_WITH_RCCL)
+using gpuStream_t = hipStream_t;
+#else
+using gpuStream_t = cudaStream_t;
+#endif
+
+namespace phi {
+// forward declaration
+class DenseTensor;
+
+namespace distributed {
+struct NCCLDynamicCheck {
+  static void CheckDataType(const DenseTensor& tensor, int64_t dtype);
+
+  static void CheckDataType(const DenseTensor& tensor,
+                            int root_rank,
+                            int cur_rank,
+                            ncclComm_t comm);
+
+  static void CheckShape(const DenseTensor& tensor, int64_t shape);
+
+  static void CheckShape(const DenseTensor& tensor,
+                         int root_rank,
+                         int cur_rank,
+                         ncclComm_t comm);
+
+  static void CheckShape(const DenseTensor& out_tensor,
+                         const DenseTensor& in_tensor,
+                         const std::vector<int64_t>& in_size_each_rank,
+                         int cur_rank,
+                         int world_size,
+                         ncclComm_t comm);
+
+  static void CheckAlltoAllShape(const std::vector<DenseTensor>& out_tensor,
+                                 const std::vector<DenseTensor>& in_tensor,
+                                 int cur_rank,
+                                 int world_size,
+                                 ncclComm_t comm);
+
+  // can be used to check gather and all gather
+  static void CheckGatherShape(const DenseTensor& in_tensor,
+                               const std::vector<DenseTensor>& out_tensors,
+                               int root_rank,
+                               int cur_rank,
+                               int world_size,
+                               ncclComm_t comm);
+
+ private:
+  // `0` represents default stream for both cuda & hip
+  static constexpr gpuStream_t kDefaultStream = 0;
+};
+
+}  //  namespace distributed
+}  //  namespace phi

paddle/phi/core/distributed/check/static_check.cc

@@ -0,0 +1,215 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/check/static_check.h"
+
+#include <cstdlib>
+#include <cstring>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi::distributed {
+
+void CommStaticCheck::CheckRank(int rank, int world_size) {
+  PADDLE_ENFORCE_GE(rank,
+                    0,
+                    common::errors::InvalidArgument(
+                        "Rank should be greater than or equal to 0."));
+  PADDLE_ENFORCE_LT(
+      rank,
+      world_size,
+      common::errors::InvalidArgument("Rank is out of the process group."));
+}
+
+void CommStaticCheck::CheckPlace(const DenseTensor& tensor,
+                                 AllocationType place) {
+  PADDLE_ENFORCE_EQ(
+      tensor.place().GetType(),
+      place,
+      common::errors::InvalidArgument("Tensor should be in backend's place."));
+}
+
+void CommStaticCheck::CheckPlace(const DenseTensor& out_tensor,
+                                 const DenseTensor& in_tensor,
+                                 AllocationType place) {
+  CheckPlace(out_tensor, place);
+  CheckPlace(in_tensor, place);
+  PADDLE_ENFORCE_EQ(
+      out_tensor.place(),
+      in_tensor.place(),
+      common::errors::InvalidArgument(
+          "Input and output tensors should be on the same place."));
+}
+
+void CommStaticCheck::CheckDataType(const DenseTensor& out_tensor,
+                                    const DenseTensor& in_tensor) {
+  PADDLE_ENFORCE_EQ(
+      out_tensor.dtype(),
+      in_tensor.dtype(),
+      common::errors::InvalidArgument(
+          "Input and output tensors should have the same data type."));
+}
+
+void CommStaticCheck::CheckDataType(
+    const std::vector<DenseTensor>& out_tensors,
+    const std::vector<DenseTensor>& in_tensors) {
+  if (in_tensors.empty() && out_tensors.empty()) {
+    return;
+  }
+
+  auto first_dtype =
+      (in_tensors.empty() ? out_tensors[0] : in_tensors[0]).dtype();
+
+  for (const auto& tensor : in_tensors) {
+    PADDLE_ENFORCE_EQ(
+        tensor.dtype(),
+        first_dtype,
+        common::errors::InvalidArgument(
+            "Input and output tensors should have the same data type."));
+  }
+  for (const auto& tensor : out_tensors) {
+    PADDLE_ENFORCE_EQ(
+        tensor.dtype(),
+        first_dtype,
+        common::errors::InvalidArgument(
+            "Input and output tensors should have the same data type."));
+  }
+}
+
+void CommStaticCheck::CheckShape(const DenseTensor& tensor) {
+  PADDLE_ENFORCE_GT(tensor.numel(),
+                    0,
+                    common::errors::InvalidArgument(
+                        "Size of tensor should be greater than 0."));
+}
+
+void CommStaticCheck::CheckShape(const DenseTensor& out_tensor,
+                                 const DenseTensor& in_tensor,
+                                 int out_size_factor,
+                                 int in_size_factor) {
+  CheckShape(out_tensor);
+  CheckShape(in_tensor);
+  int64_t out_size = out_tensor.numel(), in_size = in_tensor.numel();
+  PADDLE_ENFORCE_EQ(
+      out_size * out_size_factor,
+      in_size * in_size_factor,
+      common::errors::InvalidArgument(
+          "Input and output tensors should have matching sizes. "
+          "out_size=%ld, out_size_factor=%d, in_size=%ld, in_size_factor=%d",
+          out_size,
+          out_size_factor,
+          in_size,
+          in_size_factor));
+}
+
+void CommStaticCheck::CheckShape(const DenseTensor& out_tensor,
+                                 const DenseTensor& in_tensor,
+                                 int dst_rank,
+                                 int cur_rank,
+                                 int world_size,
+                                 int out_size_factor,
+                                 int in_size_factor,
+                                 AllocationType place) {
+  CheckRank(dst_rank, world_size);
+  CheckRank(cur_rank, world_size);
+
+  CheckPlace(out_tensor, in_tensor, place);
+  CheckDataType(out_tensor, in_tensor);
+
+  if (dst_rank == cur_rank) {
+    CheckShape(out_tensor, in_tensor, out_size_factor, in_size_factor);
+  } else {
+    CheckShape(out_tensor);
+    CheckShape(in_tensor);
+  }
+}
+
+void CommStaticCheck::CheckShape(const DenseTensor& tensor,
+                                 int rank,
+                                 int world_size,
+                                 AllocationType place) {
+  CheckPlace(tensor, place);
+  CheckRank(rank, world_size);
+}
+
+void CommStaticCheck::SameShape(const DenseTensor& out_tensor,
+                                const DenseTensor& in_tensor,
+                                int dst_rank,
+                                int cur_rank,
+                                int world_size,
+                                AllocationType place) {
+  CheckShape(out_tensor,
+             in_tensor,
+             dst_rank,
+             cur_rank,
+             world_size,
+             /*out_size_factor*/ 1,
+             /*in_size_factor*/ 1,
+             place);
+}
+
+void CommStaticCheck::ScatterLikeShape(const DenseTensor& out_tensor,
+                                       const DenseTensor& in_tensor,
+                                       int dst_rank,
+                                       int cur_rank,
+                                       int world_size,
+                                       AllocationType place) {
+  CheckShape(out_tensor,
+             in_tensor,
+             dst_rank,
+             cur_rank,
+             world_size,
+             /*out_size_factor*/ world_size,
+             /*in_size_factor*/ 1,
+             place);
+}
+
+void CommStaticCheck::GatherLikeShape(const DenseTensor& out_tensor,
+                                      const DenseTensor& in_tensor,
+                                      int dst_rank,
+                                      int cur_rank,
+                                      int world_size,
+                                      AllocationType place) {
+  CheckRank(dst_rank, world_size);
+  CheckRank(cur_rank, world_size);
+
+  CheckPlace(out_tensor, in_tensor, place);
+  CheckDataType(out_tensor, in_tensor);
+
+  CheckGatherShape(out_tensor);
+  CheckGatherShape(in_tensor);
+  int64_t out_size = out_tensor.numel(), in_size = in_tensor.numel();
+  PADDLE_ENFORCE_EQ(
+      out_size,
+      in_size * world_size,
+      common::errors::InvalidArgument(
+          "Input and output tensors should have matching sizes. "
+          "out_size=%ld, out_size_factor=%d, in_size=%ld, in_size_factor=%d",
+          out_size,
+          1,
+          in_size,
+          world_size));
+}
+
+void CommStaticCheck::CheckGatherShape(const phi::DenseTensor& tensor) {
+  PADDLE_ENFORCE_GE(
+      tensor.numel(),
+      0,
+      common::errors::InvalidArgument("Size of tensor should be greater equal "
+                                      "than 0 in gather-liked communication."));
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/check/static_check.h

@@ -0,0 +1,87 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include "paddle/phi/common/place.h"
+
+namespace phi {
+// forward declaration
+class DenseTensor;
+
+namespace distributed {
+struct CommStaticCheck {
+  static void CheckRank(int rank, int world_size);
+
+  static void CheckPlace(const DenseTensor& tensor, AllocationType place);
+
+  static void CheckPlace(const DenseTensor& out_tensor,
+                         const DenseTensor& in_tensor,
+                         AllocationType place);
+
+  static void CheckDataType(const DenseTensor& out_tensor,
+                            const DenseTensor& in_tensor);
+
+  static void CheckDataType(const std::vector<DenseTensor>& out_tensors,
+                            const std::vector<DenseTensor>& in_tensors);
+
+  static void CheckShape(const DenseTensor& tensor);
+
+  static void CheckShape(const DenseTensor& out_tensor,
+                         const DenseTensor& in_tensor,
+                         int out_size_factor,
+                         int in_size_factor);
+
+  static void CheckShape(const DenseTensor& out_tensor,
+                         const DenseTensor& in_tensor,
+                         int dst_rank,
+                         int cur_rank,
+                         int world_size,
+                         int out_size_factor,
+                         int in_size_factor,
+                         AllocationType place = AllocationType::GPU);
+
+  // for p2p
+  static void CheckShape(const DenseTensor& tensor,
+                         int rank,
+                         int world_size,
+                         AllocationType place = AllocationType::GPU);
+
+  // for collective
+  static void SameShape(const DenseTensor& out_tensor,
+                        const DenseTensor& in_tensor,
+                        int dst_rank,
+                        int cur_rank,
+                        int world_size,
+                        AllocationType place = AllocationType::GPU);
+
+  static void ScatterLikeShape(const DenseTensor& out_tensor,
+                               const DenseTensor& in_tensor,
+                               int dst_rank,
+                               int cur_rank,
+                               int world_size,
+                               AllocationType place = AllocationType::GPU);
+
+  static void GatherLikeShape(const DenseTensor& out_tensor,
+                              const DenseTensor& in_tensor,
+                              int dst_rank,
+                              int cur_rank,
+                              int world_size,
+                              AllocationType place = AllocationType::GPU);
+
+  static void CheckGatherShape(const phi::DenseTensor& tensor);
+};
+
+}  //  namespace distributed
+}  //  namespace phi

paddle/phi/core/distributed/collective/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(COLLECTIVE_SRCS process_group.cc)
+
+collect_srcs(core_srcs SRCS ${COLLECTIVE_SRCS})

paddle/phi/core/distributed/collective/process_group.cc

@@ -0,0 +1,55 @@
+// Copyright (c) 2024 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.
+
+#include "paddle/phi/core/distributed/collective/process_group.h"
+
+namespace phi::distributed {
+
+bool ProcessGroup::Task::IsCompleted() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  return is_completed_;
+}
+
+ProcessGroup::ProcessGroup(int rank, int size, int gid)
+    : rank_(rank), size_(size), gid_(gid) {
+  if (gid != kIgnoreId) {
+    auto map = ProcessGroupMapFromGid::getInstance();
+    map->insert(gid_, this);
+  }
+  const char* global_rank = std::getenv("PADDLE_TRAINER_ID");
+  PADDLE_ENFORCE_NOT_NULL(
+      global_rank,
+      common::errors::NotFound(
+          "The environment variable 'PADDLE_TRAINER_ID' cannot be found."));
+  global_rank_ = std::atoi(global_rank);
+}
+
+// TODO(sunyilun): methods below will be removed later
+ProcessGroupIdMap& ProcessGroupIdMap::GetInstance() {
+  static ProcessGroupIdMap instance;
+  return instance;
+}
+
+void ProcessGroupIdMap::DestroyProcessGroup() {
+  auto& id_map = ProcessGroupIdMap::GetInstance();
+  for (auto& item : id_map) {
+    auto use_count = item.second.use_count();
+    for (int i = 0; i < use_count; ++i) {
+      item.second.reset();
+    }
+  }
+  id_map.clear();
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/collective/process_group.h

@@ -0,0 +1,598 @@
+// Copyright (c) 2024 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.
+
+#pragma once
+
+#include <chrono>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/distributed/types.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/enforce.h"
+
+constexpr auto kWaitTimeout = std::chrono::milliseconds(0);
+
+namespace phi {
+namespace distributed {
+
+using phi::distributed::AllreduceOptions;
+using phi::distributed::BarrierOptions;
+using phi::distributed::BroadcastOptions;
+using phi::distributed::CommType;
+using phi::distributed::GatherOptions;
+using phi::distributed::GetPartialTensor;
+using phi::distributed::ReduceOp;
+using phi::distributed::ReduceOptions;
+using phi::distributed::ReduceScatterOptions;
+using phi::distributed::ScatterOptions;
+constexpr int kIgnoreId = -1;
+
+class ProcessGroup {
+ public:
+  class Task {
+   public:
+    Task(int rank, CommType comm_type, bool sync_op)
+        : rank_(rank), comm_type_(comm_type), sync_op_(sync_op) {}
+    virtual ~Task() = default;
+
+    virtual bool IsCompleted();
+    virtual bool Wait(std::chrono::milliseconds timeout UNUSED = kWaitTimeout) {
+      return false;
+    }
+    virtual void Synchronize() {}
+    virtual void UpdateWaitChain(const phi::DeviceContext& ctx UNUSED) {}
+
+    bool IsSync() const { return sync_op_; }
+
+    // TODO(sunyilun): methods below will be removed later
+    Task(int rank,
+         const std::vector<DenseTensor>& inputs UNUSED,
+         CommType comm_type)
+        : rank_(rank), comm_type_(comm_type) {}
+    Task(int rank,
+         const std::vector<DenseTensor>& inputs UNUSED,
+         CommType comm_type,
+         bool sync_op)
+        : rank_(rank), comm_type_(comm_type), sync_op_(sync_op) {}
+
+   protected:
+    const int rank_;
+    CommType comm_type_{CommType::UNKNOWN};
+    std::mutex mutex_;
+    bool is_completed_{false};
+
+   private:
+    bool sync_op_{true};
+  };
+
+ public:
+  ProcessGroup(int rank, int size, int gid);
+  virtual ~ProcessGroup() = default;
+
+  int GetRank() const { return rank_; }
+
+  int GetSize() const { return size_; }
+
+  int GetGid() const { return gid_; }
+
+  std::string GetGroupMessage() const {
+    return std::string("rank_in_group: ") + std::to_string(rank_) +
+           std::string(", nranks: ") + std::to_string(size_) +
+           std::string(", gid: ") + std::to_string(gid_) +
+           std::string(", backend: ") + GetBackendName();
+  }
+
+  virtual std::string GetBackendName() const = 0;
+
+  virtual DeviceContext* GetDeviceContext(const Place& place UNUSED) const {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support get device_context.",
+        GetBackendName()));
+  }
+
+  virtual DeviceContext* GetDeviceContext(const Place& place UNUSED,
+                                          bool use_calc_stream UNUSED) const {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support get device_context.",
+        GetBackendName()));
+  }
+
+  virtual void StartCoalescing() {}
+
+  virtual void EndCoalescing(
+      std::optional<std::vector<std::shared_ptr<ProcessGroup::Task>>>
+          tasks_opt = std::nullopt) {}
+
+  virtual void EagerConnect() {}
+
+  virtual void EagerConnectRingExchange() {}
+
+  // without stream APIs
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_gather with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      int64_t offset UNUSED,
+      int64_t numel UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_gather with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllReduce(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const AllreduceOptions& opts UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_reduce with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllToAll(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const std::vector<int64_t>& out_size_each_rank UNUSED,
+      const std::vector<int64_t>& in_size_each_rank UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_to_all with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllToAll(
+      std::vector<DenseTensor>* out_tensors UNUSED,
+      const std::vector<DenseTensor>& in_tensors UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_to_all with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Barrier(
+      const BarrierOptions& UNUSED = BarrierOptions()) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support barrier.", GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Broadcast(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const BroadcastOptions& opts UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support broadcast with sync_op flag",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Reduce(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ReduceOptions& opts UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support reduce with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> ReduceScatter(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ReduceScatterOptions& opts UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support reduce_scatter with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Scatter(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ScatterOptions& opts UNUSED,
+      bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support scatter with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Recv(DenseTensor* tensor UNUSED,
+                                                   int src_rank UNUSED,
+                                                   bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support recv with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Recv(DenseTensor* tensor UNUSED,
+                                                   int src_rank UNUSED,
+                                                   int64_t offset UNUSED,
+                                                   int64_t numel UNUSED,
+                                                   bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support recv with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Send(const DenseTensor& tensor
+                                                       UNUSED,
+                                                   int dst_rank UNUSED,
+                                                   bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support send with sync_op flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Send(const DenseTensor& tensor
+                                                       UNUSED,
+                                                   int dst_rank UNUSED,
+                                                   int64_t offset UNUSED,
+                                                   int64_t numel UNUSED,
+                                                   bool sync_op UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support send with sync_op flag.",
+        GetBackendName()));
+  }
+
+  // stream APIs
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_gather "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      int64_t offset UNUSED,
+      int64_t numel UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_gather "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllReduce(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const AllreduceOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_reduce "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllToAll(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const std::vector<int64_t>& out_size_each_rank UNUSED,
+      const std::vector<int64_t>& in_size_each_rank UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_to_all "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllToAll(
+      std::vector<DenseTensor>* out_tensors UNUSED,
+      const std::vector<DenseTensor>& in_tensors UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support all_to_all "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Broadcast(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const BroadcastOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support broadcast "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Reduce(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ReduceOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support reduce "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> ReduceScatter(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ReduceScatterOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support reduce_scatter "
+        "with sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Scatter(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const ScatterOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support scatter "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Gather(
+      DenseTensor* out_tensor UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const GatherOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support gather "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Gather(
+      std::vector<DenseTensor>* gather_tensors_ptr UNUSED,
+      const DenseTensor& in_tensor UNUSED,
+      const GatherOptions& opts UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support gather "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Recv(DenseTensor* tensor UNUSED,
+                                                   int src_rank UNUSED,
+                                                   bool sync_op UNUSED,
+                                                   bool use_calc_stream
+                                                       UNUSED) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "ProcessGroup%s does not support recv with "
+        "sync_op and use_calc_stream flag.",
+        GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Recv(DenseTensor* tensor UNUSED,
+                                                   int src_rank UNUSED,
+                                                   int64_t offset UNUSED,
+                                                   int64_t numel UNUSED,
+                                                   bool sync_op UNUSED,
+                                                   bool use_calc_stream
+                                                       UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support recv "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Send(
+      const DenseTensor& tensor UNUSED,
+      int dst_rank UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support send "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Send(
+      const DenseTensor& tensor UNUSED,
+      int dst_rank UNUSED,
+      int64_t offset UNUSED,
+      int64_t numel UNUSED,
+      bool sync_op UNUSED,
+      bool use_calc_stream UNUSED) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("ProcessGroup%s does not support send "
+                                      "with sync_op and use_calc_stream flag.",
+                                      GetBackendName()));
+  }
+
+  // legacy APIs
+  // TODO(liyurui): This API will be moved later
+  virtual std::shared_ptr<ProcessGroup::Task> AllReduce(
+      std::vector<DenseTensor>& inputs,   // NOLINT
+      std::vector<DenseTensor>& outputs,  // NOLINT
+      const AllreduceOptions& options = AllreduceOptions()) {
+    return AllReduce(outputs.data(), inputs.front(), options, false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllReduce(
+      std::vector<DenseTensor>& inputs,   // NOLINT
+      std::vector<DenseTensor>& outputs,  // NOLINT
+      const AllreduceOptions& options,
+      bool sync_op) {
+    return AllReduce(outputs.data(), inputs.front(), options, sync_op);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllReduce(
+      std::vector<DenseTensor>& inputs,   // NOLINT
+      std::vector<DenseTensor>& outputs,  // NOLINT
+      const AllreduceOptions& options,
+      bool use_calc_stream,
+      bool sync_op) {
+    return AllReduce(
+        outputs.data(), inputs.front(), options, use_calc_stream, sync_op);
+  }
+
+  // TODO(sunyilun): methods below will be removed later
+  virtual std::shared_ptr<ProcessGroup::Task> Broadcast(
+      std::vector<DenseTensor>& inputs,   // NOLINT
+      std::vector<DenseTensor>& outputs,  // NOLINT
+      const BroadcastOptions& options = BroadcastOptions()) {
+    return Broadcast(outputs.data(), inputs.front(), options, false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Broadcast(
+      std::vector<DenseTensor>& inputs,   // NOLINT
+      std::vector<DenseTensor>& outputs,  // NOLINT
+      const BroadcastOptions& options,
+      bool sync_op) {
+    return Broadcast(outputs.data(), inputs.front(), options, sync_op);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Send(
+      std::vector<DenseTensor>& tensors, int dst_rank) {  // NOLINT
+    return Send(tensors.front(), dst_rank, false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Recv(
+      std::vector<DenseTensor>& tensors, int src_rank) {  // NOLINT
+    return Recv(&tensors.front(), src_rank, false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      std::vector<DenseTensor>& in_tensors,     // NOLINT
+      std::vector<DenseTensor>& out_tensors) {  // NOLINT
+    return AllGather(out_tensors.data(), in_tensors.front(), false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      std::vector<DenseTensor>& in_tensors,   // NOLINT
+      std::vector<DenseTensor>& out_tensors,  // NOLINT
+      bool sync_op) {
+    return AllGather(out_tensors.data(), in_tensors.front(), sync_op);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllGather(
+      std::vector<DenseTensor>& in_tensors,   // NOLINT
+      std::vector<DenseTensor>& out_tensors,  // NOLINT
+      bool use_calc_stream,
+      bool sync_op) {
+    return AllGather(
+        out_tensors.data(), in_tensors.front(), use_calc_stream, sync_op);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> AllToAll(
+      std::vector<DenseTensor>&,    // NOLINT
+      std::vector<DenseTensor>&) {  // NOLINT
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "ProcessGroup%s does not support AllToAll", GetBackendName()));
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Reduce(
+      std::vector<DenseTensor>& ins,   // NOLINT
+      std::vector<DenseTensor>& outs,  // NOLINT
+      const ReduceOptions& opts) {
+    return Reduce(outs.data(), ins.front(), opts, false);
+  }
+
+  virtual std::shared_ptr<ProcessGroup::Task> Scatter(
+      std::vector<DenseTensor>& ins,   // NOLINT
+      std::vector<DenseTensor>& outs,  // NOLINT
+      const ScatterOptions& opts) {
+    return Scatter(outs.data(), ins.front(), opts, false);
+  }
+
+ protected:
+  int global_rank_{-1};
+  int rank_;
+  int size_;
+  int gid_;
+};
+
+class PADDLE_API ProcessGroupIdMap
+    : public std::unordered_map<int, std::shared_ptr<ProcessGroup>> {
+ public:
+  static ProcessGroupIdMap& GetInstance();
+  static void DestroyProcessGroup();
+};
+
+// TODO(dev): The following method will be removed soon.
+class ProcessGroupMapFromGid {
+ public:
+  bool has(int gid) { return map_.find(gid) != map_.end(); }
+
+  void insert(int gid, ProcessGroup* pg) { map_[gid] = pg; }
+
+  ProcessGroup* get(int gid) {
+    auto it = map_.find(gid);
+    if (it == map_.end()) {
+      return nullptr;
+    }
+    return it->second;
+  }
+
+  static std::shared_ptr<ProcessGroupMapFromGid> getInstance() {
+    static auto s_instance = std::make_shared<ProcessGroupMapFromGid>();
+    return s_instance;
+  }
+
+  ProcessGroupMapFromGid() = default;
+  ~ProcessGroupMapFromGid() = default;
+
+ private:
+  std::unordered_map<int, ProcessGroup*> map_;
+};
+
+static void CheckTensorContiguous(const DenseTensor& tensor) {
+  if (!tensor.meta().is_contiguous()) {
+    PADDLE_THROW(
+        common::errors::InvalidArgument("The tensor must be contiguous"));
+  }
+}
+
+static void CheckTensorContiguous(const std::vector<DenseTensor>& inputs) {
+  for (const auto& tensor : inputs) {
+    if (!tensor.meta().is_contiguous()) {
+      PADDLE_THROW(
+          common::errors::InvalidArgument("The tensor must be contiguous"));
+    }
+  }
+}
+
+}  //  namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/comm_context.h

@@ -0,0 +1,39 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/common/macros.h"
+
+namespace phi {
+namespace distributed {
+
+class CommContext {
+ public:
+  CommContext(int rank, int size) : rank_(rank), size_(size) {}
+  virtual ~CommContext() = default;
+
+  int GetRank() const { return rank_; }
+  int GetSize() const { return size_; }
+
+ protected:
+  int rank_;
+  int size_;
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(CommContext);
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/comm_context_manager.cc

@@ -0,0 +1,437 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+
+#include <memory>
+#include <string>
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/core/distributed/store/store.h"
+#include "paddle/phi/core/enforce.h"
+
+#if defined(PADDLE_WITH_GLOO)
+#include <gloo/rendezvous/prefix_store.h>
+#include "paddle/phi/core/distributed/gloo_comm_context.h"
+#include "paddle/phi/core/distributed/gloo_utils.h"
+#include "paddle/phi/core/distributed/store/gloo_store.h"
+#endif
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/distributed/nccl_comm_context.h"
+#include "paddle/phi/core/distributed/nccl_tools.h"
+#elif defined(PADDLE_WITH_XPU_BKCL)
+#include "paddle/phi/backends/xpu/xpu_info.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/distributed/bkcl_comm_context.h"
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/core/distributed/xccl_comm_context.h"
+#endif
+
+#if defined(PADDLE_WITH_FLAGCX)
+#if !defined(PADDLE_WITH_XPU)
+#include "paddle/phi/core/distributed/flagcx_comm_context.h"
+#endif
+#include "paddle/phi/backends/dynload/flagcx.h"
+#include "paddle/phi/core/distributed/flagcx_tools.h"
+#endif
+
+namespace phi::distributed {
+
+int CommContextManager::device_id = -1;
+
+void CommContextManager::SetDeviceId(int dev_id) {
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+  phi::backends::gpu::SetDeviceId(dev_id);
+  CommContextManager::device_id = dev_id;
+#elif defined(PADDLE_WITH_XPU_BKCL)
+  phi::backends::xpu::SetXPUDeviceId(dev_id);
+  CommContextManager::device_id = dev_id;
+#endif
+}
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+void CommContextManager::CreateNCCLCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    int rank,
+    int size,
+    const std::string& hash_key,
+    const P2POption* p2p_opt,
+    int nccl_comm_init_option,
+    std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  if (comm_context_manager.Has(unique_comm_key)) {
+    return;
+  }
+  ncclUniqueId nccl_id;
+  if (rank == 0 || (p2p_opt && p2p_opt->is_p2p_op && p2p_opt->p2p_rank == 0)) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclGetUniqueId(&nccl_id));
+  }
+
+  std::string unique_key = "NCCLCommContext/" + unique_comm_key + hash_key;
+  if (rank == 0 || (p2p_opt && p2p_opt->is_p2p_op && p2p_opt->p2p_rank == 0)) {
+    std::vector<uint8_t> nccl_id_wrapper(
+        reinterpret_cast<uint8_t*>(&nccl_id),
+        reinterpret_cast<uint8_t*>(&nccl_id) + NCCL_UNIQUE_ID_BYTES);
+    store->set(unique_key, nccl_id_wrapper);
+  } else {
+    const auto& nccl_id_wrapper = store->get(unique_key);
+    std::memcpy(&nccl_id, nccl_id_wrapper.data(), nccl_id_wrapper.size());
+  }
+
+  if (p2p_opt) {
+    rank = p2p_opt->rank;
+    size = p2p_opt->num_ranks;
+  }
+  VLOG(3) << "init NCCLCommContext rank: " << rank << ", size: " << size
+          << ", unique_comm_key: " << unique_comm_key
+          << ", unique_key: " << unique_key
+          << ", nccl_id: " << SerializeNCCLUniqueId(nccl_id);
+  auto nccl_comm_context = std::make_unique<NCCLCommContext>(
+      rank, size, nccl_id, nccl_comm_init_option, nccl_config_ptr);
+  if (CommContextManager::device_id != -1) {
+    std::unique_ptr<phi::GPUContext> dev_ctx(
+        new phi::GPUContext(phi::GPUPlace(CommContextManager::device_id)));
+    dev_ctx->SetAllocator(phi::memory_utils::GetAllocator(
+        CommContextManager::device_id, dev_ctx->stream()));
+    dev_ctx->SetHostAllocator(phi::memory_utils::GetHostAllocator());
+    dev_ctx->SetZeroAllocator(
+        phi::memory_utils::GetZeroAllocator(CommContextManager::device_id));
+    dev_ctx->SetHostZeroAllocator(phi::memory_utils::GetHostZeroAllocator());
+    dev_ctx->SetPinnedAllocator(phi::memory_utils::GetPinnedAllocator());
+    dev_ctx->PartialInitWithAllocator();
+    auto compute_event =
+        phi::memory_utils::GetCudaEvent(CommContextManager::device_id);
+    auto comm_event =
+        phi::memory_utils::GetCudaEvent(CommContextManager::device_id);
+
+    nccl_comm_context->SetDevContext(std::move(dev_ctx));
+    nccl_comm_context->SetComputeEvent(std::move(compute_event));
+    nccl_comm_context->SetCommEvent(std::move(comm_event));
+  }
+
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(nccl_comm_context));
+}
+
+void CommContextManager::RecreateNCCLComm(const std::shared_ptr<Store>& store,
+                                          const std::string& unique_comm_key,
+                                          int rank,
+                                          const std::string& recreate_key,
+                                          const P2POption* p2p_opt) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+
+  ncclUniqueId nccl_id;
+  if (rank == 0 || (p2p_opt && p2p_opt->is_p2p_op && p2p_opt->p2p_rank == 0)) {
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclGetUniqueId(&nccl_id));
+  }
+
+  std::string unique_key =
+      "NCCLCommContext/" + unique_comm_key + "/" + recreate_key;
+  if (rank == 0 || (p2p_opt && p2p_opt->is_p2p_op && p2p_opt->p2p_rank == 0)) {
+    std::vector<uint8_t> nccl_id_wrapper(
+        reinterpret_cast<uint8_t*>(&nccl_id),
+        reinterpret_cast<uint8_t*>(&nccl_id) + NCCL_UNIQUE_ID_BYTES);
+    store->set(unique_key, nccl_id_wrapper);
+  } else {
+    const auto& nccl_id_wrapper = store->get(unique_key);
+    std::memcpy(&nccl_id, nccl_id_wrapper.data(), nccl_id_wrapper.size());
+  }
+
+  VLOG(3) << "RecreateNCCLComm nccl_id: " << SerializeNCCLUniqueId(nccl_id);
+
+  auto comm_context = static_cast<phi::distributed::NCCLCommContext*>(
+      comm_context_manager.Get(unique_comm_key));
+  comm_context->CreateNCCLComm(nccl_id);
+
+  comm_context_manager.SetStore(store);
+}
+#endif
+
+#if defined(PADDLE_WITH_GLOO)
+void CommContextManager::CreateGlooCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    int rank,
+    int size) {
+  GlooStore store_wrapper(store);
+  auto gloo_store = std::make_shared<gloo::rendezvous::PrefixStore>(
+      unique_comm_key, store_wrapper);
+
+  auto gloo_device = CreateGlooDevice();
+
+  auto gloo_comm_context =
+      std::make_unique<GlooCommContext>(rank, size, gloo_store, gloo_device);
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  // set actual store to manager
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(gloo_comm_context));
+}
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+void CommContextManager::CreateXCCLCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    const Place& place,
+    int rank,
+    int size,
+    const std::string& hash_key) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  if (comm_context_manager.Has(unique_comm_key)) {
+    return;
+  }
+  phi::ccl::CCLRootId xccl_root_id;
+  if (rank == 0) {
+    phi::DeviceManager::CCLGetUniqueId(place.GetDeviceType(), &xccl_root_id);
+  }
+
+  std::string unique_key = "XCCLCommContext/" + unique_comm_key;
+  if (!hash_key.empty()) {
+    unique_key += "/" + hash_key;
+  }
+  if (rank == 0) {
+    store->set(unique_key, xccl_root_id);
+  } else {
+    xccl_root_id = store->get(unique_key);
+  }
+  VLOG(3) << "init xccl rank: " << rank << ", nranks: " << size
+          << ", unique_comm_key: " << unique_comm_key << ", xccl uniqueid: "
+          << phi::ccl::SerializeXCCLUniqueId(xccl_root_id);
+  auto xccl_comm_context =
+      std::make_unique<XCCLCommContext>(place, rank, size, xccl_root_id);
+  if (CommContextManager::device_id != -1) {
+    std::unique_ptr<phi::CustomContext> dev_ctx(
+
+        new phi::CustomContext(CustomPlace(place)));
+    // dev_ctx->SetAllocator(phi::memory_utils::GetAllocator(
+    //     CommContextManager::device_id, dev_ctx->GetStream()));
+    // dev_ctx->SetHostAllocator(phi::memory_utils::GetHostAllocator());
+    // dev_ctx->SetZeroAllocator(
+    //     phi::memory_utils::GetZeroAllocator(CommContextManager::device_id));
+    // dev_ctx->SetHostZeroAllocator(phi::memory_utils::GetHostZeroAllocator());
+    // dev_ctx->SetPinnedAllocator(phi::memory_utils::GetPinnedAllocator());
+    // dev_ctx->PartialInitWithAllocator();
+    // auto compute_event =
+    //     phi::memory_utils::GetCudaEvent(CommContextManager::device_id);
+    // auto comm_event =
+    //     phi::memory_utils::GetCudaEvent(CommContextManager::device_id);
+
+    xccl_comm_context->SetDevContext(std::move(dev_ctx));
+    // xccl_comm_context->SetComputeEvent(std::move(compute_event));
+    // xccl_comm_context->SetCommEvent(std::move(comm_event));
+  }
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(xccl_comm_context));
+}
+#endif
+
+#if defined(PADDLE_WITH_XPU_BKCL)
+void CommContextManager::CreateBKCLCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    int rank,
+    int size,
+    const std::string& hash_key) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  if (comm_context_manager.Has(unique_comm_key)) {
+    return;
+  }
+#if defined(PADDLE_WITH_FLAGCX)
+  flagcxHandlerGroup_t flagcx_handler;
+  phi::dynload::flagcxHandleInit(&flagcx_handler);
+  if (rank == 0) {
+    phi::dynload::flagcxGetUniqueId(&flagcx_handler->uniqueId);
+  }
+#else
+  BKCLUniqueId bkcl_id;
+  if (rank == 0) {
+    PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_get_unique_id(&bkcl_id));
+  }
+#endif
+
+  std::string unique_key = "BKCLCommContext/" + unique_comm_key + hash_key;
+#if defined(PADDLE_WITH_FLAGCX)
+  if (rank == 0) {
+    std::vector<uint8_t> bkcl_id_wrapper(
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId) +
+            sizeof(flagcxUniqueId));
+    store->set(unique_key, bkcl_id_wrapper);
+  } else {
+    const auto& bkcl_id_wrapper = store->get(unique_key);
+    std::memcpy(reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+                bkcl_id_wrapper.data(),
+                bkcl_id_wrapper.size());
+  }
+#else
+  if (rank == 0) {
+    std::vector<uint8_t> bkcl_id_wrapper(
+        reinterpret_cast<uint8_t*>(&bkcl_id),
+        reinterpret_cast<uint8_t*>(&bkcl_id) + BKCL_UNIQUE_ID_BYTES);
+    store->set(unique_key, bkcl_id_wrapper);
+  } else {
+    const auto& bkcl_id_wrapper = store->get(unique_key);
+    std::memcpy(&bkcl_id, bkcl_id_wrapper.data(), bkcl_id_wrapper.size());
+  }
+#endif
+
+  VLOG(3) << "init BKCLCommContext rank: " << rank << ", size: " << size
+          << ", unique_comm_key: " << unique_comm_key
+          << ", unique_key: " << unique_key;
+#if defined(PADDLE_WITH_FLAGCX)
+  auto bkcl_comm_context =
+      std::make_unique<BKCLCommContext>(rank, size, flagcx_handler);
+#else
+  auto bkcl_comm_context =
+      std::make_unique<BKCLCommContext>(rank, size, bkcl_id);
+#endif
+
+  if (CommContextManager::device_id != -1) {
+    std::unique_ptr<phi::XPUContext> dev_ctx(new phi::XPUContext(
+        phi::XPUPlace(CommContextManager::device_id), true));
+    dev_ctx->CreateStream();
+    dev_ctx->SetAllocator(phi::memory_utils::GetAllocator(
+        CommContextManager::device_id, dev_ctx->stream()));
+    dev_ctx->SetHostAllocator(phi::memory_utils::GetHostAllocator());
+    dev_ctx->SetZeroAllocator(
+        phi::memory_utils::GetZeroAllocator(CommContextManager::device_id));
+    dev_ctx->SetHostZeroAllocator(phi::memory_utils::GetHostZeroAllocator());
+    // XPUs do not have the concept of pinned memory,
+    // so the get_pinned_allocator function is not set.
+
+    // It currently does not support dev_ctx->PartialInitWithAllocator().
+    auto compute_event =
+        phi::memory_utils::GetXpuEvent(CommContextManager::device_id);
+    auto comm_event =
+        phi::memory_utils::GetXpuEvent(CommContextManager::device_id);
+
+    bkcl_comm_context->SetDevContext(std::move(dev_ctx));
+    bkcl_comm_context->SetComputeEvent(std::move(compute_event));
+    bkcl_comm_context->SetCommEvent(std::move(comm_event));
+  }
+
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(bkcl_comm_context));
+}
+#endif
+
+#if defined(PADDLE_WITH_FLAGCX) && !defined(PADDLE_WITH_XPU)
+void CommContextManager::CreateFlagcxCommContext(
+    const std::shared_ptr<Store>& store,
+    const std::string& unique_comm_key,
+    int rank,
+    int size,
+    const std::string& hash_key) {
+  auto& comm_context_manager = CommContextManager::GetInstance();
+  if (comm_context_manager.Has(unique_comm_key)) {
+    return;
+  }
+  flagcxHandlerGroup_t flagcx_handler;
+  phi::dynload::flagcxHandleInit(&flagcx_handler);
+  if (rank == 0) {
+    phi::dynload::flagcxGetUniqueId(&flagcx_handler->uniqueId);
+  }
+
+  std::string unique_key = "XCCLCommContext/" + unique_comm_key + hash_key;
+  if (rank == 0) {
+    std::vector<uint8_t> flagcx_id_wrapper(
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+        reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId) +
+            sizeof(flagcxUniqueId));
+    store->set(unique_key, flagcx_id_wrapper);
+  } else {
+    const auto& flagcx_id_wrapper = store->get(unique_key);
+    std::memcpy(reinterpret_cast<uint8_t*>(flagcx_handler->uniqueId),
+                flagcx_id_wrapper.data(),
+                flagcx_id_wrapper.size());
+  }
+
+  VLOG(3) << "init FlagcxCommContext rank: " << rank << ", size: " << size
+          << ", unique_comm_key: " << unique_comm_key
+          << ", unique_key: " << unique_key << ", flagcx_id: "
+          << SerializeFlagcxUniqueId(*flagcx_handler->uniqueId);
+  auto flagcx_comm_context =
+      std::make_unique<FlagcxCommContext>(rank, size, flagcx_handler);
+  // TODO(changtao): find a way to manage different device context,
+  //  now we use cuda device context as default
+  comm_context_manager.SetStore(store);
+  comm_context_manager.Emplace(unique_comm_key, std::move(flagcx_comm_context));
+}
+#endif
+
+CommContext* CommContextManager::Emplace(
+    const std::string& unique_comm_key,
+    std::unique_ptr<CommContext> comm_context) {
+  PADDLE_ENFORCE_EQ(
+      id_to_comm_context_.find(unique_comm_key),
+      id_to_comm_context_.end(),
+      errors::AlreadyExists("The unique key %s already exists in the map.",
+                            unique_comm_key));
+  id_to_comm_context_.emplace(unique_comm_key, std::move(comm_context));
+  return id_to_comm_context_.at(unique_comm_key).get();
+}
+
+CommContext* CommContextManager::Get(const std::string& unique_comm_key) const {
+  PADDLE_ENFORCE_NE(
+      id_to_comm_context_.find(unique_comm_key),
+      id_to_comm_context_.end(),
+      errors::NotFound("Can not find unique key %s in map.", unique_comm_key));
+
+  return id_to_comm_context_.at(unique_comm_key).get();
+}
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+int CommContextManager::GetRingId(const ncclComm_t& comm) const {
+  for (const auto& iter : id_to_comm_context_) {
+    if (static_cast<phi::distributed::NCCLCommContext*>(iter.second.get())
+            ->GetNcclComm() == comm) {
+      return std::stoi(iter.first);
+    }
+  }
+  return -1;
+}
+#endif
+
+bool CommContextManager::Has(const std::string& unique_comm_key) const {
+  return id_to_comm_context_.find(unique_comm_key) != id_to_comm_context_.end();
+}
+
+void CommContextManager::SetGroupSize(const std::string& pg_key, int size) {
+  pg_key_size_[pg_key] = size;
+}
+
+void CommContextManager::AddGroupRanks(const std::string& pg_key,
+                                       std::vector<int> global_ranks) {
+  if (pg_key_ranks_.find(pg_key) == pg_key_ranks_.end()) {
+    pg_key_ranks_[pg_key] = global_ranks;
+  }
+}
+
+std::vector<int> CommContextManager::GetGroupRanks(
+    const std::string& pg_key) const {
+  PADDLE_ENFORCE_NE(
+      pg_key_ranks_.find(pg_key),
+      pg_key_ranks_.end(),
+      errors::NotFound("Can not find pg_key %d in GroupRanks.", pg_key));
+  return pg_key_ranks_.at(pg_key);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/comm_context_manager.h

@@ -0,0 +1,144 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <iostream>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/gpu/forwards.h"
+#include "paddle/phi/core/distributed/nccl_config.h"
+#endif
+
+#if defined(PADDLE_WITH_FLAGCX)
+#include <flagcx.h>
+#endif
+
+namespace phi {
+namespace distributed {
+
+struct P2POption {
+  bool is_p2p_op;
+  int p2p_rank;
+  int num_ranks;
+  int rank;
+};
+
+class Store;
+
+class CommContextManager {
+ public:
+  CommContextManager() = default;
+  ~CommContextManager() = default;
+
+  static CommContextManager& GetInstance() {
+    static CommContextManager instance;
+    return instance;
+  }
+
+  void SetStore(const std::shared_ptr<Store>& store) { store_ = store; }
+
+  CommContext* Emplace(const std::string& unique_comm_key,
+                       std::unique_ptr<CommContext> comm_context);
+
+  PADDLE_API CommContext* Get(const std::string& unique_comm_key) const;
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+  int GetRingId(const ncclComm_t& comm) const;
+#endif
+
+  PADDLE_API bool Has(const std::string& unique_comm_key) const;
+
+  PADDLE_API static void SetDeviceId(int dev_id);
+
+  PADDLE_API void SetGroupSize(const std::string& pg_key, int size);
+
+  PADDLE_API void AddGroupRanks(const std::string& pg_key,
+                                std::vector<int> global_ranks);
+
+  PADDLE_API std::vector<int> GetGroupRanks(const std::string& pg_key) const;
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+  static void CreateNCCLCommContext(
+      const std::shared_ptr<Store>& store,
+      const std::string& unique_comm_key,
+      int rank,
+      int size,
+      const std::string& hash_key = "",
+      const P2POption* opt = nullptr,
+      int nccl_comm_init_option = 0,
+      std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr = nullptr);
+  static void RecreateNCCLComm(const std::shared_ptr<Store>& store,
+                               const std::string& unique_comm_key,
+                               int rank,
+                               const std::string& hash_key = "",
+                               const P2POption* opt = nullptr);
+#endif
+
+#if defined(PADDLE_WITH_GLOO)
+  static void CreateGlooCommContext(const std::shared_ptr<Store>& store,
+                                    const std::string& unique_comm_key,
+                                    int rank,
+                                    int size);
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  static void CreateXCCLCommContext(const std::shared_ptr<Store>& store,
+                                    const std::string& unique_comm_key,
+                                    const Place& place,
+                                    int rank,
+                                    int size,
+                                    const std::string& hash_key = "");
+#endif
+
+#if defined(PADDLE_WITH_XPU_BKCL)
+  static void CreateBKCLCommContext(const std::shared_ptr<Store>& store,
+                                    const std::string& unique_comm_key,
+                                    int rank,
+                                    int size,
+                                    const std::string& hash_key = "");
+#endif
+
+#if defined(PADDLE_WITH_FLAGCX)
+  static void CreateFlagcxCommContext(const std::shared_ptr<Store>& store,
+                                      const std::string& unique_comm_key,
+                                      int rank,
+                                      int size,
+                                      const std::string& hash_key = "");
+#endif
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(CommContextManager);
+
+  std::unordered_map<std::string, std::unique_ptr<CommContext>>
+      id_to_comm_context_;
+  std::shared_ptr<Store> store_;
+  static int device_id;
+
+  // process group key to global ranks map
+  std::unordered_map<std::string, std::vector<int>> pg_key_ranks_;
+  // process group key to group size map
+  std::unordered_map<std::string, int> pg_key_size_;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/comm_task.h

@@ -0,0 +1,182 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <chrono>
+#include <condition_variable>
+#include <cstdint>
+#include <exception>
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/enforce.h"
+
+#if defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/dynload/rccl.h"
+#endif
+#if defined(PADDLE_WITH_NCCL)
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+namespace phi {
+namespace distributed {
+
+class Store;
+class CommTask {
+ public:
+  CommTask(const std::string& backend = "",
+           const Place& place = Place(),
+           const std::string& group_key = "",
+           int rank = -1,
+           int size = 0,
+           int gid = 0,
+           uint64_t seq = 0,
+           int64_t numel = 0,
+           ncclComm_t nccl_comm = nullptr,
+           gpuStream_t nccl_stream = nullptr,
+           CommType comm_type = CommType::UNKNOWN)
+      : backend_(backend),
+        place_(place),
+        group_key_(group_key),
+        rank_(rank),
+        size_(size),
+        gid_(gid),
+        seq_(seq),
+        numel_(numel),
+        nccl_comm_(nccl_comm),
+        nccl_stream_(nccl_stream),
+        comm_type_(comm_type) {
+    const char* global_rank = std::getenv("PADDLE_TRAINER_ID");
+    PADDLE_ENFORCE_NOT_NULL(
+        global_rank,
+        common::errors::NotFound(
+            "The environment variable 'PADDLE_TRAINER_ID' cannot be found."));
+    global_rank_ = std::atoi(global_rank);
+  }
+  virtual ~CommTask() = default;
+
+  std::string UniqueKey() {
+    return "group_key:" + group_key_ + ",op:" + CommTypeToString(comm_type_) +
+           ",gid:" + std::to_string(gid_) + ",seq:" + std::to_string(seq_);
+  }
+
+  std::string GroupKey() { return group_key_; }
+  std::string GetBackend() { return backend_; }
+  Place GetPlace() { return place_; }
+  int GetGlobalRank() { return global_rank_; }
+  int GetRank() { return rank_; }
+  int GetSize() { return size_; }
+  int GetGid() { return gid_; }
+  int64_t GetNumel() { return numel_; }
+  uint64_t GetSeq() { return seq_; }
+  CommType GetCommType() { return comm_type_; }
+  bool GetTraceUpdated() { return start_trace_updated_; }
+  void SetTraceUpdated() { start_trace_updated_ = true; }
+  std::chrono::time_point<std::chrono::steady_clock> GetStartTime() {
+    return start_time_;
+  }
+  std::shared_ptr<Store> GetStore() { return store_; }
+  void SetStore(std::shared_ptr<Store> store) { store_ = store; }
+
+  ncclComm_t nccl_comm() { return nccl_comm_; }
+  gpuStream_t nccl_stream() { return nccl_stream_; }
+
+  virtual std::string GetTraceMsg() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return "";
+  }
+  virtual void StartRecord() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return;
+  }
+  virtual void EndRecord() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return;
+  }
+
+  virtual void ClearRecord() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return;
+  }
+
+  virtual std::string GetCommErrors() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return "";
+  }
+  virtual bool IsStarted() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return false;
+  }
+  virtual bool IsTimeout() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return false;
+  }
+  virtual bool IsCompleted() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return false;
+  }
+  virtual void SetUpdated(bool updated) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return;
+  }
+  virtual bool IsUpdated() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return false;
+  }
+  virtual void AbortComm() {
+    PADDLE_THROW(
+        common::errors::Unimplemented("%s is not implemented.", __func__));
+    return;
+  }
+
+ protected:
+  std::string backend_;
+  Place place_;
+  std::string group_key_;
+  int global_rank_;
+  int rank_;
+  int size_;
+  int gid_;
+  uint64_t seq_{0};
+  int64_t numel_;
+  ncclComm_t nccl_comm_;
+  gpuStream_t nccl_stream_;
+  CommType comm_type_;
+  bool start_trace_updated_{false};
+
+  // task status
+  bool started_ = false;
+  bool completed_ = false;
+  // task status changed
+  bool updated_ = true;
+  bool aborted_{false};
+  std::chrono::time_point<std::chrono::steady_clock> start_time_;
+  std::shared_ptr<Store> store_;
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(CommTask);
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/comm_task_manager.cc

@@ -0,0 +1,278 @@
+// Copyright (c) 2023 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.
+
+#if defined(PADDLE_WITH_GLOO)
+#include <gloo/rendezvous/prefix_store.h>
+
+#include "paddle/phi/core/distributed/gloo_comm_context.h"
+#include "paddle/phi/core/distributed/gloo_utils.h"
+#include "paddle/phi/core/distributed/store/gloo_store.h"
+#endif
+
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+
+#include <future>
+#include <memory>
+#include <string>
+
+#include "gflags/gflags.h"
+#include "glog/logging.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/core/distributed/store/store.h"
+#include "paddle/phi/core/enforce.h"
+
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/core/distributed/comm_task_manager.h"
+#include "paddle/phi/core/distributed/nccl_comm_context.h"
+#endif
+
+namespace phi::distributed {
+
+std::thread CommTaskManager::comm_task_loop_thread_;
+std::thread CommTaskManager::comm_task_clear_loop_thread_;
+const int64_t CommTaskManager::loop_thread_sleep_millis = 10000;
+
+std::atomic<bool> CommTaskManager::terminated_;
+std::mutex CommTaskManager::comm_task_list_mutex_;
+std::condition_variable CommTaskManager::comm_task_list_cv_;
+std::list<std::shared_ptr<CommTask>> CommTaskManager::comm_task_list_;
+
+std::mutex CommTaskManager::comm_task_clear_list_mutex_;
+std::condition_variable CommTaskManager::comm_task_clear_list_cv_;
+std::list<std::shared_ptr<CommTask>> CommTaskManager::comm_task_clear_list_;
+
+std::unordered_map<std::string, std::shared_ptr<CommTask>>
+    CommTaskManager::init_comm_task_map_;
+std::unordered_map<std::string, std::shared_ptr<CommTask>>
+    CommTaskManager::start_comm_task_map_;
+std::unordered_map<std::string, std::shared_ptr<CommTask>>
+    CommTaskManager::group_last_comm_task_;
+std::chrono::time_point<std::chrono::steady_clock>
+    CommTaskManager::last_update_time_ = std::chrono::steady_clock::now();
+
+CommTaskManager::CommTaskManager() : timeout_(0) {
+  terminated_.store(false);
+  comm_task_loop_thread_ = std::thread(&CommTaskManager::CommTaskLoop, this);
+  comm_task_clear_loop_thread_ =
+      std::thread(&CommTaskManager::CommTaskClearLoop, this);
+  LOG(INFO) << "CommTaskManager init success.";
+}
+CommTaskManager::~CommTaskManager() {
+  terminated_.store(true);
+
+  if (comm_task_loop_thread_.joinable()) {
+    comm_task_list_cv_.notify_one();
+    comm_task_loop_thread_.join();
+  }
+
+  if (comm_task_clear_loop_thread_.joinable()) {
+    comm_task_clear_list_cv_.notify_one();
+    comm_task_clear_loop_thread_.join();
+  }
+  LOG(INFO) << "CommTaskManager destruct success.";
+}
+
+void CommTaskManager::CommTaskEnqueue(std::shared_ptr<CommTask> comm_task) {
+  if (!terminated_.load()) {
+    std::lock_guard<std::mutex> lock(comm_task_list_mutex_);
+    comm_task_list_.emplace_back(std::move(comm_task));
+  }
+}
+
+void CommTaskManager::CommTaskClearEnqueue(
+    std::shared_ptr<CommTask> comm_task) {
+  if (!terminated_.load()) {
+    std::lock_guard<std::mutex> lock(comm_task_clear_list_mutex_);
+    comm_task_clear_list_.emplace_back(comm_task);
+  }
+}
+
+void CommTaskManager::Stop() {
+  terminated_.store(true);
+
+  LOG(INFO) << "CommTaskManager stopped begin.";
+  if (comm_task_loop_thread_.joinable()) {
+    comm_task_list_cv_.notify_one();
+    comm_task_loop_thread_.join();
+  }
+
+  if (comm_task_clear_loop_thread_.joinable()) {
+    comm_task_clear_list_cv_.notify_one();
+    comm_task_clear_loop_thread_.join();
+  }
+
+  LOG(INFO) << "CommTaskManager stopped.";
+}
+
+inline void LogLongStr(const std::string prefix, const std::string& log) {
+  size_t max_log_size = 20000;
+  if (log.size() >= max_log_size) {
+    int log_count = log.size() / max_log_size + 1;
+    int index = 0;
+    int part = 0;
+    while (index + max_log_size < log.size()) {
+      LOG(INFO) << prefix << "part:" << part << "/" << log_count << ","
+                << log.substr(index, max_log_size) << std::endl;
+      index += max_log_size;
+      part++;
+    }
+    LOG(INFO) << prefix << "part:" << part << "/" << log_count << ","
+              << log.substr(index) << std::endl;
+  } else {
+    LOG(INFO) << prefix << "part:0/1," << log << std::endl;
+  }
+}
+
+void CommTaskManager::CommTaskLoop() {
+  bool done = false;
+  while (!terminated_.load() || !done) {
+    std::unique_lock<std::mutex> lock(comm_task_list_mutex_);
+    VLOG(3) << "IsTimeout: " << IsTimeout()
+            << ", comm_task_list_ size: " << comm_task_list_.size()
+            << ", init_comm_task_map_ size: " << init_comm_task_map_.size()
+            << ", start_comm_task_map_ size: " << start_comm_task_map_.size()
+            << ", logged_ " << logged_;
+
+    comm_task_list_cv_.wait_for(
+        lock,
+        std::chrono::milliseconds(loop_thread_sleep_millis),
+        [&]() -> bool { return terminated_.load(); });
+
+    if (IsTimeout() && !logged_) {
+      // case 1: all group is empty, has no task
+      // report error immediately
+      if (group_last_comm_task_.empty()) {
+        LOG(WARNING) << "Find no task started in all group";
+      } else {
+        // case 2: all group is not empty, but all last task is completed
+        // case 3: all group is not empty, some group task started but not
+        for (const auto& iter : group_last_comm_task_) {
+          LogLongStr("Find last group comm task:", iter.second->GetTraceMsg());
+        }
+      }
+      logged_ = true;
+    }
+    for (auto iter = comm_task_list_.begin(); iter != comm_task_list_.end();) {
+      auto task = *iter;
+      if (task->IsTimeout()) {
+        if (!task->IsStarted()) {
+          LOG(WARNING) << "Find timeout init but not start task:"
+                       << task->GetTraceMsg();
+          std::string task_key = task->UniqueKey();
+          init_comm_task_map_[task_key] = task;
+        } else if (!task->IsCompleted()) {
+          LOG(WARNING) << "Find timeout start but not finish task:"
+                       << task->GetTraceMsg();
+          std::string task_key = task->UniqueKey();
+          start_comm_task_map_[task_key] = task;
+        }
+        iter = comm_task_list_.erase(iter);
+      } else {
+        if (task->IsStarted()) {
+          if (task->IsCompleted()) {
+            CommTaskClearEnqueue(task);
+            iter = comm_task_list_.erase(iter);
+          } else {
+            ++iter;
+          }
+          UpdateLastCommTask(task);
+        } else {
+          ++iter;
+        }
+      }
+    }
+
+    for (auto iter = init_comm_task_map_.begin();
+         iter != init_comm_task_map_.end();) {
+      auto task = iter->second;
+      if (task->IsStarted()) {
+        std::string task_key = task->UniqueKey();
+        start_comm_task_map_[task_key] = task;
+        iter = init_comm_task_map_.erase(iter);
+        LOG(INFO) << "Start timeout task: " << task->GetTraceMsg();
+      } else {
+        ++iter;
+      }
+    }
+
+    for (auto iter = start_comm_task_map_.begin();
+         iter != start_comm_task_map_.end();) {
+      auto task = iter->second;
+      if (task->IsCompleted()) {
+        CommTaskClearEnqueue(task);
+        UpdateLastCommTask(task);
+        iter = start_comm_task_map_.erase(iter);
+        LOG(INFO) << "Finish timeout task: " << task->GetTraceMsg();
+      } else {
+        ++iter;
+      }
+    }
+
+    if (comm_task_list_.empty() && init_comm_task_map_.empty() &&
+        start_comm_task_map_.empty()) {
+      done = true;
+    } else {
+      done = false;
+    }
+  }
+}
+
+void CommTaskManager::CommTaskClearLoop() {
+  std::future<void> future;
+  while (!terminated_.load()) {
+    if (future.valid()) {
+      future.wait();
+    }
+    std::unique_lock<std::mutex> lock(comm_task_clear_list_mutex_);
+    comm_task_clear_list_cv_.wait_for(
+        lock,
+        std::chrono::milliseconds(loop_thread_sleep_millis),
+        [&]() -> bool { return terminated_.load(); });
+
+    VLOG(3) << "comm_task_clear_list_ size: " << comm_task_clear_list_.size();
+    for (auto iter = comm_task_clear_list_.begin();
+         iter != comm_task_clear_list_.end();) {
+      auto task = *iter;
+      VLOG(3) << "start clear task: " << task->GetTraceMsg();
+      future = std::async(std::launch::async, [&]() { task->ClearRecord(); });
+      if (future.wait_for(std::chrono::seconds(30)) ==
+          std::future_status::timeout) {
+        VLOG(0) << "clear task timeout, detail: " << task->GetTraceMsg();
+        break;
+      }
+      VLOG(3) << "end clear task: " << task->GetTraceMsg();
+      iter = comm_task_clear_list_.erase(iter);
+    }
+  }
+}
+
+void CommTaskManager::UpdateLastCommTask(std::shared_ptr<CommTask> task) {
+  if (!task->IsUpdated()) {
+    return;
+  }
+  group_last_comm_task_[task->GroupKey()] = task;
+  last_update_time_ = std::chrono::steady_clock::now();
+  task->SetUpdated(false);
+}
+
+void CommTaskManager::SetTimeout(int64_t timeout) {
+  timeout_ = std::chrono::milliseconds(timeout);
+}
+
+bool CommTaskManager::IsTimeout() {
+  auto current_timepoint = std::chrono::steady_clock::now();
+  return std::chrono::duration_cast<std::chrono::milliseconds>(
+             current_timepoint - last_update_time_) >= timeout_;
+}
+}  // namespace phi::distributed

paddle/phi/core/distributed/comm_task_manager.h

@@ -0,0 +1,89 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <iostream>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <thread>
+#include <unordered_map>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+#include "paddle/phi/core/distributed/comm_task.h"
+
+namespace phi {
+namespace distributed {
+
+enum ErrorHandlingMode { NoHandling = 0, TearDown = 1 };
+
+class Store;
+
+class CommTaskManager {
+ public:
+  CommTaskManager();
+  ~CommTaskManager();
+
+ public:
+  static CommTaskManager& GetInstance() {
+    static CommTaskManager instance;
+    return instance;
+  }
+
+  void CommTaskEnqueue(std::shared_ptr<CommTask> comm_task);
+  void CommTaskClearEnqueue(std::shared_ptr<CommTask> comm_task);
+  void Stop();
+  void UpdateLastCommTask(std::shared_ptr<CommTask> comm_task);
+  void SetTimeout(int64_t timeout);
+
+ private:
+  void CommTaskLoop();
+  void CommTaskClearLoop();
+  bool IsTimeout();
+
+  static std::thread comm_task_loop_thread_;
+  static std::thread comm_task_clear_loop_thread_;
+  static const int64_t loop_thread_sleep_millis;
+
+  static std::atomic<bool> terminated_;
+
+  static std::mutex comm_task_list_mutex_;
+  static std::condition_variable comm_task_list_cv_;
+  static std::list<std::shared_ptr<CommTask>> comm_task_list_;
+
+  static std::mutex comm_task_clear_list_mutex_;
+  static std::condition_variable comm_task_clear_list_cv_;
+  static std::list<std::shared_ptr<CommTask>> comm_task_clear_list_;
+
+  // not start task
+  static std::unordered_map<std::string, std::shared_ptr<CommTask>>
+      init_comm_task_map_;
+  // start but not finish task
+  static std::unordered_map<std::string, std::shared_ptr<CommTask>>
+      start_comm_task_map_;
+  std::shared_ptr<Store> store_;
+  // record last comm task in current group, eg: group_key->comm_task
+  static std::unordered_map<std::string, std::shared_ptr<CommTask>>
+      group_last_comm_task_;
+  static std::chrono::time_point<std::chrono::steady_clock> last_update_time_;
+  std::chrono::milliseconds timeout_;
+  bool logged_ = false;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/flagcx_comm_context.cc

@@ -0,0 +1,198 @@
+// Copyright (c) 2025 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.
+
+#include "paddle/phi/core/distributed/flagcx_comm_context.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+// #include "paddle/phi/core/distributed/check/nccl_dynamic_check.h"
+#include "paddle/phi/core/distributed/check/static_check.h"
+#include "paddle/phi/core/distributed/flagcx_tools.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+#ifdef PADDLE_WITH_FLAGCX
+#include <flagcx.h>
+#endif
+
+namespace phi::distributed {
+
+// set this flag to `true` and recompile to enable dynamic checks
+// constexpr bool FLAGS_enable_flagcx_dynamic_check = false;
+
+FlagcxCommContext::FlagcxCommContext(int rank,
+                                     int size,
+                                     flagcxHandlerGroup_t flagcx_handler)
+    : CommContext(rank, size),
+      flagcx_version_(0),
+      flagcx_handler_(flagcx_handler) {
+  phi::dynload::flagcxCommInitRank(
+      &flagcx_handler_->comm, size_, flagcx_handler_->uniqueId, rank_),
+      phi::dynload::flagcxGetVersion(&flagcx_version_);
+}
+
+int FlagcxCommContext::GetFlagcxVersion() { return flagcx_version_; }
+
+flagcxComm_t FlagcxCommContext::GetFlagcxComm() {
+  return flagcx_handler_->comm;
+}
+
+void FlagcxCommContext::Broadcast(DenseTensor* out_tensor,
+                                  const DenseTensor& in_tensor,
+                                  int root,
+                                  flagcxStream_t stream) {
+  CommStaticCheck::SameShape(*out_tensor,
+                             in_tensor,
+                             /*dst_rank*/ rank_,
+                             /*cur_rank*/ rank_,
+                             size_);
+  FLAGCX_CHECK(phi::dynload::flagcxBroadcast(in_tensor.data(),
+                                             out_tensor->data(),
+                                             in_tensor.numel(),
+                                             ToFlagcxDataType(in_tensor.type()),
+                                             root,
+                                             flagcx_handler_->comm,
+                                             stream));
+}
+
+void FlagcxCommContext::AllGather(DenseTensor* out_tensor,
+                                  const DenseTensor& in_tensor,
+                                  flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::GatherLikeShape(*out_tensor,
+                                                     in_tensor,
+                                                     /*dst_rank*/ rank_,
+                                                     /*cur_rank*/ rank_,
+                                                     size_);
+  FLAGCX_CHECK(phi::dynload::flagcxAllGather(in_tensor.data(),
+                                             out_tensor->data(),
+                                             in_tensor.numel(),
+                                             ToFlagcxDataType(in_tensor.type()),
+                                             flagcx_handler_->comm,
+                                             stream));
+}
+void FlagcxCommContext::ReduceScatter(DenseTensor* out_tensor,
+                                      const DenseTensor& in_tensor,
+                                      flagcxRedOp_t reduce_type,
+                                      flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::ScatterLikeShape(*out_tensor,
+                                                      in_tensor,
+                                                      /*dst_rank*/ rank_,
+                                                      /*cur_rank*/ rank_,
+                                                      size_);
+  FLAGCX_CHECK(
+      phi::dynload::flagcxReduceScatter(in_tensor.data(),
+                                        out_tensor->data(),
+                                        out_tensor->numel(),
+                                        ToFlagcxDataType(in_tensor.type()),
+                                        reduce_type,
+                                        flagcx_handler_->comm,
+                                        stream));
+}
+
+void FlagcxCommContext::Send(const DenseTensor& in_tensor,
+                             const int64_t& count,
+                             const int& peer,
+                             flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::CheckShape(in_tensor, rank_, size_);
+
+  FLAGCX_CHECK(phi::dynload::flagcxSend(in_tensor.data(),
+                                        count,
+                                        ToFlagcxDataType(in_tensor.dtype()),
+                                        peer,
+                                        flagcx_handler_->comm,
+                                        stream));
+  VLOG(3) << "rank " << GetRank() << " send "
+          << common::product(in_tensor.dims()) << " to " << peer;
+}
+
+void FlagcxCommContext::Recv(DenseTensor* out_tensor,
+                             const int64_t& count,
+                             const int& peer,
+                             flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::CheckShape(*out_tensor, rank_, size_);
+
+  FLAGCX_CHECK(phi::dynload::flagcxRecv(out_tensor->data(),
+                                        count,
+                                        ToFlagcxDataType(out_tensor->dtype()),
+                                        peer,
+                                        flagcx_handler_->comm,
+                                        stream));
+  VLOG(3) << "rank " << GetRank() << " recv "
+          << common::product(out_tensor->dims()) << " from " << peer;
+}
+
+void FlagcxCommContext::AllReduce(DenseTensor* out_tensor,
+                                  const DenseTensor& in_tensor,
+                                  flagcxRedOp_t reduce_type,
+                                  flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ rank_,
+                                               /*cur_rank*/ rank_,
+                                               size_);
+  FLAGCX_CHECK(phi::dynload::flagcxAllReduce(in_tensor.data(),
+                                             out_tensor->data(),
+                                             in_tensor.numel(),
+                                             ToFlagcxDataType(in_tensor.type()),
+                                             reduce_type,
+                                             flagcx_handler_->comm,
+                                             stream));
+}
+
+void FlagcxCommContext::Reduce(DenseTensor* out_tensor,
+                               const DenseTensor& in_tensor,
+                               flagcxRedOp_t reduce_type,
+                               int root,
+                               flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ root,
+                                               /*cur_rank*/ rank_,
+                                               size_);
+  FLAGCX_CHECK(phi::dynload::flagcxReduce(in_tensor.data(),
+                                          out_tensor->data(),
+                                          in_tensor.numel(),
+                                          ToFlagcxDataType(in_tensor.type()),
+                                          reduce_type,
+                                          root,
+                                          flagcx_handler_->comm,
+                                          stream));
+}
+
+void FlagcxCommContext::AllToAll(DenseTensor* out_tensor,
+                                 const DenseTensor& in_tensor,
+                                 flagcxStream_t stream) {
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ rank_,
+                                               /*cur_rank*/ rank_,
+                                               size_);
+  FLAGCX_CHECK(phi::dynload::flagcxAlltoAll(in_tensor.data(),
+                                            out_tensor->data(),
+                                            in_tensor.numel() / size_,
+                                            ToFlagcxDataType(in_tensor.type()),
+                                            flagcx_handler_->comm,
+                                            stream));
+}
+
+void FlagcxCommContext::GroupStart() {
+  FLAGCX_CHECK(phi::dynload::flagcxGroupStart(flagcx_handler_->comm));
+}
+void FlagcxCommContext::GroupEnd() {
+  FLAGCX_CHECK(phi::dynload::flagcxGroupEnd(flagcx_handler_->comm));
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/flagcx_comm_context.h

@@ -0,0 +1,96 @@
+// Copyright (c) 2025 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.
+#pragma once
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/dynload/flagcx.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+
+namespace phi {
+class DenseTensor;
+namespace distributed {
+
+class FlagcxCommContext final : public CommContext {
+ public:
+  FlagcxCommContext(int rank, int size, flagcxHandlerGroup_t flagcx_handler);
+  ~FlagcxCommContext() override = default;
+
+  int GetFlagcxVersion();
+
+  flagcxComm_t GetFlagcxComm();
+
+  void Broadcast(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 int root,
+                 flagcxStream_t stream);
+
+  void Send(const DenseTensor& in_tensor,
+            const int64_t& count,
+            const int& peer,
+            flagcxStream_t stream);
+
+  void Recv(DenseTensor* out_tensor,
+            const int64_t& count,
+            const int& peer,
+            flagcxStream_t stream);
+
+  void ReduceScatter(DenseTensor* out_tensor,
+                     const DenseTensor& in_tensor,
+                     flagcxRedOp_t reduce_type,
+                     flagcxStream_t stream);
+
+  void AllGather(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 flagcxStream_t stream);
+
+  void AllReduce(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 flagcxRedOp_t reduce_type,
+                 flagcxStream_t stream);
+
+  void Reduce(DenseTensor* out_tensor,
+              const DenseTensor& in_tensor,
+              flagcxRedOp_t reduce_type,
+              int root,
+              flagcxStream_t stream);
+
+  void AllToAll(DenseTensor* out_tensor,
+                const DenseTensor& in_tensor,
+                flagcxStream_t stream);
+
+  void GroupStart();
+
+  void GroupEnd();
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(FlagcxCommContext);
+
+  int flagcx_version_;
+
+  std::unique_ptr<phi::GPUContext> dev_ctx_;
+
+  // used for comm wait compute, compute_stream-->event-->comm_stream
+  std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type> compute_event_;
+
+  // used for compute wait comm, comm_stream-->event-->compute_stream
+  std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type> comm_event_;
+
+ public:
+  flagcxHandlerGroup_t flagcx_handler_;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/flagcx_tools.cc

@@ -0,0 +1,86 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/flagcx_tools.h"
+
+#include <unordered_map>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace distributed {
+
+flagcxRedOp_t ToFlagcxRedType(ReduceOp reduction) {
+  static const std::unordered_map<ReduceOp, flagcxRedOp_t> red_type = {
+      {ReduceOp::MIN, flagcxMin},
+      {ReduceOp::MAX, flagcxMax},
+      {ReduceOp::SUM, flagcxSum},
+      {ReduceOp::PRODUCT, flagcxProd},
+      {ReduceOp::AVG, flagcxAvg},
+  };
+  auto it = red_type.find(reduction);
+  PADDLE_ENFORCE_EQ(
+      it != red_type.end(),
+      true,
+      common::errors::InvalidArgument(
+          "Invalid flagcx reduction. Must be flagcxMin | flagcxMax | "
+          "flagcxProd | flagcxSum | flagcxAvg."));
+  return it->second;
+}
+
+std::string SerializeFlagcxUniqueId(const flagcxUniqueId& flagcxID) {
+  const uint8_t* bytes = reinterpret_cast<const uint8_t*>(&flagcxID);
+  std::ostringstream oss;
+  for (auto i = 0; i < FLAGCX_UNIQUE_ID_BYTES; ++i) {
+    oss << std::hex << static_cast<int>(bytes[i]);
+  }
+  return oss.str();
+}
+
+std::string FlagcxDTypeToString(flagcxDataType_t dtype) {
+#define PD_FLAGCX_DTYPE_TO_STR(__flagcx_dtype, __str_dtype) \
+  if (dtype == __flagcx_dtype) return __str_dtype;
+  PD_FLAGCX_DTYPE_TO_STR(flagcxFloat, "float32");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxFloat32, "float32");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxHalf, "float16");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxFloat16, "float16");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxBfloat16, "bfloat16");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxDouble, "float64");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxFloat64, "float64");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxInt8, "int8");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxChar, "int8");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxUint8, "uint8");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxInt32, "int32");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxInt, "int32");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxUint32, "uint32");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxInt64, "int64");
+  PD_FLAGCX_DTYPE_TO_STR(flagcxUint64, "uint64");
+
+#undef PD_FLAGCX_DTYPE_TO_STR
+  PADDLE_THROW(common::errors::InvalidArgument(
+      "This datatype %d in flagcx is not supported.", static_cast<int>(dtype)));
+}
+
+std::string FlagcxRedTypeToString(flagcxRedOp_t op) {
+  if (op == flagcxSum) return "SUM";
+  if (op == flagcxProd) return "PROD";
+  if (op == flagcxMin) return "MIN";
+  if (op == flagcxMax) return "MAX";
+  if (op == flagcxAvg) return "AVG";
+  return "UDF_" + std::to_string(op);
+}
+
+}  //  namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/flagcx_tools.h

@@ -0,0 +1,45 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <flagcx.h>
+#include <string>
+#include "paddle/phi/core/distributed/types.h"
+
+namespace phi {
+namespace distributed {
+
+#define FLAGCX_CHECK(cmd)                                                   \
+  do {                                                                      \
+    flagcxResult_t r = cmd;                                                 \
+    if (r != flagcxSuccess) {                                               \
+      PADDLE_THROW(                                                         \
+          common::errors::External("Failed, FlagCX error %s:%d '%s'\n",     \
+                                   __FILE__,                                \
+                                   __LINE__,                                \
+                                   phi::dynload::flagcxGetErrorString(r))); \
+    }                                                                       \
+  } while (0)
+
+flagcxRedOp_t ToFlagcxRedType(ReduceOp reduction);
+
+std::string SerializeFlagcxUniqueId(const flagcxUniqueId& flagcxID);
+
+std::string FlagcxDTypeToString(flagcxDataType_t dtype);
+
+std::string FlagcxRedTypeToString(flagcxRedOp_t op);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/gloo_comm_context.cc

@@ -0,0 +1,168 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/gloo_comm_context.h"
+#include "paddle/phi/core/distributed/gloo_utils.h"
+
+#include <gloo/allgather.h>
+#include <gloo/allreduce.h>
+#include <gloo/barrier.h>
+#include <gloo/broadcast.h>
+#include <gloo/gather.h>
+#include <gloo/reduce.h>
+#include <gloo/scatter.h>
+#include <gloo/types.h>
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/distributed/check/static_check.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi::distributed {
+
+GlooCommContext::GlooCommContext(
+    int rank,
+    int size,
+    std::shared_ptr<gloo::rendezvous::Store> store,
+    std::shared_ptr<gloo::transport::Device> device)
+    : CommContext(rank, size) {
+  gloo_context_ = std::make_shared<gloo::rendezvous::Context>(rank, size);
+  gloo_context_->connectFullMesh(*store, device);
+}
+
+void GlooCommContext::Broadcast(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                int root,
+                                uint32_t tag) {
+  // gloo only uses CPU now
+  CommStaticCheck::SameShape(*out_tensor,
+                             in_tensor,
+                             /*dst_rank*/ rank_,
+                             /*cur_rank*/ rank_,
+                             size_,
+                             AllocationType::CPU);
+  gloo::BroadcastOptions opts(gloo_context_);
+  const auto& dtype = in_tensor.dtype();
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  if (rank_ == root) {
+    GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  }
+  opts.setRoot(root);
+  opts.setTag(tag);
+  gloo::broadcast(opts);
+}
+
+void GlooCommContext::AllGather(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                uint32_t tag) {
+  // gloo only uses CPU now
+
+  gloo::AllgatherOptions opts(gloo_context_);
+  const auto& dtype = in_tensor.dtype();
+  opts.setTag(tag);
+  GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  gloo::allgather(opts);
+}
+
+void GlooCommContext::AllReduce(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                int reduce_type,
+                                uint32_t tag) {
+  gloo::AllreduceOptions opts(gloo_context_);
+  opts.setTag(tag);
+  const auto& dtype = in_tensor.dtype();
+  GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  GENERATE_FUNC(dtype, SetReduceFunc, &opts, reduce_type);
+  gloo::allreduce(opts);
+}
+
+void GlooCommContext::Reduce(DenseTensor* out_tensor,
+                             const DenseTensor& in_tensor,
+                             int reduce_type,
+                             int root,
+                             uint32_t tag) {
+  gloo::ReduceOptions opts(gloo_context_);
+  opts.setRoot(root);
+  opts.setTag(tag);
+  const auto& dtype = in_tensor.dtype();
+  GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  GENERATE_FUNC(dtype, SetReduceFunc, &opts, reduce_type);
+  gloo::reduce(opts);
+}
+
+void GlooCommContext::Gather(DenseTensor* out_tensor,
+                             const DenseTensor& in_tensor,
+                             int src,
+                             uint32_t tag) {
+  gloo::GatherOptions opts(gloo_context_);
+  const auto& dtype = in_tensor.dtype();
+  opts.setTag(tag);
+  opts.setRoot(src);
+  GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  if (rank_ == src) {
+    GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  }
+  gloo::gather(opts);
+}
+
+void GlooCommContext::Scatter(DenseTensor* out_tensor,
+                              const DenseTensor& in_tensor,
+                              int src,
+                              int size,
+                              uint32_t tag) {
+  gloo::ScatterOptions opts(gloo_context_);
+  const auto& dtype = in_tensor.dtype();
+  if (rank_ == src) {
+    if (size == 0) {
+      size = size_;
+    }
+    GENERATE_FUNC(dtype, SetInputForScatter, &opts, in_tensor, size);
+  }
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  opts.setRoot(src);
+  opts.setTag(tag);
+  gloo::scatter(opts);
+}
+
+void GlooCommContext::Barrier() {
+  gloo::BarrierOptions opts(gloo_context_);
+  gloo::barrier(opts);
+}
+
+void GlooCommContext::Send(const DenseTensor& in_tensor,
+                           int dst,
+                           uint32_t tag) {
+  SendRecvOptions opts(gloo_context_);
+  const auto& dtype = in_tensor.dtype();
+  GENERATE_FUNC(dtype, SetInput, &opts, in_tensor);
+  opts.setSrc(gloo_context_->rank);
+  opts.setDst(dst);
+  opts.setTag(tag);
+  send_recv(&opts);
+}
+
+void GlooCommContext::Recv(DenseTensor* out_tensor, int src, uint32_t tag) {
+  SendRecvOptions opts(gloo_context_);
+  const auto& dtype = out_tensor->dtype();
+  GENERATE_FUNC(dtype, SetOutput, &opts, out_tensor);
+  opts.setSrc(src);
+  opts.setDst(gloo_context_->rank);
+  opts.setTag(tag);
+  send_recv(&opts);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/gloo_comm_context.h

@@ -0,0 +1,80 @@
+// Copyright (c) 2023 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.
+#pragma once
+
+#include <gloo/rendezvous/context.h>
+#include <gloo/rendezvous/store.h>
+#include <gloo/transport/tcp/device.h>
+
+#include <memory>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+
+namespace phi {
+class DenseTensor;
+namespace distributed {
+
+class GlooCommContext final : public CommContext {
+ public:
+  GlooCommContext(int rank,
+                  int size,
+                  std::shared_ptr<gloo::rendezvous::Store> store,
+                  std::shared_ptr<gloo::transport::Device> device);
+
+  void Broadcast(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 int root,
+                 uint32_t tag = 0);
+
+  void AllReduce(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 int reduce_type,
+                 uint32_t tag = 0);
+
+  void Reduce(DenseTensor* out_tensor,
+              const DenseTensor& in_tensor,
+              int reduce_type,
+              int root,
+              uint32_t tag = 0);
+
+  void AllGather(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 uint32_t tag = 0);
+
+  void Gather(DenseTensor* out_tensor,
+              const DenseTensor& in_tensor,
+              int src,
+              uint32_t tag = 0);
+
+  void Scatter(DenseTensor* out_tensor,
+               const DenseTensor& in_tensor,
+               int src,
+               int size = 0,
+               uint32_t tag = 0);
+
+  void Barrier();
+
+  void Send(const DenseTensor& in_tensor, int dst, uint32_t tag = 0);
+
+  void Recv(DenseTensor* out_tensor, int src, uint32_t tag = 0);
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(GlooCommContext);
+
+  std::shared_ptr<gloo::rendezvous::Context> gloo_context_;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/gloo_utils.cc

@@ -0,0 +1,108 @@
+// Copyright (c) 2023 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.
+
+#ifdef _WIN32
+#include <gloo/common/win.h>
+#include <winsock2.h>
+#include <ws2tcpip.h>
+#else
+#include <netdb.h>
+#include <sys/socket.h>
+#include <unistd.h>
+#include <array>
+#endif
+
+#include <cstdlib>
+#include <cstring>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/distributed/gloo_utils.h"
+#include "paddle/phi/core/distributed/store/tcp_utils.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi::distributed {
+std::shared_ptr<gloo::transport::Device> CreateDeviceForInterface(
+    const std::string& ifname) {
+  gloo::transport::tcp::attr attr;
+  attr.iface = ifname;
+  return gloo::transport::tcp::CreateDevice(attr);
+}
+
+std::shared_ptr<gloo::transport::Device> CreateDeviceForHostname(
+    const std::string& hostname) {
+  gloo::transport::tcp::attr attr;
+  attr.hostname = hostname;
+  return gloo::transport::tcp::CreateDevice(attr);
+}
+
+std::shared_ptr<gloo::transport::Device> CreateDefaultDevice() {
+  std::array<char, HOST_NAME_MAX> hostname = {};
+  auto ret = ::gethostname(hostname.data(), HOST_NAME_MAX);
+  PADDLE_ENFORCE_EQ(
+      ret,
+      0,
+      common::errors::Fatal("Get hostname error for createDefaultDevice."));
+  ::addrinfo* result;
+  result = phi::distributed::tcputils::get_addr_info(
+      hostname.data(), "", 0, AF_UNSPEC);
+  ::addrinfo* cur;
+  for (cur = result; cur != nullptr; cur = cur->ai_next) {
+    phi::distributed::SocketType socket =
+        ::socket(cur->ai_family, cur->ai_socktype, cur->ai_protocol);
+    if (socket == -1) {
+      continue;
+    }
+    ret = ::bind(socket, cur->ai_addr, cur->ai_addrlen);
+#ifdef _WIN32
+    closesocket(socket);
+#else
+    close(socket);
+#endif
+    if (ret == -1) {
+      continue;
+    }
+    break;
+  }
+  freeaddrinfo(result);
+  if (cur != nullptr) {
+    return CreateDeviceForHostname(hostname.data());
+  }
+  return CreateDeviceForHostname("127.0.0.1");
+}
+
+std::shared_ptr<gloo::transport::Device> CreateGlooDevice() {
+  char* ifname = std::getenv("GLOO_SOCKET_IFNAME");
+  if (ifname && std::strlen(ifname) > 1) {
+    return CreateDeviceForInterface(std::string(ifname));
+  } else {
+    return CreateDefaultDevice();
+  }
+}
+
+void send_recv(SendRecvOptions* opts) {
+  const auto& context = opts->context;
+  gloo::transport::UnboundBuffer* in = opts->in.get();
+  gloo::transport::UnboundBuffer* out = opts->out.get();
+  const auto slot = gloo::Slot::build(kSendRecvSlotPrefix, opts->tag);
+
+  if (context->rank == opts->src) {
+    in->send(opts->dst, slot);
+    in->waitSend(opts->timeout);
+  } else if (context->rank == opts->dst) {
+    out->recv(opts->src, slot);
+    out->waitRecv(opts->timeout);
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/gloo_utils.h

@@ -0,0 +1,219 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <gloo/allreduce.h>
+#include <gloo/math.h>
+#include <gloo/transport/tcp/device.h>
+#include <gloo/types.h>
+
+#include <climits>
+#include <memory>
+#include <string>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/reduce_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+namespace distributed {
+
+// data preparation
+#ifdef _WIN32
+#define GENERATE_FUNC(type, func, ...)       \
+  switch (type) {                            \
+    case DataType::FLOAT32:                  \
+      func<float>(__VA_ARGS__);              \
+      break;                                 \
+    case DataType::FLOAT64:                  \
+      func<double>(__VA_ARGS__);             \
+      break;                                 \
+    case DataType::FLOAT16:                  \
+      func<gloo::float16>(__VA_ARGS__);      \
+      break;                                 \
+    case DataType::INT32:                    \
+      func<int32_t>(__VA_ARGS__);            \
+      break;                                 \
+    case DataType::INT64:                    \
+      func<int64_t>(__VA_ARGS__);            \
+      break;                                 \
+    default:                                 \
+      VLOG(0) << "Error: Unknown DataType."; \
+      exit(-1);                              \
+  }
+#define HOST_NAME_MAX 256
+#else
+#define GENERATE_FUNC(type, func, args...)   \
+  switch (type) {                            \
+    case DataType::FLOAT32:                  \
+      func<float>(args);                     \
+      break;                                 \
+    case DataType::FLOAT64:                  \
+      func<double>(args);                    \
+      break;                                 \
+    case DataType::FLOAT16:                  \
+      func<gloo::float16>(args);             \
+      break;                                 \
+    case DataType::INT32:                    \
+      func<int32_t>(args);                   \
+      break;                                 \
+    case DataType::INT64:                    \
+      func<int64_t>(args);                   \
+      break;                                 \
+    case DataType::INT8:                     \
+      func<int8_t>(args);                    \
+      break;                                 \
+    case DataType::UINT8:                    \
+      func<uint8_t>(args);                   \
+      break;                                 \
+    case DataType::BOOL:                     \
+      func<bool>(args);                      \
+      break;                                 \
+    case DataType::BFLOAT16:                 \
+      func<phi::dtype::bfloat16>(args);      \
+      break;                                 \
+    default:                                 \
+      VLOG(0) << "Error: Unknown DataType."; \
+      exit(-1);                              \
+  }
+#endif
+
+template <typename T, typename P>
+void SetOutput(P* opts, DenseTensor* tensor) {
+  opts->setOutput(reinterpret_cast<T*>(tensor->data()), tensor->numel());
+}
+
+template <typename T, typename P>
+void SetInput(P* opts, const DenseTensor& tensor) {
+  // gloo only support mutable data input
+  opts->setInput(reinterpret_cast<T*>(const_cast<void*>(tensor.data())),
+                 tensor.numel());
+}
+
+template <typename T, typename P>
+void SetInputForScatter(P* opts, const DenseTensor& tensor, int nranks) {
+  std::vector<T*> ret;
+  ret.reserve(nranks);
+  T* raw_pointer = reinterpret_cast<T*>(const_cast<void*>(tensor.data()));
+  size_t offset = 0;
+  for (int i = 0; i < nranks; i++) {
+    ret.push_back(raw_pointer + offset);
+    offset += tensor.numel() / nranks;
+  }
+  opts->setInputs(ret, tensor.numel() / nranks);
+}
+
+template <typename T, typename P>
+void SetReduceFunc(P* opts, int reduce_type) {
+  // gloo only support mutable data input
+  ReduceType reduce_type_enum = static_cast<ReduceType>(reduce_type);
+  switch (reduce_type_enum) {
+    case ReduceType::kRedSum:
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::sum<T>));
+      break;
+    case ReduceType::kRedMax:
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::max<T>));
+      break;
+    case ReduceType::kRedMin:
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::min<T>));
+      break;
+    case ReduceType::kRedProd:
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::product<T>));
+      break;
+    case ReduceType::kRedAll:
+      // NOTE(zhonghui): There is no reduce_all math function for gloo, just use
+      // min to replace
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::min<T>));
+      break;
+    case ReduceType::kRedAny:
+      // NOTE(ooooo): There is no reduce_any math function for gloo, just use
+      // max to replace
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::max<T>));
+      break;
+    default:
+      PADDLE_THROW(
+          errors::InvalidArgument("Unsupported reduce type: %d.", reduce_type));
+  }
+}
+
+// env preparation
+std::shared_ptr<gloo::transport::Device> CreateGlooDevice();
+
+constexpr uint8_t kSendRecvSlotPrefix = 0x08;
+
+class SendRecvOptions {
+ public:
+  explicit SendRecvOptions(const std::shared_ptr<gloo::Context>& context)
+      : context(context), timeout(context->getTimeout()) {}
+
+  template <typename T>
+  void setInput(T* ptr, size_t elements) {
+    this->in = context->createUnboundBuffer(ptr, elements * sizeof(T));
+  }
+
+  template <typename T>
+  void setOutput(T* ptr, size_t elements) {
+    this->out = context->createUnboundBuffer(ptr, elements * sizeof(T));
+  }
+
+  void setSrc(int src) { this->src = src; }
+
+  void setDst(int dst) { this->dst = dst; }
+
+  void setTag(uint32_t tag) { this->tag = tag; }
+
+  void setTimeout(std::chrono::milliseconds timeout) {
+    this->timeout = timeout;
+  }
+
+ protected:
+  std::shared_ptr<gloo::Context> context;
+  std::unique_ptr<gloo::transport::UnboundBuffer> in;
+  std::unique_ptr<gloo::transport::UnboundBuffer> out;
+
+  // Rank of process to send_recv from.
+  int src = -1;
+
+  // Rank of process to send_recv to.
+  int dst = -1;
+
+  // Tag for this operation.
+  // Must be unique across operations executing in parallel.
+  uint32_t tag = 0;
+
+  // End-to-end timeout for this operation.
+  std::chrono::milliseconds timeout;
+
+  friend void send_recv(SendRecvOptions*);
+};
+
+void send_recv(SendRecvOptions* opts);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/nccl_comm_context.cc

@@ -0,0 +1,308 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/nccl_comm_context.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/distributed/check/nccl_dynamic_check.h"
+#include "paddle/phi/core/distributed/check/static_check.h"
+#include "paddle/phi/core/distributed/nccl_tools.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi::distributed {
+
+// set this flag to `true` and recompile to enable dynamic checks
+constexpr bool FLAGS_enable_nccl_dynamic_check = false;
+
+NCCLCommContext::NCCLCommContext(
+    int rank,
+    int size,
+    ncclUniqueId nccl_id,
+    int nccl_comm_init_option,
+    std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr)
+    : CommContext(rank, size),
+      nccl_version_(0),
+      nccl_comm_(nullptr),
+      nranks(size_),
+      myrank(rank_),
+      param(nccl_comm_init_option) {
+  this->CreateNCCLComm(nccl_id, nccl_config_ptr);
+  NCCL_CHECK(phi::dynload::ncclGetVersion(&nccl_version_));
+}
+
+void NCCLCommContext::CreateNCCLComm(
+    ncclUniqueId nccl_id,
+    std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr) {
+  if (param > 0 && phi::dynload::ncclCommInitRank2.IsValid()) {
+    LOG(WARNING) << "Creating modified qp with ncclCommInitRank2.";
+    NCCL_CHECK(phi::dynload::ncclCommInitRank2(
+        &nccl_comm_, nranks, nccl_id, myrank, param));
+  } else {
+    if (param > 0) {
+      LOG(WARNING) << "ncclCommInitRank2 is not supported.";
+    }
+    if (nccl_config_ptr != nullptr &&
+        phi::dynload::ncclCommInitRankConfigMemOpt.IsValid()) {
+      NCCL_CHECK(phi::dynload::ncclCommInitRankConfigMemOpt(
+          &nccl_comm_,
+          nranks,
+          nccl_id,
+          myrank,
+          nccl_config_ptr->GetOrigin(),
+          nccl_config_ptr->GetMemOpt()));
+    } else {
+      if (nccl_config_ptr != nullptr) {
+        LOG(WARNING) << "ncclCommInitRankConfigMemOpt is not supported.";
+      }
+      NCCL_CHECK(
+          phi::dynload::ncclCommInitRank(&nccl_comm_, nranks, nccl_id, myrank));
+    }
+  }
+}
+
+void NCCLCommContext::DestroyNCCLComm() {
+  if (nccl_comm_ != nullptr) {
+    NCCL_CHECK(phi::dynload::ncclCommDestroy(nccl_comm_));
+    nccl_comm_ = nullptr;
+  }
+}
+
+int NCCLCommContext::GetNcclVersion() { return nccl_version_; }
+
+ncclComm_t NCCLCommContext::GetNcclComm() { return nccl_comm_; }
+
+gpuStream_t NCCLCommContext::GetStream() { return dev_ctx_->stream(); }
+
+phi::GPUContext* NCCLCommContext::GetDevContext() { return dev_ctx_.get(); }
+
+void NCCLCommContext::SetDevContext(
+    std::unique_ptr<phi::GPUContext>&& dev_ctx) {
+  dev_ctx_ = std::move(dev_ctx);
+}
+
+gpuEvent_t NCCLCommContext::GetComputeEvent() { return compute_event_.get(); }
+
+void NCCLCommContext::SetComputeEvent(
+    std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type>&&
+        compute_event) {
+  compute_event_ = std::move(compute_event);
+}
+
+gpuEvent_t NCCLCommContext::GetCommEvent() { return comm_event_.get(); }
+
+void NCCLCommContext::SetCommEvent(
+    std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type>&& comm_event) {
+  comm_event_ = std::move(comm_event);
+}
+
+void NCCLCommContext::Broadcast(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                int root,
+                                gpuStream_t stream) {
+  CommStaticCheck::SameShape(*out_tensor,
+                             in_tensor,
+                             /*dst_rank*/ rank_,
+                             /*cur_rank*/ rank_,
+                             size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    NCCLDynamicCheck::CheckShape(*out_tensor, root, rank_, nccl_comm_);
+  }
+  NCCL_CHECK(phi::dynload::ncclBroadcast(in_tensor.data(),
+                                         out_tensor->data(),
+                                         in_tensor.numel(),
+                                         ToNCCLDataType(in_tensor.type()),
+                                         root,
+                                         nccl_comm_,
+                                         stream));
+}
+
+void NCCLCommContext::AllGather(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                gpuStream_t stream) {
+  phi::distributed::CommStaticCheck::GatherLikeShape(*out_tensor,
+                                                     in_tensor,
+                                                     /*dst_rank*/ rank_,
+                                                     /*cur_rank*/ rank_,
+                                                     size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    phi::distributed::NCCLDynamicCheck::CheckShape(*out_tensor,
+                                                   /*root_rank*/ 0,
+                                                   rank_,
+                                                   nccl_comm_);
+  }
+  NCCL_CHECK(phi::dynload::ncclAllGather(in_tensor.data(),
+                                         out_tensor->data(),
+                                         in_tensor.numel(),
+                                         ToNCCLDataType(in_tensor.type()),
+                                         nccl_comm_,
+                                         stream));
+}
+void NCCLCommContext::ReduceScatter(DenseTensor* out_tensor,
+                                    const DenseTensor& in_tensor,
+                                    ncclRedOp_t reduce_type,
+                                    gpuStream_t stream) {
+  PADDLE_ENFORCE_EQ(
+      in_tensor.dtype() != DataType::FLOAT8_E4M3FN &&
+          in_tensor.dtype() != DataType::FLOAT8_E5M2,
+      true,
+      common::errors::InvalidArgument(
+          "float8 dtypes are not currently supported for NCCL reductions"));
+
+  phi::distributed::CommStaticCheck::ScatterLikeShape(*out_tensor,
+                                                      in_tensor,
+                                                      /*dst_rank*/ rank_,
+                                                      /*cur_rank*/ rank_,
+                                                      size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    phi::distributed::NCCLDynamicCheck::CheckShape(*out_tensor,
+                                                   /*root_rank*/ 0,
+                                                   rank_,
+                                                   nccl_comm_);
+  }
+  NCCL_CHECK(phi::dynload::ncclReduceScatter(in_tensor.data(),
+                                             out_tensor->data(),
+                                             out_tensor->numel(),
+                                             ToNCCLDataType(in_tensor.type()),
+                                             reduce_type,
+                                             nccl_comm_,
+                                             stream));
+}
+
+void NCCLCommContext::Send(const DenseTensor& in_tensor,
+                           const int64_t& count,
+                           const int& peer,
+                           gpuStream_t stream) {
+  phi::distributed::CommStaticCheck::CheckShape(in_tensor, rank_, size_);
+
+  if (FLAGS_enable_nccl_dynamic_check) {
+    NCCLDynamicCheck::CheckShape(in_tensor, rank_, rank_, nccl_comm_);
+  }
+
+  NCCL_CHECK(phi::dynload::ncclSend(in_tensor.data(),
+                                    count,
+                                    ToNCCLDataType(in_tensor.dtype()),
+                                    peer,
+                                    nccl_comm_,
+                                    stream));
+  VLOG(3) << "rank " << GetRank() << " send "
+          << common::product(in_tensor.dims()) << " to " << peer;
+}
+
+void NCCLCommContext::Recv(DenseTensor* out_tensor,
+                           const int64_t& count,
+                           const int& peer,
+                           gpuStream_t stream) {
+  phi::distributed::CommStaticCheck::CheckShape(*out_tensor, rank_, size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    NCCLDynamicCheck::CheckShape(*out_tensor, peer, rank_, nccl_comm_);
+  }
+
+  NCCL_CHECK(phi::dynload::ncclRecv(out_tensor->data(),
+                                    count,
+                                    ToNCCLDataType(out_tensor->dtype()),
+                                    peer,
+                                    nccl_comm_,
+                                    stream));
+  VLOG(3) << "rank " << GetRank() << " recv "
+          << common::product(out_tensor->dims()) << " from " << peer;
+}
+
+void NCCLCommContext::AllReduce(DenseTensor* out_tensor,
+                                const DenseTensor& in_tensor,
+                                ncclRedOp_t reduce_type,
+                                gpuStream_t stream) {
+  PADDLE_ENFORCE_EQ(
+      in_tensor.dtype() != DataType::FLOAT8_E4M3FN &&
+          in_tensor.dtype() != DataType::FLOAT8_E5M2,
+      true,
+      common::errors::InvalidArgument(
+          "float8 dtypes are not currently supported for NCCL reductions"));
+
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ rank_,
+                                               /*cur_rank*/ rank_,
+                                               size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    phi::distributed::NCCLDynamicCheck::CheckShape(*out_tensor,
+                                                   /*root_rank*/ 0,
+                                                   rank_,
+                                                   nccl_comm_);
+  }
+  NCCL_CHECK(phi::dynload::ncclAllReduce(in_tensor.data(),
+                                         out_tensor->data(),
+                                         in_tensor.numel(),
+                                         ToNCCLDataType(in_tensor.type()),
+                                         reduce_type,
+                                         nccl_comm_,
+                                         stream));
+}
+
+void NCCLCommContext::Reduce(DenseTensor* out_tensor,
+                             const DenseTensor& in_tensor,
+                             ncclRedOp_t reduce_type,
+                             int root,
+                             gpuStream_t stream) {
+  PADDLE_ENFORCE_EQ(
+      in_tensor.dtype() != DataType::FLOAT8_E4M3FN &&
+          in_tensor.dtype() != DataType::FLOAT8_E5M2,
+      true,
+      common::errors::InvalidArgument(
+          "float8 dtypes are not currently supported for NCCL reductions"));
+
+  phi::distributed::CommStaticCheck::SameShape(*out_tensor,
+                                               in_tensor,
+                                               /*dst_rank*/ root,
+                                               /*cur_rank*/ rank_,
+                                               size_);
+  if (FLAGS_enable_nccl_dynamic_check) {
+    phi::distributed::NCCLDynamicCheck::CheckShape(*out_tensor,
+                                                   /*root_rank*/ root,
+                                                   rank_,
+                                                   nccl_comm_);
+  }
+  NCCL_CHECK(phi::dynload::ncclReduce(in_tensor.data(),
+                                      out_tensor->data(),
+                                      in_tensor.numel(),
+                                      ToNCCLDataType(in_tensor.type()),
+                                      reduce_type,
+                                      root,
+                                      nccl_comm_,
+                                      stream));
+}
+
+void NCCLCommContext::GroupStart() {
+  NCCL_CHECK(phi::dynload::ncclGroupStart());
+}
+void NCCLCommContext::GroupEnd() { NCCL_CHECK(phi::dynload::ncclGroupEnd()); }
+
+#if NCCL_VERSION_CODE >= 21100
+void NCCLCommContext::RedOpCreatePreMulSum(ncclRedOp_t* op,
+                                           void* scalar,
+                                           ncclDataType_t dtype,
+                                           ncclScalarResidence_t residence) {
+  PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclRedOpCreatePreMulSum(
+      op, scalar, dtype, residence, nccl_comm_));
+}
+
+void NCCLCommContext::RedOpDestroy(ncclRedOp_t op) {
+  PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclRedOpDestroy(op, nccl_comm_));
+}
+#endif
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/nccl_comm_context.h

@@ -0,0 +1,151 @@
+// Copyright (c) 2022 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.
+#pragma once
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+#include "paddle/phi/core/distributed/nccl_config.h"
+
+#if defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/dynload/rccl.h"
+#else
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+namespace phi {
+class DenseTensor;
+namespace distributed {
+
+class NCCLCommContext final : public CommContext {
+ public:
+  NCCLCommContext(
+      int rank,
+      int size,
+      ncclUniqueId nccl_id,
+      int nccl_comm_init_option = 0,
+      std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr = nullptr);
+  ~NCCLCommContext() override = default;
+
+  int GetNcclVersion();
+
+  ncclComm_t GetNcclComm();
+
+  void CreateNCCLComm(
+      ncclUniqueId nccl_id,
+      std::shared_ptr<phi::distributed::NCCLConfig> nccl_config_ptr = nullptr);
+
+  void DestroyNCCLComm();
+
+  gpuStream_t GetStream();
+
+  gpuEvent_t GetComputeEvent();
+
+  void SetComputeEvent(
+      std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type>&&
+          compute_event);
+
+  gpuEvent_t GetCommEvent();
+
+  void SetCommEvent(
+      std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type>&& comm_event);
+
+  phi::GPUContext* GetDevContext();
+
+  void SetDevContext(std::unique_ptr<phi::GPUContext>&& dev_ctx);
+
+  void Broadcast(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 int root,
+                 gpuStream_t stream);
+
+  void Send(const DenseTensor& in_tensor,
+            const int64_t& count,
+            const int& peer,
+            gpuStream_t stream);
+
+  void Recv(DenseTensor* out_tensor,
+            const int64_t& count,
+            const int& peer,
+            gpuStream_t stream);
+
+  void ReduceScatter(DenseTensor* out_tensor,
+                     const DenseTensor& in_tensor,
+                     ncclRedOp_t reduce_type,
+                     gpuStream_t stream);
+
+  void AllGather(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 gpuStream_t stream);
+
+  void AllReduce(DenseTensor* out_tensor,
+                 const DenseTensor& in_tensor,
+                 ncclRedOp_t reduce_type,
+                 gpuStream_t stream);
+
+  void Reduce(DenseTensor* out_tensor,
+              const DenseTensor& in_tensor,
+              ncclRedOp_t reduce_type,
+              int root,
+              gpuStream_t stream);
+
+  void GroupStart();
+
+  void GroupEnd();
+
+#if NCCL_VERSION_CODE >= 21100
+  // Creates a new reduction operator which pre-multiplies input values by a
+  // given scalar locally before reducing them with peer values via summation.
+  void RedOpCreatePreMulSum(ncclRedOp_t* op,
+                            void* scalar,
+                            ncclDataType_t dtype,
+                            ncclScalarResidence_t residence);
+
+  // Destroys the reduction operator op. The operator must have been created by
+  // ncclRedOpCreatePreMul with the matching communicator comm.
+  void RedOpDestroy(ncclRedOp_t op);
+#endif
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(NCCLCommContext);
+
+  int nccl_version_;
+
+  ncclComm_t nccl_comm_;
+
+  std::unique_ptr<phi::GPUContext> dev_ctx_;
+
+  // used for comm wait compute, compute_stream-->event-->comm_stream
+  std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type> compute_event_;
+
+  // used for compute wait comm, comm_stream-->event-->compute_stream
+  std::shared_ptr<std::remove_pointer<phi::gpuEvent_t>::type> comm_event_;
+
+  int nranks;
+  int myrank;
+  int param;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/nccl_comm_task.cc

@@ -0,0 +1,268 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/nccl_comm_task.h"
+
+#include "gflags/gflags.h"
+
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+#include "paddle/phi/core/distributed/nccl_tools.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi::distributed {
+
+NCCLCommTask::NCCLCommTask(const Place& place,
+                           const std::string& group_key,
+                           int rank,
+                           int size,
+                           int gid,
+                           uint64_t seq,
+                           int64_t numel,
+                           bool sync_op,
+                           bool use_calc_stream,
+                           ncclComm_t nccl_comm,
+                           gpuStream_t stream,
+                           CommType comm_type,
+                           int64_t timeout)
+    : CommTask("NCCL",
+               place,
+               group_key,
+               rank,
+               size,
+               gid,
+               seq,
+               numel,
+               nccl_comm,
+               stream,
+               comm_type),
+      timeout_(std::chrono::milliseconds(timeout)),
+      sync_op_(sync_op),
+      use_calc_stream_(use_calc_stream),
+      nccl_start_event_(nullptr),
+      nccl_end_event_(nullptr) {
+  start_trace_updated_ = false;
+  start_event_created_ = false;
+  end_event_created_ = false;
+  start_time_ = std::chrono::steady_clock::now();
+}
+
+void NCCLCommTask::StartRecord() {
+  backends::gpu::GPUDeviceGuard guard(place_.device);
+  if (!start_event_created_) {
+#ifdef PADDLE_WITH_CUDA
+    CUDA_CHECK(cudaEventCreateWithFlags(&nccl_start_event_, cuda_event_flags_));
+#else  // PADDLE_WITH_HIP
+    HIP_CHECK(hipEventCreateWithFlags(&nccl_start_event_, hip_event_flags_));
+#endif
+    start_event_created_ = true;
+  }
+#ifdef PADDLE_WITH_CUDA
+  CUDA_CHECK(cudaEventRecord(nccl_start_event_, nccl_stream_));
+#else  // PADDLE_WITH_HIP
+  HIP_CHECK(hipEventRecord(nccl_start_event_, nccl_stream_));
+#endif
+}
+void NCCLCommTask::EndRecord() {
+  backends::gpu::GPUDeviceGuard guard(place_.device);
+  if (!end_event_created_) {
+#ifdef PADDLE_WITH_CUDA
+    CUDA_CHECK(cudaEventCreateWithFlags(&nccl_end_event_, cuda_event_flags_));
+#else  // PADDLE_WITH_HIP
+    HIP_CHECK(hipEventCreateWithFlags(&nccl_end_event_, hip_event_flags_));
+#endif
+    end_event_created_ = true;
+  }
+#ifdef PADDLE_WITH_CUDA
+  CUDA_CHECK(cudaEventRecord(nccl_end_event_, nccl_stream_));
+#else  // PADDLE_WITH_HIP
+  HIP_CHECK(hipEventRecord(nccl_end_event_, nccl_stream_));
+#endif
+}
+
+#ifdef PADDLE_WITH_CUDA
+void NCCLCommTask::ClearRecord() {
+  if (start_event_created_) {
+    backends::gpu::GPUDeviceGuard guard(place_.device);
+    CUDA_CHECK(cudaEventDestroy(nccl_start_event_));
+    start_event_created_ = false;
+  }
+  if (end_event_created_) {
+    backends::gpu::GPUDeviceGuard guard(place_.device);
+    CUDA_CHECK(cudaEventDestroy(nccl_end_event_));
+    end_event_created_ = false;
+  }
+}
+#else  // PADDLE_WITH_HIP
+void NCCLCommTask::ClearRecord() {
+  if (start_event_created_) {
+    backends::gpu::GPUDeviceGuard guard(place_.device);
+    HIP_CHECK(hipEventDestroy(nccl_start_event_));
+    start_event_created_ = false;
+  }
+  if (end_event_created_) {
+    backends::gpu::GPUDeviceGuard guard(place_.device);
+    HIP_CHECK(hipEventDestroy(nccl_end_event_));
+    end_event_created_ = false;
+  }
+}
+#endif
+
+bool NCCLCommTask::CudaEventQuery(gpuEvent_t event) {
+#ifdef PADDLE_WITH_CUDA
+  cudaError_t ret = cudaEventQuery(event);
+  if (ret == cudaSuccess) {
+    return true;
+  } else if (ret != cudaErrorNotReady) {
+    CUDA_CHECK(ret);
+  } else {
+    // ignore and clear the error if not ready
+    CUDA_CHECK(cudaGetLastError());
+  }
+#else  // PADDLE_WITH_HIP
+  hipError_t ret = hipEventQuery(event);
+  if (ret == hipSuccess) {
+    return true;
+  } else if (ret != hipErrorNotReady) {
+    HIP_CHECK(ret);
+  } else {
+    // ignore and clear the error if not ready
+    HIP_CHECK(hipGetLastError());
+  }
+#endif
+  return false;
+}
+
+std::string GetNCCLErrorDetail(ncclResult_t result) {
+  std::string detail;
+  std::string last_error;
+#ifdef ENABLE_NCCL_GET_LAST_ERROR
+  last_error =
+      ", Last error: " + std::string(phi::dynload::ncclGetLastError(NULL));
+#endif
+  switch (result) {
+    case ncclUnhandledCudaError:
+      detail = "ncclUnhandledCudaError: Call to CUDA function failed.";
+      break;
+    case ncclSystemError:
+      detail =
+          "ncclSystemError: System call (e.g. socket, malloc) or external "
+          "library call failed or device error. ";
+#ifndef NCCL_REMOTE_ERROR
+      // Before ncclRemoteError was created, unexpected remote disconnect was
+      // categorized as ncclSystemError
+      detail += "It can be also caused by unexpected exit of a remote peer.";
+#endif
+      break;
+    case ncclInternalError:
+      detail = "ncclInternalError: Internal check failed.";
+      break;
+    case ncclInvalidArgument:
+      detail = "ncclInvalidArgument: Invalid value for an argument.";
+      break;
+    case ncclInvalidUsage:
+      detail =
+          "ncclInvalidUsage: This usually reflects invalid usage of NCCL "
+          "library.";
+      break;
+#ifdef NCCL_REMOTE_ERROR
+    case ncclRemoteError:
+      detail =
+          "ncclRemoteError: A call failed possibly due to a network error or a "
+          "remote process exiting prematurely.";
+      break;
+#endif
+    default:
+      detail = "Unknown NCCL error!";
+  }
+  return detail + last_error;
+}
+
+std::string NCCLCommTask::GetCommErrors() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (!comm_error_.empty()) {
+    return comm_error_;
+  }
+
+  ncclResult_t nccl_async_error;
+  NCCL_CHECK(
+      phi::dynload::ncclCommGetAsyncError(nccl_comm_, &nccl_async_error));
+  if (nccl_async_error != ncclSuccess) {
+    comm_error_ =
+        "\n\t Find nccl comm error: " + GetNCCLErrorDetail(nccl_async_error);
+  }
+  return comm_error_;
+}
+
+bool NCCLCommTask::IsStarted() {
+  if (started_) {
+    return true;
+  }
+  if (start_event_created_ && CudaEventQuery(nccl_start_event_)) {
+    started_ = true;
+    updated_ = true;
+  }
+  return started_;
+}
+
+bool NCCLCommTask::IsCompleted() {
+  if (completed_) {
+    return true;
+  }
+  if (end_event_created_ && CudaEventQuery(nccl_end_event_)) {
+    completed_ = true;
+    updated_ = true;
+  }
+  return completed_;
+}
+
+void NCCLCommTask::SetUpdated(bool updated) { updated_ = updated; }
+
+bool NCCLCommTask::IsUpdated() { return updated_; }
+
+bool NCCLCommTask::IsTimeout() {
+  auto current_timepoint = std::chrono::steady_clock::now();
+  return std::chrono::duration_cast<std::chrono::milliseconds>(
+             current_timepoint - start_time_) >= timeout_;
+}
+
+void NCCLCommTask::AbortComm() {
+  std::unique_lock<std::mutex> lock(mutex_);
+  if (aborted_) {
+    return;
+  }
+  NCCL_CHECK(phi::dynload::ncclCommAbort(nccl_comm_));
+
+  aborted_ = true;
+  nccl_comm_ = nullptr;
+  return;
+}
+
+std::string NCCLCommTask::GetTraceMsg() {
+  auto global_ranks =
+      phi::distributed::CommContextManager::GetInstance().GetGroupRanks(
+          group_key_);
+  return "group_key:" + group_key_ +
+         ",group_ranks:" + VectorToString(global_ranks) +
+         ",global_rank:" + std::to_string(global_rank_) +
+         ",local_rank:" + std::to_string(rank_) +
+         ",comm_count:" + std::to_string(seq_) +
+         ",op:" + CommTypeToString(comm_type_) +
+         ",started:" + std::to_string(started_) +
+         ",completed:" + std::to_string(completed_) +
+         ",numel:" + std::to_string(numel_) +
+         ",nranks:" + std::to_string(size_);
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/nccl_comm_task.h

@@ -0,0 +1,93 @@
+// Copyright (c) 2023 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.
+#pragma once
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/backends/gpu/gpu_decls.h"
+#include "paddle/phi/core/distributed/comm_context.h"
+#include "paddle/phi/core/distributed/comm_task.h"
+#include "paddle/phi/core/distributed/utils.h"
+
+#if defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/dynload/rccl.h"
+#else
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+namespace phi {
+class DenseTensor;
+namespace distributed {
+
+static int64_t DefaultTimeout = 30 * 60 * 1000;
+
+class NCCLCommTask : public CommTask {
+ public:
+  NCCLCommTask(const Place& place = Place(),
+               const std::string& group_key = "",
+               int rank = -1,
+               int size = 0,
+               int gid = 0,
+               uint64_t seq = 0,
+               int64_t numel = 0,
+               bool sync_op = true,
+               bool use_calc_stream = false,
+               ncclComm_t = nullptr,
+               gpuStream_t = nullptr,
+               CommType comm_type = CommType::UNKNOWN,
+               int64_t timeout = DefaultTimeout);
+  ~NCCLCommTask() override = default;
+
+  // check whether the nccl kernel started
+  bool IsStarted() override;
+  bool IsTimeout() override;
+  bool IsCompleted() override;
+  void SetUpdated(bool updated) override;
+  bool IsUpdated() override;
+
+  std::string GetTraceMsg() override;
+  std::string GetCommErrors() override;
+  void AbortComm() override;
+
+  void StartRecord() override;
+  void EndRecord() override;
+  void ClearRecord() override;
+
+  bool CudaEventQuery(gpuEvent_t event);
+
+ protected:
+  std::mutex mutex_;
+  std::chrono::milliseconds timeout_;
+
+#ifdef PADDLE_WITH_CUDA
+  unsigned int cuda_event_flags_ = cudaEventDisableTiming;
+#else  // PADDLE_WITH_HIP
+  unsigned int hip_event_flags_ = hipEventDisableTiming;
+#endif
+
+  bool sync_op_;
+  bool use_calc_stream_;
+
+  bool start_event_created_;
+  bool end_event_created_;
+  gpuEvent_t nccl_start_event_;
+  gpuEvent_t nccl_end_event_;
+
+  std::string comm_error_;
+
+ private:
+  DISABLE_COPY_AND_ASSIGN(NCCLCommTask);
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/nccl_config.cc

@@ -0,0 +1,81 @@
+// Copyright (c) 2025 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.
+
+#include "paddle/phi/core/distributed/nccl_config.h"
+
+#include "glog/logging.h"
+
+namespace phi::distributed {
+
+std::shared_ptr<NCCLConfig> NCCLConfig::CreateNCCLConfig(
+    const std::string& commName,
+    const int ll_buffsize,
+    const int ll128_buffsize,
+    const int simple_buffsize,
+    const int buffsize_align,
+    const int nchannels,
+    const std::string& algoStr,
+    const std::string& protoStr) {
+  return std::make_shared<NCCLConfig>(commName,
+                                      ll_buffsize,
+                                      ll128_buffsize,
+                                      simple_buffsize,
+                                      buffsize_align,
+                                      nchannels,
+                                      algoStr,
+                                      protoStr);
+}
+
+NCCLConfig::NCCLConfig(const std::string& commName,
+                       const int ll_buffsize,
+                       const int ll128_buffsize,
+                       const int simple_buffsize,
+                       const int buffsize_align,
+                       const int nchannels,
+                       const std::string& algoStr,
+                       const std::string& protoStr)
+    : commName_(commName),
+      ll_buffsize_(ll_buffsize),
+      ll128_buffsize_(ll128_buffsize),
+      simple_buffsize_(simple_buffsize),
+      buffsize_align_(buffsize_align),
+      nchannels_(nchannels),
+      algoStr_(algoStr),
+      protoStr_(protoStr),
+      nccl_memopt_config_ptr(nullptr) {
+  if (phi::dynload::ncclCommGenMemOptConfig.IsValid()) {
+    nccl_memopt_config_ptr = phi::dynload::ncclCommGenMemOptConfig(
+        commName_.empty() ? nullptr : commName_.c_str(),
+        ll_buffsize_ >= 0 ? ll_buffsize_ : -1,
+        ll128_buffsize_ >= 0 ? ll128_buffsize_ : -1,
+        simple_buffsize_ >= 0 ? simple_buffsize_ : -1,
+        buffsize_align_ >= 0 ? buffsize_align_ : -1,
+        nchannels_ >= 0 ? nchannels_ : -1,
+        algoStr_.empty() ? nullptr : algoStr_.c_str(),
+        protoStr_.empty() ? nullptr : protoStr_.c_str());
+  }
+}
+
+ncclConfig_t* NCCLConfig::GetOrigin() { return nullptr; }
+
+ncclMemOptConfig_t* NCCLConfig::GetMemOpt() { return nccl_memopt_config_ptr; }
+
+NCCLConfig::~NCCLConfig() {
+  if (phi::dynload::ncclCommFreeMemOptConfig.IsValid() &&
+      nccl_memopt_config_ptr != nullptr) {
+    phi::dynload::ncclCommFreeMemOptConfig(nccl_memopt_config_ptr);
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/nccl_config.h

@@ -0,0 +1,64 @@
+// Copyright (c) 2025 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.
+#pragma once
+
+#include <memory>
+#if defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/backends/dynload/rccl.h"
+#else
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+namespace phi {
+namespace distributed {
+
+class NCCLConfig {
+ public:
+  static std::shared_ptr<NCCLConfig> CreateNCCLConfig(
+      const std::string& commName = "",
+      const int ll_buffsize = -1,
+      const int ll128_buffsize = -1,
+      const int simple_buffsize = -1,
+      const int buffsize_align = -1,
+      const int nchannels = -1,
+      const std::string& algoStr = "",
+      const std::string& protoStr = "");
+  ncclConfig_t* GetOrigin();
+  ncclMemOptConfig_t* GetMemOpt();
+
+  NCCLConfig(const std::string& commName,
+             const int ll_buffsize,
+             const int ll128_buffsize,
+             const int simple_buffsize,
+             const int buffsize_align,
+             const int nchannels,
+             const std::string& algoStr,
+             const std::string& protoStr);
+  ~NCCLConfig();
+
+ private:
+  const std::string commName_;
+  const int ll_buffsize_;
+  const int ll128_buffsize_;
+  const int simple_buffsize_;
+  const int buffsize_align_;
+  const int nchannels_;
+  const std::string algoStr_;
+  const std::string protoStr_;
+
+  ncclMemOptConfig_t* nccl_memopt_config_ptr{nullptr};
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/nccl_tools.cc

@@ -0,0 +1,98 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/nccl_tools.h"
+
+#include <unordered_map>
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/enforce.h"
+
+#if NCCL_VERSION_CODE >= 21300
+#define ENABLE_NCCL_GET_LAST_ERROR
+#define NCCL_REMOTE_ERROR
+#endif
+
+namespace phi {
+namespace distributed {
+
+ncclRedOp_t ToNCCLRedType(ReduceOp reduction) {
+  static const std::unordered_map<ReduceOp, ncclRedOp_t> red_type = {
+    {ReduceOp::MIN, ncclMin},
+    {ReduceOp::MAX, ncclMax},
+    {ReduceOp::SUM, ncclSum},
+    {ReduceOp::PRODUCT, ncclProd},
+#if NCCL_VERSION_CODE >= 21000
+    {ReduceOp::AVG, ncclAvg},
+#endif
+  };
+  auto it = red_type.find(reduction);
+  PADDLE_ENFORCE_EQ(it != red_type.end(),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Invalid nccl reduction. Must be ncclMin | ncclMax | "
+                        "ncclProd | ncclSum | ncclAvg."));
+  return it->second;
+}
+
+std::string SerializeNCCLUniqueId(const ncclUniqueId& ncclID) {
+  const uint8_t* bytes = reinterpret_cast<const uint8_t*>(&ncclID);
+  std::ostringstream oss;
+  for (auto i = 0; i < NCCL_UNIQUE_ID_BYTES; ++i) {
+    oss << std::hex << static_cast<int>(bytes[i]);
+  }
+  return oss.str();
+}
+
+std::string NCCLDTypeToString(ncclDataType_t dtype) {
+#define PD_NCCL_DTYPE_TO_STR(__nccl_dtype, __str_dtype) \
+  if (dtype == __nccl_dtype) return __str_dtype;
+  PD_NCCL_DTYPE_TO_STR(ncclFloat, "float32");
+  PD_NCCL_DTYPE_TO_STR(ncclFloat32, "float32");
+  PD_NCCL_DTYPE_TO_STR(ncclHalf, "float16");
+  PD_NCCL_DTYPE_TO_STR(ncclFloat16, "float16");
+#if (NCCL_VERSION_CODE >= 21000 && CUDA_VERSION >= 11000) || \
+    defined(PADDLE_WITH_HIP)
+  PD_NCCL_DTYPE_TO_STR(ncclBfloat16, "bfloat16");
+#endif
+  PD_NCCL_DTYPE_TO_STR(ncclDouble, "float64");
+  PD_NCCL_DTYPE_TO_STR(ncclFloat64, "float64");
+
+  PD_NCCL_DTYPE_TO_STR(ncclInt8, "int8");
+  PD_NCCL_DTYPE_TO_STR(ncclChar, "int8");
+  PD_NCCL_DTYPE_TO_STR(ncclUint8, "uint8");
+  PD_NCCL_DTYPE_TO_STR(ncclInt32, "int32");
+  PD_NCCL_DTYPE_TO_STR(ncclInt, "int32");
+  PD_NCCL_DTYPE_TO_STR(ncclUint32, "uint32");
+  PD_NCCL_DTYPE_TO_STR(ncclInt64, "int64");
+  PD_NCCL_DTYPE_TO_STR(ncclUint64, "uint64");
+
+#undef PD_NCCL_DTYPE_TO_STR
+  PADDLE_THROW(common::errors::InvalidArgument(
+      "This datatype %d in nccl is not supported.", static_cast<int>(dtype)));
+}
+
+std::string NCCLRedTypeToString(ncclRedOp_t op) {
+  if (op == ncclSum) return "SUM";
+  if (op == ncclProd) return "PROD";
+  if (op == ncclMin) return "MIN";
+  if (op == ncclMax) return "MAX";
+#if NCCL_VERSION_CODE >= 21000
+  if (op == ncclAvg) return "AVG";
+#endif
+  return "UDF_" + std::to_string(op);
+}
+
+}  //  namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/nccl_tools.h

@@ -0,0 +1,77 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <string>
+
+#include "paddle/phi/core/distributed/types.h"
+
+#ifdef PADDLE_WITH_RCCL
+#include <hip/hip_runtime.h>
+#include "paddle/phi/backends/dynload/rccl.h"
+#else
+#include <cuda_runtime.h>
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+namespace phi {
+namespace distributed {
+
+#define NCCL_CHECK(cmd)                                                   \
+  do {                                                                    \
+    ncclResult_t r = cmd;                                                 \
+    if (r != ncclSuccess) {                                               \
+      PADDLE_THROW(                                                       \
+          common::errors::External("Failed, NCCL error %s:%d '%s'\n",     \
+                                   __FILE__,                              \
+                                   __LINE__,                              \
+                                   phi::dynload::ncclGetErrorString(r))); \
+    }                                                                     \
+  } while (0)
+
+#ifdef PADDLE_WITH_NCCL
+#define CUDA_CHECK(expr)                                                       \
+  do {                                                                         \
+    cudaError_t r = expr;                                                      \
+    if (r != cudaSuccess) {                                                    \
+      PADDLE_THROW(common::errors::External("Failed, cuda error %s:%d '%s'\n", \
+                                            __FILE__,                          \
+                                            __LINE__,                          \
+                                            cudaGetErrorString(r)));           \
+    }                                                                          \
+  } while (0)
+#else  // PADDLE_WITH_RCCL
+#define HIP_CHECK(expr)                                                       \
+  do {                                                                        \
+    hipError_t r = expr;                                                      \
+    if (r != hipSuccess) {                                                    \
+      PADDLE_THROW(common::errors::External("Failed, hip error %s:%d '%s'\n", \
+                                            __FILE__,                         \
+                                            __LINE__,                         \
+                                            hipGetErrorString(r)));           \
+    }                                                                         \
+  } while (0)
+#endif
+
+ncclRedOp_t ToNCCLRedType(ReduceOp reduction);
+
+std::string SerializeNCCLUniqueId(const ncclUniqueId& ncclID);
+
+std::string NCCLDTypeToString(ncclDataType_t dtype);
+
+std::string NCCLRedTypeToString(ncclRedOp_t op);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/CMakeLists.txt

@@ -0,0 +1,8 @@
+set(STORE_COMMON_SRCS tcp_store.cc tcp_store_libuv.cc tcp_utils.cc socket.cpp
+                      store.cc store_utils.cc)
+
+if(WITH_GLOO)
+  list(APPEND STORE_COMMON_SRCS gloo_store.cc)
+endif()
+
+collect_srcs(core_srcs SRCS ${STORE_COMMON_SRCS})

paddle/phi/core/distributed/store/gloo_store.cc

@@ -0,0 +1,45 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/store/gloo_store.h"
+
+namespace phi::distributed {
+
+GlooStore::GlooStore(const std::shared_ptr<phi::distributed::Store>& store)
+    : store_(store) {}
+
+std::vector<char> GlooStore::get(const std::string& key) {
+  auto value = store_->get(key);
+  return std::vector<char>(value.begin(), value.end());
+}
+
+void GlooStore::wait(const std::vector<std::string>& keys) {
+  for (auto& key : keys) {
+    store_->wait(key);
+  }
+}
+
+void GlooStore::set(const std::string& key, const std::vector<char>& value) {
+  std::vector<uint8_t> tmp(value.begin(), value.end());
+  store_->set(key, tmp);
+}
+
+void GlooStore::wait(const std::vector<std::string>& keys,
+                     const std::chrono::milliseconds& timeout) {
+  for (auto& key : keys) {
+    store_->wait(key);
+  }
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/store/gloo_store.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <chrono>
+#include <memory>
+#include <vector>
+
+#include "gloo/rendezvous/store.h"
+
+#include "paddle/phi/core/distributed/store/store.h"
+
+namespace phi {
+namespace distributed {
+
+class GlooStore : public gloo::rendezvous::Store {
+ public:
+  explicit GlooStore(const std::shared_ptr<phi::distributed::Store>& store);
+
+  ~GlooStore() = default;
+
+  std::vector<char> get(const std::string& key) override;
+
+  void wait(const std::vector<std::string>& keys) override;
+
+  void set(const std::string& key, const std::vector<char>& value) override;
+
+  void wait(const std::vector<std::string>& keys,
+            const std::chrono::milliseconds& timeout) override;
+
+ protected:
+  std::shared_ptr<phi::distributed::Store> store_;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/socket.cpp

@@ -0,0 +1,78 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/store/socket.h"
+#include <array>
+
+#ifndef _WIN32
+#include <arpa/inet.h>
+#include <netinet/ip.h>
+#include <sys/socket.h>
+#include <unistd.h>
+#endif
+#include <cerrno>
+#include <cstdio>
+
+namespace phi::distributed {
+
+#ifdef _WIN32
+static int _get_sockname_of_win(int sock, char* out, int out_len) {
+  snprintf(out, out_len, "not support win now");
+  return 0;
+}
+#else
+static int _get_sockname(int sock, char *out, int out_len) {
+  struct sockaddr_in addr = {};
+  socklen_t s_len = sizeof(addr);
+
+  if (::getpeername(sock, reinterpret_cast<sockaddr *>(&addr), &s_len)) {
+    ::snprintf(
+        out, out_len, "can't getsocketname of %d, errno:%d", sock, errno);
+    return -1;
+  }
+
+  std::array<char, 128> ip = {};
+  int port = 0;
+
+  // deal with both IPv4 and IPv6:
+  if (addr.sin_family == AF_INET) {
+    struct sockaddr_in *s = (struct sockaddr_in *)&addr;
+    port = ntohs(s->sin_port);
+    ::inet_ntop(AF_INET, &s->sin_addr, ip.data(), sizeof(ip));
+  } else {  // AF_INET6
+    struct sockaddr_in6 *s = (struct sockaddr_in6 *)&addr;
+    port = ntohs(s->sin6_port);
+    ::inet_ntop(AF_INET6, &s->sin6_addr, ip.data(), sizeof(ip));
+  }
+
+  ::snprintf(out, out_len, "%s:%d", ip.data(), port);
+  return 0;
+}
+#endif
+
+int GetSockName(int sock, char* out, int out_len) {
+#ifdef _WIN32
+  return _get_sockname_of_win(sock, out, out_len);
+#else
+  return _get_sockname(sock, out, out_len);
+#endif
+}
+
+std::string GetSockName(int fd) {
+  std::array<char, 256> out = {};
+  GetSockName(fd, out.data(), sizeof(out));
+  return std::string(out.data());
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/store/socket.h

@@ -0,0 +1,26 @@
+// 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.
+
+#pragma once
+
+#include <string>
+
+namespace phi {
+namespace distributed {
+
+int GetSockName(int fd, char* out, int out_len);
+
+std::string GetSockName(int fd);
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/store.cc

@@ -0,0 +1,45 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/store/store.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace phi::distributed {
+
+int64_t Store::add(const std::string& key, int64_t value) {
+  PADDLE_THROW(
+      errors::InvalidArgument("Implement the add method in the subclass."));
+}
+
+std::vector<uint8_t> Store::get(const std::string& key) {
+  PADDLE_THROW(
+      errors::InvalidArgument("Implement the get method in the subclass."));
+}
+
+bool Store::check(const std::string& key) {
+  PADDLE_THROW(
+      errors::InvalidArgument("Implement the get method in the subclass."));
+}
+
+void Store::wait(const std::string& key) {
+  PADDLE_THROW(
+      errors::InvalidArgument("Implement the wait method in the subclass."));
+}
+
+void Store::set(const std::string& key, const std::vector<uint8_t>& value) {
+  PADDLE_THROW(
+      errors::InvalidArgument("Implement the set method in the subclass."));
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/store/store.h

@@ -0,0 +1,44 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+#include <chrono>
+#include <iostream>
+#include <string>
+#include <vector>
+#include "paddle/common/macros.h"
+
+namespace phi {
+namespace distributed {
+
+class PADDLE_API Store {
+ public:
+  Store() : _timeout(900) {}
+  explicit Store(const int timeout) : _timeout(timeout) {}
+  virtual ~Store() = default;
+
+  virtual int64_t add(const std::string& key, int64_t value);
+  virtual std::vector<uint8_t> get(const std::string& key);
+  virtual bool check(const std::string& key);
+  virtual void wait(const std::string& key);
+  virtual void set(const std::string& key, const std::vector<uint8_t>& value);
+
+  virtual int timeout() { return _timeout; }
+
+ protected:
+  int _timeout;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/store_utils.cc

@@ -0,0 +1,86 @@
+// Copyright (c) 2023 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.
+
+#include "paddle/phi/core/distributed/store/store_utils.h"
+
+#include <cstdlib>
+
+// the <winsock2.h> needs to be included before <winsock.h>, otherwise
+// there will be symbol redefinition error on windows
+#include "paddle/phi/core/distributed/store/tcp_store.h"
+
+#include "paddle/phi/core/distributed/auto_parallel/utils.h"
+
+namespace phi::distributed {
+using auto_parallel::str_split;
+
+namespace {
+std::string GetMasterEndpoint() {
+  const char* master_endpoint = std::getenv("PADDLE_MASTER");
+  if (!master_endpoint) {
+    const char* trainer_endpoints = std::getenv("PADDLE_TRAINER_ENDPOINTS");
+    PADDLE_ENFORCE_NOT_NULL(trainer_endpoints,
+                            common::errors::NotFound(
+                                "The environment variable "
+                                "'PADDLE_TRAINER_ENDPOINTS' cannot be found."));
+    return str_split(trainer_endpoints, ",")[0];
+  }
+
+  PADDLE_ENFORCE_NOT_NULL(
+      master_endpoint,
+      common::errors::NotFound(
+          "The environment variable 'PADDLE_MASTER' cannot be found."));
+  return master_endpoint;
+}
+}  // namespace
+
+int64_t GetCurGlobalRank() {
+  const char* cur_rank = std::getenv("PADDLE_TRAINER_ID");
+  if (cur_rank == nullptr) {
+    return 0;
+  }
+  return std::atoi(cur_rank);
+}
+
+int64_t GetGlobalWorldSize() {
+  const char* world_size = std::getenv("PADDLE_TRAINERS_NUM");
+  if (world_size == nullptr) {
+    return 1;
+  }
+  return std::atoi(world_size);
+}
+
+std::string GetMasterAddr() {
+  std::string master_endpoint = GetMasterEndpoint();
+  return str_split(master_endpoint, ":")[0];
+}
+
+uint16_t GetMasterPort() {
+  std::string master_endpoint = GetMasterEndpoint();
+  return std::stoi(str_split(master_endpoint, ":")[1]);
+}
+
+std::shared_ptr<Store> CreateOrGetGlobalTCPStore() {
+  std::string host = GetMasterAddr();
+  uint16_t port = GetMasterPort();
+  int64_t cur_rank = GetCurGlobalRank();
+  int64_t world_size = GetGlobalWorldSize();
+  bool is_master = (cur_rank == 0);
+
+  static std::shared_ptr<TCPStore> store =
+      std::make_shared<TCPStore>(host, port, is_master, world_size);
+  return store;
+}
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/store/store_utils.h

@@ -0,0 +1,37 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include "paddle/common/macros.h"
+
+namespace phi {
+namespace distributed {
+class Store;
+
+PADDLE_API int64_t GetCurGlobalRank();
+
+PADDLE_API std::string GetMasterAddr();
+
+PADDLE_API int64_t GetGlobalWorldSize();
+
+PADDLE_API uint16_t GetMasterPort();
+
+PADDLE_API std::shared_ptr<Store> CreateOrGetGlobalTCPStore();
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/tcp_store.cc

@@ -0,0 +1,494 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/core/distributed/store/tcp_store.h"
+
+#include <chrono>
+#include <iostream>
+#include <thread>
+
+#include "glog/logging.h"
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/distributed/store/tcp_utils.h"
+
+COMMON_DECLARE_bool(tcp_store_using_libuv);
+namespace phi::distributed::detail {
+
+// DaemonThread thread parent class methods
+DaemonThread::~DaemonThread() = default;
+
+void DaemonThread::start() {
+  daemonThread_ = std::thread{&DaemonThread::run, this};
+  is_running_.store(true);
+}
+
+void DaemonThread::cleanup() {
+  stop();
+  daemonThread_.join();
+}
+
+bool DaemonThread::is_running() { return is_running_.load(); }
+
+constexpr int INFTIME = 10000;  // 10 seconds
+
+std::unique_ptr<MasterDaemon> MasterDaemon::createDaemon(SocketType socket,
+                                                         int nranks,
+                                                         int timeout) {
+  VLOG(8) << ("begin to run start");
+  return std::make_unique<MasterDaemon>(socket, nranks, timeout);
+}
+
+MasterDaemon::MasterDaemon(SocketType socket, int nranks, int timeout)
+    : _listen_socket(socket), _nranks(nranks), _timeout(timeout) {
+  InitControlFd();
+}
+
+MasterDaemon::~MasterDaemon() {  // NOLINT
+  VLOG(8) << ("begin to destruct MasterDaemon");
+  StopByControlFd();
+  cleanup();
+  tcputils::close_socket(_listen_socket);
+  for (SocketType socket : _sockets) {
+    tcputils::close_socket(socket);
+  }
+  CloseControlFd();
+}
+
+void MasterDaemon::_do_add(SocketType socket) {
+  int64_t new_value{};
+  std::string key = tcputils::receive_string(socket);
+  new_value = tcputils::receive_value<int64_t>(socket);
+  std::vector<uint8_t> old_value;
+  auto it = _store.find(key);
+  if (it != _store.end()) {
+    old_value = it->second;
+    char* buffer = reinterpret_cast<char*>(it->second.data());
+    size_t len = old_value.size();
+    new_value += std::stoll(std::string(buffer, len));
+  }
+
+  std::string new_value_str = std::to_string(new_value);
+  _store[key] =
+      std::vector<uint8_t>(new_value_str.begin(), new_value_str.end());
+  VLOG(8) << "TCPStore: new value (" << new_value << ") for key (" << key
+          << ") " << GetSockName(socket);
+  tcputils::send_value<int64_t>(socket, new_value);
+  _notify_waiting_sockets(key);
+}
+
+void MasterDaemon::_do_set(SocketType socket) {
+  std::string key = tcputils::receive_string(socket);
+  VLOG(8) << "MasterDaemon::_do_set key(" << key << ") " << GetSockName(socket);
+
+  auto value = tcputils::receive_vector<uint8_t>(socket);
+  _store[key] = value;
+  _notify_waiting_sockets(key);
+}
+
+void MasterDaemon::_notify_waiting_sockets(const std::string& key) {
+  if (_waiting_sockets.find(key) != _waiting_sockets.end()) {
+    for (auto waiting_socket : _waiting_sockets.at(key)) {
+      auto reply = ReplyType::STOP_WAIT;
+      VLOG(7) << "TCPStore: notify the socket: " << GetSockName(waiting_socket)
+              << " that key: " << key << " is ready.";
+      tcputils::send_value<ReplyType>(waiting_socket, reply);
+    }
+    _waiting_sockets.erase(key);
+  }
+}
+
+void MasterDaemon::_do_get(SocketType socket) {
+  std::string key = tcputils::receive_string(socket);
+  VLOG(8) << "MasterDaemon::_do_get key(" << key << ") " << GetSockName(socket);
+
+  auto iter = _store.find(key);
+  PADDLE_ENFORCE_NE(
+      iter,
+      _store.end(),
+      common::errors::InvalidArgument("Key %s not found in TCPStore.", key));
+  std::vector<uint8_t> value = iter->second;
+  tcputils::send_vector<uint8_t>(socket, value);
+}
+
+void MasterDaemon::_do_check(SocketType socket) {
+  std::string key = tcputils::receive_string(socket);
+  VLOG(4) << "MasterDaemon::_do_check key(" << key << ") "
+          << GetSockName(socket);
+
+  auto iter = _store.find(key);
+  if (iter != _store.end()) {
+    tcputils::send_value<ReplyType>(socket, ReplyType::READY);
+  } else {
+    tcputils::send_value<ReplyType>(socket, ReplyType::NOT_READY);
+  }
+}
+
+#ifndef _WIN32
+void MasterDaemon::InitControlFd() {
+  PADDLE_ENFORCE_NE(
+      pipe(_control_fd.data()),
+      -1,
+      common::errors::Fatal("failed to cread control pipe errno:%d", errno));
+}
+void MasterDaemon::CloseControlFd() {
+  for (int fd : _control_fd) {
+    if (fd != -1) {
+      ::close(fd);
+    }
+  }
+}
+void MasterDaemon::StopByControlFd() {
+  VLOG(8) << ("begin to run StopByControlFd");
+  if (_control_fd[1] != -1) {
+    PADDLE_ENFORCE_NE(
+        ::write(_control_fd[1], "\0", 1),
+        -1,
+        common::errors::Fatal("failed to write control pipe errno:%d", errno));
+    // close the write end of the pipe
+    ::close(_control_fd[1]);
+    _control_fd[1] = -1;
+  }
+}
+#else
+void MasterDaemon::InitControlFd() {
+  ghStopEvent_ = CreateEvent(NULL, TRUE, FALSE, NULL);
+  PADDLE_ENFORCE_NE(ghStopEvent_,
+                    nullptr,
+                    common::errors::Fatal("failed to cread control pipe"));
+}
+void MasterDaemon::CloseControlFd() { CloseHandle(ghStopEvent_); }
+void MasterDaemon::StopByControlFd() { SetEvent(ghStopEvent_); }
+#endif
+
+void MasterDaemon::_do_wait(SocketType socket) {
+  std::string key = tcputils::receive_string(socket);
+  VLOG(8) << "MasterDaemon::_do_wait key(" << key << ") "
+          << GetSockName(socket);
+
+  auto iter = _store.find(key);
+  if (iter == _store.end()) {
+    // The key can not be found in store currently. Record and check later.
+    _waiting_sockets[key].emplace_back(socket);
+  } else {
+    auto reply = ReplyType::STOP_WAIT;
+    VLOG(7) << "TCPStore: wait reply (" << static_cast<int>(reply)
+            << ") for key (" << key << ").";
+    tcputils::send_value<ReplyType>(socket, reply);
+  }
+}
+
+void MasterDaemon::ProcessCommands(std::vector<struct pollfd>* p_fds) {
+  std::vector<struct pollfd>& fds = *p_fds;
+  // FIXME(gongwb): Don't loop all fds of set just the fds who have event.
+#ifdef _WIN32
+  // 0: listen socket, so loop from 1.
+  for (size_t i = 1; i < fds.size(); i++) {
+#else
+  // 0: listen socket, 1:controller pipe, so loop from 2.
+  for (uint i = 2; i < fds.size(); i++) {
+#endif
+    try {
+      if (fds[i].revents == 0) {
+        continue;
+      }
+
+      // VLOG(8) << "Plan to receive command from " << GetSockName(fds[i].fd);
+      Command command = tcputils::receive_value<Command>(fds[i].fd);
+      // VLOG(7) << "TCPStore: recv command: " << static_cast<int>(command) <<
+      // ".";
+
+      switch (command) {
+        case Command::ADD:
+          _do_add(fds[i].fd);
+          break;
+        case Command::GET:
+          _do_get(fds[i].fd);
+          break;
+        case Command::CHECK:
+          _do_check(fds[i].fd);
+          break;
+        case Command::SET:
+          _do_set(fds[i].fd);
+          break;
+        case Command::WAIT:
+          _do_wait(fds[i].fd);
+          break;
+        default:
+          break;
+          // VLOG(8) << "Unknown command: " << static_cast<int>(command)
+          //         << " from addr info:" << GetSockName(fds[i].fd);
+      }
+    } catch (const std::exception& ex) {
+      auto map_iter = _waiting_sockets.begin();
+      while (map_iter != _waiting_sockets.end()) {
+        auto vec_iter = map_iter->second.begin();
+        while (vec_iter != map_iter->second.end()) {
+          if (*vec_iter == fds[i].fd) {
+            vec_iter = map_iter->second.erase(vec_iter);
+          } else {
+            ++vec_iter;
+          }
+        }
+        if (map_iter->second.empty()) {
+          map_iter = _waiting_sockets.erase(map_iter);
+        } else {
+          ++map_iter;
+        }
+      }
+
+      tcputils::close_socket(fds[i].fd);
+      fds.erase(fds.begin() + i);
+#ifdef _WIN32
+      _sockets.erase(_sockets.begin() + i - 1);
+#else
+      _sockets.erase(_sockets.begin() + i - 2);
+#endif
+      std::string s(ex.what());
+      // if (s.find("TCP connection reset by peer") != std::string::npos) {
+      //   VLOG(5) << "TCP connection reset by peer";
+      // } else {
+      //   VLOG(5) << "Meet some exceptions during run:" << ex.what();
+      // }
+    }
+  }
+}
+
+void MasterDaemon::run() {
+  std::vector<struct pollfd> fds;
+#ifdef _WIN32
+  fds.push_back({_listen_socket, POLLIN});
+#else
+  fds.push_back({.fd = _listen_socket, .events = POLLIN, .revents = 0});
+  fds.push_back(
+      {.fd = _control_fd[0], .events = POLLIN | POLLHUP, .revents = 0});
+#endif
+
+  bool finished = false;
+  while (!finished) {
+    for (auto& item : fds) {
+      item.revents = 0;
+    }
+
+    VLOG(9) << "begin to poll fds_size:"
+            << paddle::string::Sprintf("%d", fds.size());
+#ifdef _WIN32
+    int res = ::WSAPoll(fds.data(), fds.size(), INFTIME);
+    if (res == 0) {
+      auto rv = WaitForSingleObject(ghStopEvent_, 0);
+      if (rv != WAIT_TIMEOUT) {
+        finished = true;
+        break;
+      }
+      continue;
+    }
+#else
+    ::poll(fds.data(), fds.size(), INFTIME);
+
+    VLOG(9) << "begin to fds[1].revents:"
+            << paddle::string::Sprintf("%d", fds[1].revents);
+    // The control pipe receive shutdown event, and begin to close it.
+    if (fds[1].revents != 0) {
+      if (fds[1].revents & ~(POLLIN | POLLHUP)) {
+        PADDLE_THROW(
+            common::errors::Fatal("Undefined event type:%d", fds[1].revents));
+      }
+      VLOG(0)
+          << "receive shutdown event and so quit from MasterDaemon run loop";
+      finished = true;  // NOLINT
+      break;
+    }
+#endif
+
+    // accept connect request.
+    if (fds[0].revents != 0) {
+      auto socket = tcputils::tcp_accept(_listen_socket);
+      _sockets.emplace_back(socket);
+#ifdef _WIN32
+      fds.push_back({socket, POLLIN});
+#else
+      fds.push_back({.fd = socket, .events = POLLIN, .revents = 0});
+#endif
+    }
+
+    ProcessCommands(&fds);
+  }
+}
+
+std::unique_ptr<TCPServer> TCPServer::create(uint16_t port,
+                                             int nranks,
+                                             int stop_check_timeout,
+                                             bool use_libuv) {
+  auto server = std::make_unique<TCPServer>();
+  if (use_libuv) {
+    // start libuv server
+    VLOG(0) << "create libuv server at port: " << port;
+    server->_master_daemon = create_libuv_tcpstore(port);
+    server->_master_daemon->start();
+  } else {
+    int socket = tcputils::tcp_listen("", std::to_string(port), AF_INET);
+    server->_master_daemon =
+        MasterDaemon::createDaemon(socket, nranks, stop_check_timeout);
+    server->_master_daemon->start();
+  }
+  return server;
+}
+
+std::unique_ptr<TCPClient> TCPClient::connect(const std::string host,
+                                              uint16_t port) {
+  int socket = tcputils::tcp_connect(host, std::to_string(port), AF_INET);
+  return std::make_unique<TCPClient>(socket);
+}
+
+void TCPClient::send_command_for_key(Command type, const std::string& key) {
+  tcputils::send_value<Command>(_socket, type);
+  if (key.empty()) {
+    return;
+  }
+  tcputils::send_string(_socket, key);
+}
+
+template <typename T>
+void TCPClient::send_value(const T& value) {
+  tcputils::send_bytes<T>(_socket, &value, 1);
+}
+
+template <typename T>
+T TCPClient::receive_value() {
+  T res;
+  tcputils::receive_bytes<T>(_socket, &res, 1);
+  return res;
+}
+
+template <typename T>
+void TCPClient::send_vector(const std::vector<T>& value) {
+  tcputils::send_vector<T>(_socket, value);
+}
+
+template <typename T>
+std::vector<T> TCPClient::receive_vector() {
+  return tcputils::receive_vector<T>(_socket);
+}
+
+}  // namespace phi::distributed::detail
+namespace phi::distributed {
+
+TCPStore::TCPStore(std::string host,
+                   uint16_t port,
+                   bool is_master,
+                   size_t num_workers,
+                   int timeout)
+    : Store(timeout),
+      _is_master(is_master),
+      _num_workers(static_cast<int>(num_workers)) {
+  _timeout = timeout;
+  PADDLE_ENFORCE_GT(
+      timeout,
+      0,
+      common::errors::InvalidArgument("timeout must >= %d", timeout));
+
+  VLOG(7) << "input timeout" << timeout << ", member timeout:" << _timeout;
+  if (_is_master) {
+    _server = detail::TCPServer::create(
+        port, this->_num_workers, timeout, FLAGS_tcp_store_using_libuv);
+  }
+
+  _client = detail::TCPClient::connect(host, port);
+  waitWorkers();
+}
+
+void TCPStore::waitWorkers() {
+  if (_num_workers == 0) {
+    return;
+  }
+  add(_init_key, 1);
+
+  if (_is_master) {
+    VLOG(7) << paddle::string::Sprintf("_timeout:%d", _timeout);
+    auto begin = std::chrono::steady_clock::now();
+    do {
+      auto value = get(_init_key);
+      int completed = std::stoi(std::string(value.begin(), value.end()));
+      VLOG(7) << completed << " worker ready, total " << _num_workers
+              << ", _timeout:" << _timeout;
+      if (completed >= _num_workers) {
+        break;
+      }
+      const auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(
+          std::chrono::steady_clock::now() - begin);
+
+      std::this_thread::sleep_for(std::chrono::milliseconds(10));
+      if (_timeout != 0 && elapsed.count() > _timeout) {
+        PADDLE_THROW(common::errors::Fatal(paddle::string::Sprintf(
+            "_timeout:%d elapsed:%d (elapsed > _timeout)=%d",
+            _timeout,
+            elapsed.count(),
+            elapsed.count() > _timeout)));
+
+        PADDLE_ENFORCE_EQ(
+            completed,
+            _num_workers,
+            common::errors::InvalidArgument(
+                "TCPStore timeouted and not all workers got ready."));
+      }
+    } while (true);
+  }
+  VLOG(7) << "TCPStore initialized.";
+}
+
+int64_t TCPStore::add(const std::string& key, int64_t value) {
+  VLOG(7) << "TCPStore add.";
+  _client->send_command_for_key(Command::ADD, _key_prefix + key);
+  _client->send_value<std::int64_t>(value);
+  return _client->receive_value<std::int64_t>();
+}
+
+void TCPStore::set(const std::string& key, const std::vector<uint8_t>& value) {
+  VLOG(7) << "TCPStore set.";
+  _client->send_command_for_key(Command::SET, _key_prefix + key);
+  _client->send_vector<uint8_t>(value);
+}
+
+std::vector<uint8_t> TCPStore::get(const std::string& key) {
+  wait(key);
+  _client->send_command_for_key(Command::GET, _key_prefix + key);
+  VLOG(7) << "TCPStore get.";
+  return _client->receive_vector<uint8_t>();
+}
+
+bool TCPStore::check(const std::string& key) {
+  _client->send_command_for_key(Command::CHECK, _key_prefix + key);
+  VLOG(3) << "TCPStore check.";
+  auto response = _client->receive_value<ReplyType>();
+  if (response == ReplyType::READY) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+void TCPStore::wait(const std::string& key) {
+  ReplyType reply;  // NOLINT
+  VLOG(7) << "TCPStore wait.";
+  _client->send_command_for_key(Command::WAIT, _key_prefix + key);
+  reply = _client->receive_value<ReplyType>();
+  PADDLE_ENFORCE_EQ(
+      reply == ReplyType::STOP_WAIT,
+      true,
+      common::errors::InvalidArgument("Stop_waiting response is expected"));
+}
+
+TCPStore::~TCPStore() { VLOG(7) << "TCPStore destructure"; }
+
+}  // namespace phi::distributed

paddle/phi/core/distributed/store/tcp_store.h

@@ -0,0 +1,179 @@
+// 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.
+
+#pragma once
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/types.h>
+
+#ifndef _WIN32
+#include <unistd.h>
+#endif
+
+#include <array>
+#include <atomic>
+#include <iostream>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <thread>
+#include <unordered_map>
+
+#include "paddle/phi/core/distributed/store/socket.h"
+#include "paddle/phi/core/distributed/store/store.h"
+#include "paddle/phi/core/distributed/store/tcp_utils.h"
+
+namespace phi {
+namespace distributed {
+
+enum class ReplyType { WAITING, STOP_WAIT, READY, NOT_READY };
+enum class Command { ADD, GET, CHECK, SET, WAIT, STOP };
+
+namespace detail {
+
+// Abstract base class to handle thread state for TCPStoreMasterDaemon.
+// Contains the windows/unix implementations to signal a
+// shutdown sequence for the thread
+class DaemonThread {
+ public:
+  DaemonThread() = default;
+  virtual ~DaemonThread() = 0;
+  void start();
+
+ protected:
+  void cleanup();
+  virtual void run() = 0;
+  virtual void stop() = 0;
+  bool is_running();
+
+ private:
+  std::atomic<bool> is_running_{false};
+  std::thread daemonThread_{};
+};
+
+std::unique_ptr<DaemonThread> create_libuv_tcpstore(const std::uint16_t& port);
+
+class MasterDaemon : public DaemonThread {
+ public:
+  static std::unique_ptr<MasterDaemon> createDaemon(SocketType listen_socket,
+                                                    int nranks,
+                                                    int timeout);
+  MasterDaemon() = delete;
+  explicit MasterDaemon(SocketType listen_socket,
+                        int nranks,
+                        int stop_check_timeout);
+  ~MasterDaemon() override;
+
+ protected:
+  void run() override;
+  void stop() override{};
+
+ private:
+  void ProcessCommands(std::vector<struct pollfd>* p_fds);
+  void _do_add(SocketType socket);
+  void _do_wait(SocketType socket);
+  void _do_get(SocketType socket);
+  void _do_check(SocketType socket);
+  void _do_set(SocketType socket);
+  void _notify_waiting_sockets(const std::string&);
+  SocketType _listen_socket;
+  std::vector<SocketType> _sockets;
+  std::unordered_map<std::string, std::vector<uint8_t>> _store;
+  std::thread _background_thread{};
+  int _nranks = -1;
+  int _timeout = 0;
+  std::unordered_map<std::string, std::vector<SocketType>>
+      _waiting_sockets;  // key -> list of waiting sockets
+
+  void InitControlFd();
+  void CloseControlFd();
+  void StopByControlFd();
+#ifdef _WIN32
+  HANDLE ghStopEvent_{};
+#else
+  std::array<int, 2> _control_fd{{-1, -1}};
+#endif
+};
+
+class TCPServer {
+ public:
+  TCPServer() = default;
+  static std::unique_ptr<TCPServer> create(std::uint16_t port,
+                                           int nranks,
+                                           int stop_check_timeout,
+                                           bool use_libuv);
+
+ private:
+  std::unique_ptr<DaemonThread> _master_daemon;
+};
+
+class TCPClient {
+ public:
+  explicit TCPClient(SocketType socket) : _socket{socket} {}
+  static std::unique_ptr<TCPClient> connect(const std::string host,
+                                            uint16_t port);
+  ~TCPClient() { tcputils::close_socket(_socket); }
+  void send_command_for_key(Command type, const std::string& key);
+
+  template <typename T>
+  void send_value(const T& value);
+
+  template <typename T>
+  void send_vector(const std::vector<T>& value);
+  template <typename T>
+  std::vector<T> receive_vector();
+
+  template <typename T>
+  T receive_value();
+
+ private:
+  SocketType _socket;
+};
+
+}  // namespace detail
+
+// TODO(gongwb) :Add IP6 support.
+class PADDLE_API TCPStore : public Store {
+ public:
+  static constexpr std::uint16_t kDefaultPort = 6170;
+  explicit TCPStore(std::string host,
+                    uint16_t port = kDefaultPort,
+                    bool is_master = false,
+                    size_t num_workers = 1,
+                    int timeout = 900);
+
+  ~TCPStore();
+
+  int64_t add(const std::string& key, int64_t value) override;
+  std::vector<uint8_t> get(const std::string& key) override;
+  bool check(const std::string& key) override;
+  void wait(const std::string& key) override;
+  void set(const std::string& key, const std::vector<uint8_t>& value) override;
+
+ private:
+  void waitWorkers();
+  std::unique_ptr<detail::TCPServer> _server;
+  std::unique_ptr<detail::TCPClient> _client;
+
+  const std::string _init_key = "init/";
+  const std::string _key_prefix = "/";
+
+  bool _is_master;
+  int _num_workers;
+};
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/core/distributed/store/tcp_store_libuv.cc

@@ -0,0 +1,738 @@
+// Copyright (c) 2025 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.
+#include "glog/logging.h"
+
+#include "paddle/phi/core/distributed/store/tcp_store_libuv.h"
+
+namespace phi::distributed::detail {
+
+// SegmentedDataStream
+void SegmentedDataStream::append(uv_buf_t buf) {
+  if (buf.len == 0) {
+    free(buf.base);
+  } else {
+    capacity += buf.len;
+    _buffers.push_back(buf);
+  }
+}
+
+bool SegmentedDataStream::readMany(char* dest, size_t size) {
+  if (available() < size) {
+    return false;
+  }
+
+  size_t remaining = size;
+  char* write_base = dest;
+  while (remaining > 0) {
+    auto to_read = std::min(_buffers[_buff_idx].len - _buff_offset, remaining);
+    ::memcpy(write_base, _buffers[_buff_idx].base + _buff_offset, to_read);
+    _buff_offset += to_read;
+    remaining -= to_read;
+    write_base += to_read;
+    if (_buff_offset >= _buffers[_buff_idx].len) {
+      _buff_offset = 0;
+      ++_buff_idx;
+      if (_buff_idx >= _buffers.size() && remaining > 0) {
+        PADDLE_THROW(common::errors::Fatal(paddle::string::Sprintf(
+            "Read operation exceeds buffer boundary. ",
+            "buffer index: %d, available: %d, remaining: %d",
+            _buff_idx,
+            _buffers.size(),
+            remaining)));
+      }
+    }
+  }
+  _read_offset += size;
+  return true;
+}
+
+template <typename T>
+bool SegmentedDataStream::readValue(T& value) {
+  return readMany(reinterpret_cast<char*>(&value), sizeof(T));
+}
+
+bool SegmentedDataStream::readKey(std::string& str) {
+  uint64_t size = 0;
+  if (!readValue(size)) return false;
+  PADDLE_ENFORCE_LE(size,
+                    phi::distributed::detail::MAX_KEY_LEN,
+                    common::errors::InvalidArgument(paddle::string::Sprintf(
+                        "Key size validation failed. size: %d, max: %d",
+                        size,
+                        phi::distributed::detail::MAX_KEY_LEN)));
+
+  if (available() < size) return false;
+  str.resize(size);
+  return readMany(reinterpret_cast<char*>(str.data()), size);
+}
+
+bool SegmentedDataStream::readContent(std::vector<uint8_t>& data) {
+  uint64_t size = 0;
+  if (!readValue(size)) return false;
+  auto size_in_bytes = size * sizeof(uint8_t);
+  PADDLE_ENFORCE_LE(size_in_bytes,
+                    MAX_CONTENT_LEN,
+                    common::errors::InvalidArgument(paddle::string::Sprintf(
+                        "Content size validation failed. size: %d, max: %d",
+                        size_in_bytes,
+                        MAX_CONTENT_LEN)));
+
+  if (available() < size_in_bytes) return false;
+  data.resize(size);
+  return readMany(reinterpret_cast<char*>(data.data()), size_in_bytes);
+}
+
+size_t SegmentedDataStream::available() { return capacity - _read_offset; }
+
+void SegmentedDataStream::commit() {
+  if (_buff_idx >= _buffers.size() || _buff_offset >= _buffers[_buff_idx].len) {
+    _buff_offset = 0;
+    if (_buff_idx < _buffers.size()) ++_buff_idx;
+  }
+
+  for (size_t i = 0; i < _buff_idx; ++i) {
+    free(_buffers[0].base);
+    capacity -= _buffers[0].len;
+    _buffers.pop_front();
+  }
+  _buff_idx = 0;
+  _read_offset = _buff_offset_commit = _buff_offset;
+}
+
+void SegmentedDataStream::reset() {
+  _buff_idx = 0;
+  _read_offset = _buff_offset = _buff_offset_commit;
+}
+
+// LibUVHandle
+std::shared_ptr<LibUVHandle> LibUVHandle::ptr() { return shared_from_this(); }
+
+void LibUVHandle::close() {
+  if (uv_is_closing(getRawHandle())) {
+    return;
+  }
+  uv_close(getRawHandle(), handleClose);
+}
+
+void LibUVHandle::handleAvailable() {
+  uv_handle_set_data(getRawHandle(), this);
+}
+
+void LibUVHandle::handleClose(uv_handle_t* uv_handle) {
+  auto h = reinterpret_cast<LibUVHandle*>(uv_handle_get_data(uv_handle));
+  h->onClose();
+}
+
+//   ==== LibUVTCPSocket ====
+LibUVTCPSocket::LibUVTCPSocket(uv_loop_t* loop) {
+  uv_tcp_init(loop, &client);
+  if (int err = uv_tcp_nodelay(&client, 1)) {
+    VLOG(2) << "The no-delay option is unavailable. err: " << err;
+  }
+}
+
+uv_handle_t* LibUVTCPSocket::getRawHandle() {
+  return reinterpret_cast<uv_handle_t*>(&client);
+}
+
+std::shared_ptr<LibUVTCPSocket> LibUVTCPSocket::ptr() {
+  return std::static_pointer_cast<LibUVTCPSocket>(shared_from_this());
+}
+
+std::shared_ptr<LibUVTCPSocket> LibUVTCPSocket::getTCPSocket(
+    uv_stream_t* handle) {
+  auto h = reinterpret_cast<LibUVTCPSocket*>(
+      uv_handle_get_data(reinterpret_cast<uv_handle_t*>(handle)));
+  return h->ptr();
+}
+
+//   LibUVTCPServer
+void LibUVTCPServer::setCallback(LibUVCallback&& callback) {
+  _on_connect_callback = std::move(callback);
+}
+
+std::shared_ptr<LibUVTCPServer> LibUVTCPServer::createServer(uv_loop_t* loop,
+                                                             std::uint16_t port,
+                                                             bool useIpv6) {
+  auto res = std::make_shared<LibUVTCPServer>(loop);
+  res->handleAvailable();
+  try {
+    struct sockaddr_storage addr {};
+    int uv_res = 0;
+    if (useIpv6) {
+      uv_res = uv_ip6_addr("::", port, (struct sockaddr_in6*)&addr);
+    } else {
+      uv_res = uv_ip4_addr("0.0.0.0", port, (struct sockaddr_in*)&addr);
+    }
+    PADDLE_ENFORCE_EQ(uv_res,
+                      0,
+                      common::errors::InvalidArgument(paddle::string::Sprintf(
+                          "sockaddr parsing failure. port: %d, useIpv6:%d, "
+                          "code: %d, name: %s, message: %s",
+                          port,
+                          useIpv6,
+                          uv_res,
+                          uv_err_name(uv_res),
+                          uv_strerror(uv_res))));
+
+    uv_res =
+        uv_tcp_bind(res->getRawSocket(), (const struct ::sockaddr*)&addr, 0);
+    PADDLE_ENFORCE_EQ(
+        uv_res,
+        0,
+        common::errors::InvalidArgument(paddle::string::Sprintf(
+            "Bind operation failed for the server socket. port: %d, "
+            "useIpv6: %d, code: %d, name: %s, message: %s",
+            port,
+            useIpv6,
+            uv_res,
+            uv_err_name(uv_res),
+            uv_strerror(uv_res))));
+
+    uv_res = uv_listen(
+        res->getRawStream(), FLAGS_tcp_max_syn_backlog, onNewConnection);
+    PADDLE_ENFORCE_EQ(
+        uv_res,
+        0,
+        common::errors::InvalidArgument(paddle::string::Sprintf(
+            "Server socket unable to listen on local network interfaces. "
+            "port: %d, useIpv6: %d, code: %d, name: %s, message: %s",
+            port,
+            useIpv6,
+            uv_res,
+            uv_err_name(uv_res),
+            uv_strerror(uv_res))));
+    res->setSocketPort();
+  } catch (std::exception& ex) {
+    res->close();
+    throw;
+  }
+  return res;
+}
+
+void LibUVTCPServer::accept(const std::shared_ptr<LibUVTCPSocket>& socket) {
+  int res = uv_accept(getRawStream(),
+                      reinterpret_cast<uv_stream_t*>(socket->getRawHandle()));
+  PADDLE_ENFORCE_EQ(
+      res,
+      0,
+      common::errors::InvalidArgument(paddle::string::Sprintf(
+          "Socket accept operation failed. code: %d, name: %s, message: %s",
+          res,
+          uv_err_name(res),
+          uv_strerror(res))));
+}
+
+void LibUVTCPServer::setSocketPort() {
+  sockaddr_storage addr_s{};
+  int addr_len = sizeof(addr_s);
+  if (uv_tcp_getsockname(reinterpret_cast<uv_tcp_t*>(getRawStream()),
+                         reinterpret_cast<::sockaddr*>(&addr_s),
+                         &addr_len) != 0) {
+    throw std::runtime_error("the port number cannot be retrieved.");
+  }
+  if (addr_s.ss_family == AF_INET) {
+    _port = ntohs(reinterpret_cast<sockaddr_in*>(&addr_s)->sin_port);
+  } else {
+    _port = ntohs(reinterpret_cast<sockaddr_in6*>(&addr_s)->sin6_port);
+  }
+}
+
+void LibUVTCPServer::onNewConnection(uv_stream_t* server, int status) {
+  auto h = reinterpret_cast<LibUVTCPServer*>(
+      uv_handle_get_data(reinterpret_cast<uv_handle_t*>(server)));
+  h->_on_connect_callback(status);
+}
+
+// WriteUVContent
+WriteUVContent::WriteUVContent(std::vector<uint8_t>&& in_data,
+                               std::shared_ptr<LibUVHandle> handle)
+    : data(std::move(in_data)), handle(std::move(handle)) {
+  uv_req_set_data(reinterpret_cast<uv_req_t*>(&req), new RequestData());
+}
+
+void WriteUVContent::writeDone(uv_write_t* req, int status) {
+  auto data_ptr = static_cast<RequestData*>(
+      uv_req_get_data(reinterpret_cast<uv_req_t*>(req)));
+  if (!data_ptr) return;
+
+  auto self = std::move(data_ptr->strong_self);
+  delete data_ptr;
+  uv_req_set_data(reinterpret_cast<uv_req_t*>(req), nullptr);
+  if (self && status) {
+    VLOG(2) << "Write to client failed. code:" << status
+            << " desc:" << uv_strerror(status)
+            << " name:" << uv_err_name(status);
+    self->handle->close();
+  }
+}
+
+WriteUVContent::~WriteUVContent() {
+  // safely clean up pending request data
+  if (auto data = static_cast<RequestData*>(
+          uv_req_get_data(reinterpret_cast<uv_req_t*>(&req)))) {
+    delete data;
+    uv_req_set_data(reinterpret_cast<uv_req_t*>(&req), nullptr);
+  }
+}
+
+void WriteUVContent::send() {
+  if (data.empty()) return;
+  buf = uv_buf_init(reinterpret_cast<char*>(data.data()), data.size());
+  int res = uv_write(&req,
+                     reinterpret_cast<uv_stream_t*>(handle->getRawHandle()),
+                     &buf,
+                     1,
+                     writeDone);
+  if (res) {
+    VLOG(2) << "Write failed. code:" << res << " desc:" << uv_strerror(res)
+            << " name:" << uv_err_name(res);
+    handle->close();
+  } else {
+    auto data_ptr = static_cast<RequestData*>(
+        uv_req_get_data(reinterpret_cast<uv_req_t*>(&req)));
+    if (data_ptr) {
+      data_ptr->strong_self = shared_from_this();
+    }
+  }
+}
+
+// UVWriter
+template <typename T>
+void UVWriter::writeValue(T val) {
+  uint8_t* val_ptr = reinterpret_cast<uint8_t*>(&val);
+  data.insert(data.end(), val_ptr, val_ptr + sizeof(T));
+}
+
+void UVWriter::writeVector(const std::vector<uint8_t>& val) {
+  writeValue<uint64_t>(val.size());
+  data.insert(data.end(), val.begin(), val.end());
+}
+
+void UVWriter::writeString(const std::string& val) {
+  writeValue<uint64_t>(val.size());
+  data.insert(data.end(), val.data(), val.data() + val.size());
+}
+
+void UVWriter::send() {
+  auto wd = std::make_shared<WriteUVContent>(std::move(data), handle);
+  wd->send();
+}
+
+// LibUVClient
+void LibUVClient::allocBuffer(uv_handle_t* handle,
+                              size_t buf_size,
+                              uv_buf_t* buf) {
+  buf_size = std::min(buf_size, MAX_BUFFER_SIZE);
+  buf->base = reinterpret_cast<char*>(malloc(buf_size));
+  buf->len = buf_size;
+}
+
+void LibUVClient::readCallback(uv_stream_t* client,
+                               ssize_t nread,
+                               const uv_buf_t* buf) {
+  auto uv_socket = LibUVTCPSocket::getTCPSocket(client);
+  if (nread > 0) {
+    try {
+      uv_socket->doProcess(buf, nread);
+      return;
+    } catch (std::exception& ex) {
+      VLOG(2) << "Failed to process incoming client message: " << ex.what();
+      uv_socket->close();
+    }
+  } else if (nread == UV_EOF) {
+    // EOF
+    VLOG(5) << "Remote peer closed the connection.";
+    uv_socket->close();
+  } else if (nread < 0) {
+    // error and EOF
+    VLOG(5) << "Read callback handler exception. code:" << nread
+            << " desc:" << uv_strerror(nread) << " name:" << uv_err_name(nread);
+    uv_socket->close();
+  }
+  free(buf->base);
+}
+
+void LibUVClient::doProcess(const uv_buf_t* buf, size_t nread) {
+  auto tmp = *buf;
+  tmp.len = nread;
+  stream.append(tmp);
+
+  VLOG(5) << "process: " << std::string(buf->base, nread)
+          << ", nread: " << nread;
+  while (true) {
+    stream.reset();
+    uint32_t command = -1;
+    if (!stream.readValue(command)) break;
+
+    VLOG(5) << "Client parse command" << command;
+    switch ((Command)command) {
+      case Command::ADD:
+        if (!doAddCommand()) return;
+        break;
+      case Command::GET:
+        if (!doGetCommand()) return;
+        break;
+      case Command::CHECK:
+        if (!doCheckCommand()) return;
+        break;
+      case Command::SET:
+        if (!doSetCommand()) return;
+        break;
+      case Command::WAIT:
+        if (!doWaitCommand()) return;
+        break;
+      default:
+        VLOG(4) << "invalid command from Client, command: " << command;
+        close();
+        return;
+    }
+    stream.commit();
+  }
+}
+
+bool LibUVClient::doSetCommand() {
+  std::string key;
+  if (!stream.readKey(key)) return false;
+
+  std::vector<uint8_t> newData;
+  if (!stream.readContent(newData)) return false;
+  VLOG(7) << "set key:" << key << " address:" << this->address();
+  store->set(key, newData);
+  return true;
+}
+
+bool LibUVClient::doGetCommand() {
+  std::string key;
+  if (!stream.readKey(key)) return false;
+
+  VLOG(7) << "get key: " << key << " address:" << this->address();
+  const auto& data = store->get(key);
+  UVWriter sw(ptr());
+  sw.writeVector(data);
+  sw.send();
+  return true;
+}
+
+bool LibUVClient::doAddCommand() {
+  std::string key;
+  if (!stream.readKey(key)) return false;
+  int64_t addVal = 0;
+  if (!stream.readValue(addVal)) return false;
+
+  addVal = store->add(key, addVal);
+  VLOG(7) << "add key:" << key << " val: " << addVal
+          << " address:" << this->address();
+  UVWriter sw(ptr());
+  sw.writeValue(addVal);
+  sw.send();
+  return true;
+}
+
+bool LibUVClient::doCheckCommand() {
+  std::string key;
+  if (!stream.readKey(key)) return false;
+
+  VLOG(7) << "check key:" << key << " address:" << this->address();
+  std::vector<std::string> keys = {key};
+  UVWriter sw(ptr());
+  if (store->checkKeys(keys)) {
+    sw.writeValue(ReplyType::READY);
+  } else {
+    sw.writeValue(ReplyType::NOT_READY);
+  }
+  sw.send();
+  return true;
+}
+
+bool LibUVClient::doWaitCommand() {
+  std::string key;
+  if (!stream.readKey(key)) return false;
+
+  VLOG(7) << "wait key:  " << key << " address:" << this->address();
+  if (store->waitKey(key, ptr())) {
+    UVWriter sw(ptr());
+    sw.writeValue(ReplyType::STOP_WAIT);
+    sw.send();
+    VLOG(7) << "wait send:  " << key;
+  }
+  return true;
+}
+
+void LibUVClient::onClose() { store->removeClient(ptr()); }
+
+PADDLE_API std::string fmtSockAddr(const struct ::sockaddr* addr,
+                                   socklen_t len) {
+  char host[NI_MAXHOST], port[NI_MAXSERV];  // NOLINT
+  int flags = NI_NUMERICSERV;
+  int err =
+      ::getnameinfo(addr, len, host, sizeof(host), port, sizeof(port), flags);
+  if (err) {
+    VLOG(1) << "Cannot resolve hostname, fallback to numeric. Error: " << err;
+    // fallback to numeric
+    flags |= NI_NUMERICHOST;
+    err =
+        ::getnameinfo(addr, len, host, sizeof(host), port, sizeof(port), flags);
+    if (err) {
+      VLOG(1) << "Numeric address resolution failed. Error: " << err;
+      return "?UNKNOWN?";
+    }
+  }
+  switch (addr->sa_family) {
+    case AF_INET:
+      return paddle::string::Sprintf("%s:%s", host, port);
+    case AF_INET6:
+      return paddle::string::Sprintf("[%s]:%s", host, port);
+    default:
+      return paddle::string::Sprintf("[%s]:%s", host, port);
+  }
+}
+
+void LibUVClient::readStart() {
+  struct ::sockaddr_storage addr {};
+  int addrLen{sizeof(struct ::sockaddr_storage)};
+
+  if (int err = uv_tcp_getpeername(
+          &client, reinterpret_cast<struct ::sockaddr*>(&addr), &addrLen)) {
+    VLOG(2) << "Remote endpoint resolution failed. err=" << uv_strerror(err);
+  } else {
+    _address =
+        fmtSockAddr(reinterpret_cast<struct ::sockaddr*>(&addr), addrLen);
+  }
+  int res = uv_read_start(
+      reinterpret_cast<uv_stream_t*>(&client), allocBuffer, readCallback);
+  if (res) {
+    VLOG(2) << "Read callback initialization failure. client:"
+            << reinterpret_cast<void*>(this) << " code:" << res
+            << " desc:" << uv_strerror(res) << " name:" << uv_err_name(res);
+    close();
+  }
+}
+
+std::shared_ptr<LibUVClient> LibUVClient::make(uv_loop_t* loop,
+                                               LibUVMasterDaemon* store) {
+  auto res = std::make_shared<LibUVClient>(loop, store);
+  res->handleAvailable();
+  return res;
+}
+
+std::shared_ptr<LibUVClient> LibUVClient::ptr() {
+  return std::static_pointer_cast<LibUVClient>(shared_from_this());
+}
+
+//  LibUVMasterDaemon
+void LibUVMasterDaemon::onConnect(int status) {
+  auto client = LibUVClient::make(&loop_, this);
+  addClient(client);
+  try {
+    _tcp_server->accept(client);
+    client->readStart();
+  } catch (std::exception& e) {
+    VLOG(2) << "Accept client failed, err: " << e.what();
+    client->close();
+  }
+}
+
+void LibUVMasterDaemon::onExitRequest() {
+  VLOG(4) << "begin to exit requested";
+  uv_close(reinterpret_cast<uv_handle_t*>(&_exit_handle), nullptr);
+  uv_stop(&loop_);
+}
+
+void LibUVMasterDaemon::init(const std::uint16_t& port) {
+  try {
+    _tcp_server = LibUVTCPServer::createServer(&loop_, port, /*useIpv6=*/false);
+  } catch (std::exception& ex) {
+    PADDLE_THROW(common::errors::Fatal(
+        paddle::string::Sprintf("Bind to ipv4 address failed: %s", ex.what())));
+  }
+  _tcp_server->setCallback([this](auto status) { this->onConnect(status); });
+
+  port_ = _tcp_server->port();
+  PADDLE_ENFORCE_EQ(
+      port_,
+      port,
+      common::errors::InvalidArgument(paddle::string::Sprintf(
+          "listen fd is bound to port %d, but expected port %d", port_, port)));
+}
+
+LibUVMasterDaemon::LibUVMasterDaemon(int port) : port_(port) {
+  // uv loop init
+  PADDLE_ENFORCE_EQ(uv_loop_init(&loop_),
+                    0,
+                    common::errors::InvalidArgument("init libuv loop failed"));
+  // uv async init
+  PADDLE_ENFORCE_EQ(
+      uv_async_init(&loop_, &_exit_handle, LibUVMasterDaemon::on_exit_request),
+      0,
+      common::errors::InvalidArgument("init libuv async event failed"));
+  uv_handle_set_data(reinterpret_cast<uv_handle_t*>(&_exit_handle), this);
+}
+
+LibUVMasterDaemon::~LibUVMasterDaemon() {
+  if (!is_running()) {
+    uv_close(reinterpret_cast<uv_handle_t*>(&_exit_handle), nullptr);
+    uv_run(&loop_, UV_RUN_NOWAIT);
+    if (uv_loop_close(&loop_) != 0) {
+      VLOG(0) << "uv loop close failed";
+    }
+  } else {
+    // the daemon thread cleanup libuv
+    cleanup();
+  }
+}
+
+void LibUVMasterDaemon::run() {
+  VLOG(4) << "start LibUV master daemon loop";
+  int res = uv_run(&loop_, UV_RUN_DEFAULT);
+  if (res) {
+    VLOG(4) << "LibUV master daemon loop done: " << res;
+  }
+
+  for (const auto& client : _clients) {
+    client->close();
+  }
+  _tcp_server->close();
+
+  while (true) {
+    res = uv_loop_close(&loop_);
+    if (res == 0) {
+      break;
+    }
+    VLOG(3) << "uv_loop_close failed with:" << res
+            << " err: " << uv_err_name(res)
+            << " std error:" << uv_strerror(res);
+    res = uv_run(&loop_, UV_RUN_NOWAIT);
+    if (res != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(300));
+    }
+  }
+  VLOG(3) << "LibUV master daemon loop cleanup finished.";
+}
+
+void LibUVMasterDaemon::stop() {
+  int res = uv_async_send(&_exit_handle);
+  if (res) {
+    VLOG(2) << "stop with uv_async_send failed:" << res
+            << " err:" << uv_err_name(res) << " std error:" << uv_strerror(res);
+  }
+}
+
+void LibUVMasterDaemon::addClient(const std::shared_ptr<LibUVHandle>& client) {
+  _clients.insert(client);
+}
+
+void LibUVMasterDaemon::removeClient(
+    const std::shared_ptr<LibUVHandle>& client) {
+  _clients.erase(client);
+  clearWaitState(client);
+}
+
+void LibUVMasterDaemon::clearWaitState(
+    const std::shared_ptr<LibUVHandle>& client) {
+  if (_awaited_keys.find(client) == _awaited_keys.end()) {
+    return;
+  }
+  _awaited_keys.erase(client);
+  for (auto it = _waiting_sockets.begin(); it != _waiting_sockets.end();) {
+    for (auto vecIt = it->second.begin(); vecIt != it->second.end();) {
+      if (*vecIt == client) {
+        vecIt = it->second.erase(vecIt);
+      } else {
+        ++vecIt;
+      }
+    }
+    if (it->second.empty()) {
+      it = _waiting_sockets.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
+void LibUVMasterDaemon::set(const std::string& key,
+                            const std::vector<uint8_t>& value) {
+  _tcp_store[key] = value;
+  // notify all clients that have been waiting
+  notifyWaitingClients(key);
+}
+
+const std::vector<uint8_t>& LibUVMasterDaemon::get(const std::string& key) {
+  static std::vector<uint8_t> missing_key;
+  return _tcp_store.count(key) ? _tcp_store.at(key) : missing_key;
+}
+
+int64_t LibUVMasterDaemon::add(const std::string& key, int64_t addVal) {
+  std::vector<uint8_t> old_data;
+  auto it = _tcp_store.find(key);
+  if (it != _tcp_store.end()) {
+    old_data = it->second;
+    auto buf = reinterpret_cast<const char*>(it->second.data());
+    auto len = it->second.size();
+    addVal += std::stoll(std::string(buf, len));
+  }
+  auto addValStr = std::to_string(addVal);
+  std::vector<uint8_t> newData =
+      std::vector<uint8_t>(addValStr.begin(), addValStr.end());
+  _tcp_store[key] = newData;
+
+  // notify all clients that have been waiting
+  notifyWaitingClients(key);
+  return addVal;
+}
+
+bool LibUVMasterDaemon::checkKeys(const std::vector<std::string>& keys) {
+  return std::all_of(keys.begin(), keys.end(), [&](const std::string& s) {
+    if (_tcp_store.count(s) > 0) {
+      return true;
+    }
+    return false;
+  });
+}
+
+bool LibUVMasterDaemon::waitKey(const std::string& key,
+                                const std::shared_ptr<LibUVHandle>& client) {
+  int num_to_await = 0;
+  if (_tcp_store.find(key) == _tcp_store.end()) {
+    _waiting_sockets[key].push_back(client);
+    num_to_await++;
+    VLOG(7) << "add to wait key:  " << key;
+  } else {
+    return true;
+  }
+  _awaited_keys[client] = num_to_await;
+  return false;
+}
+
+void LibUVMasterDaemon::notifyWaitingClients(const std::string& key) {
+  auto sockets_to_wait = _waiting_sockets.find(key);
+  if (sockets_to_wait != _waiting_sockets.end()) {
+    for (const auto& client : sockets_to_wait->second) {
+      if (--_awaited_keys[client] == 0) {
+        UVWriter sw(client->ptr());
+        sw.writeValue(ReplyType::STOP_WAIT);
+        sw.send();
+      }
+    }
+    _waiting_sockets.erase(sockets_to_wait);
+  }
+}
+
+std::unique_ptr<phi::distributed::detail::DaemonThread> create_libuv_tcpstore(
+    const std::uint16_t& port) {
+  auto res = std::make_unique<LibUVMasterDaemon>(port);
+  res->init(port);
+  return res;
+}
+}  // namespace phi::distributed::detail

paddle/phi/core/distributed/store/tcp_store_libuv.h

@@ -0,0 +1,242 @@
+// Copyright (c) 2025 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.
+
+#pragma once
+
+#include <uv.h>
+
+#include <algorithm>
+#include <cstdio>
+#include <deque>
+#include <exception>
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "paddle/common/flags.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/distributed/store/tcp_store.h"
+#include "paddle/phi/core/distributed/store/tcp_utils.h"
+
+namespace phi::distributed::detail {
+auto constexpr MAX_KEY_LEN = 16 * 1024;
+auto constexpr MAX_CONTENT_LEN = 16 * 1024 * 1024;
+auto constexpr MAX_BUFFER_SIZE = size_t(4096);
+
+class PADDLE_API SegmentedDataStream {
+  std::deque<uv_buf_t> _buffers;
+  size_t _buff_idx{0};
+  size_t _buff_offset{0};
+  size_t capacity{0};
+  size_t _buff_offset_commit{0};
+  size_t _read_offset{0};
+
+ public:
+  SegmentedDataStream() = default;
+  void append(uv_buf_t buf);
+  bool readMany(char* dest, size_t size);
+  template <typename T>
+  bool readValue(T& value);  // NOLINT(runtime/references)
+
+  bool readKey(std::string& str);                // NOLINT(runtime/references)
+  bool readContent(std::vector<uint8_t>& data);  // NOLINT(runtime/references)
+  size_t available();
+  void commit();
+  void reset();
+};
+
+class PADDLE_API LibUVHandle
+    : public std::enable_shared_from_this<LibUVHandle> {
+ public:
+  ~LibUVHandle() = default;
+  std::shared_ptr<LibUVHandle> ptr();
+  virtual uv_handle_t* getRawHandle() = 0;
+  void close();
+
+ protected:
+  void handleAvailable();
+  virtual void onClose() = 0;
+
+ private:
+  static void handleClose(uv_handle_t* uv_handle);
+};
+
+class PADDLE_API LibUVTCPSocket : public LibUVHandle {
+ public:
+  explicit LibUVTCPSocket(uv_loop_t* loop);
+  uv_handle_t* getRawHandle() override;
+  std::shared_ptr<LibUVTCPSocket> ptr();
+  static std::shared_ptr<LibUVTCPSocket> getTCPSocket(uv_stream_t* handle);
+  virtual void doProcess(const uv_buf_t* buf, size_t nread) {
+    PADDLE_THROW(
+        common::errors::Fatal("Socket subclass does not implement doProcess"));
+  }
+  uv_tcp_t client{};
+
+ protected:
+  void onClose() override {}
+};
+
+class PADDLE_API LibUVTCPServer : public LibUVTCPSocket {
+ public:
+  typedef std::function<void(int)> LibUVCallback;
+  explicit LibUVTCPServer(uv_loop_t* loop)
+      : LibUVTCPSocket(loop), _on_connect_callback(defaultOnConnect) {}
+  void setCallback(LibUVCallback&& callback);
+  static std::shared_ptr<LibUVTCPServer> createServer(uv_loop_t* loop,
+                                                      std::uint16_t port,
+                                                      bool useIpv6);
+  std::uint16_t port() const { return _port; }
+  void accept(const std::shared_ptr<LibUVTCPSocket>& socket);
+
+ protected:
+  uv_tcp_t* getRawSocket() { return &client; }
+  uv_stream_t* getRawStream() {
+    return reinterpret_cast<uv_stream_t*>(&client);
+  }
+
+ private:
+  LibUVCallback _on_connect_callback;
+  std::uint16_t _port{};
+
+  void setSocketPort();
+  static void defaultOnConnect(int status) {
+    PADDLE_THROW(common::errors::Fatal(
+        "Socket accepted, but onConnect callback is undefined"));
+  }
+  static void onNewConnection(uv_stream_t* server, int status);
+};
+
+class PADDLE_API LibUVMasterDaemon : public DaemonThread {
+ public:
+  explicit LibUVMasterDaemon(int port);
+  // Disable copy constructor
+  LibUVMasterDaemon(const LibUVMasterDaemon& other) = delete;
+  // Disable move constructor
+  LibUVMasterDaemon(LibUVMasterDaemon&& other) = delete;
+  // Disable copy assignment operator
+  LibUVMasterDaemon& operator=(const LibUVMasterDaemon& other) = delete;
+  // Disable move assignment operator
+  LibUVMasterDaemon& operator=(LibUVMasterDaemon&& other) = delete;
+  ~LibUVMasterDaemon() override;
+  void init(const std::uint16_t& port);
+  // operator for key
+  void set(const std::string& key, const std::vector<uint8_t>& value);
+  const std::vector<uint8_t>& get(const std::string& key);
+  int64_t add(const std::string& key, int64_t addVal);
+  bool waitKey(const std::string& key,
+               const std::shared_ptr<LibUVHandle>& client);
+  bool checkKeys(const std::vector<std::string>& keys);
+  // client
+  void addClient(const std::shared_ptr<LibUVHandle>& client);
+  void removeClient(const std::shared_ptr<LibUVHandle>& client);
+  void clearWaitState(const std::shared_ptr<LibUVHandle>& client);
+
+ protected:
+  void run() override;
+  void stop() override;
+
+ private:
+  uv_loop_t loop_{};
+  uv_async_t _exit_handle{};
+  // tcp server
+  std::shared_ptr<LibUVTCPServer> _tcp_server;
+  // tcp store
+  std::unordered_map<std::string, std::vector<uint8_t>> _tcp_store;
+  // the list of LibUVClient waiting on the key
+  std::unordered_map<std::string, std::vector<std::shared_ptr<LibUVHandle>>>
+      _waiting_sockets;
+  // number of keys awaited
+  std::unordered_map<std::shared_ptr<LibUVHandle>, size_t> _awaited_keys;
+  std::unordered_set<std::shared_ptr<LibUVHandle>> _clients;
+  int port_;
+
+  static LibUVMasterDaemon& UVMasterDaemon(uv_handle_t* stream) {
+    return *reinterpret_cast<LibUVMasterDaemon*>(uv_handle_get_data(stream));
+  }
+  static void on_new_connection(uv_stream_t* server, int status) {
+    UVMasterDaemon(reinterpret_cast<uv_handle_t*>(server)).onConnect(status);
+  }
+  static void on_exit_request(uv_async_t* handle) {
+    UVMasterDaemon(reinterpret_cast<uv_handle_t*>(handle)).onExitRequest();
+  }
+  void onConnect(int status);
+  void onExitRequest();
+  void notifyWaitingClients(const std::string& key);
+};
+
+class PADDLE_API WriteUVContent
+    : public std::enable_shared_from_this<WriteUVContent> {
+  std::shared_ptr<WriteUVContent> ptr() { return shared_from_this(); }
+  static void writeDone(uv_write_t* req, int status);
+  struct RequestData {
+    std::shared_ptr<WriteUVContent> strong_self;
+  };
+  std::vector<uint8_t> data;
+  uv_write_t req = {};
+  uv_buf_t buf = {};
+  std::shared_ptr<LibUVHandle> handle;
+
+ public:
+  WriteUVContent(std::vector<uint8_t>&& in_data,
+                 std::shared_ptr<LibUVHandle> handle);
+  ~WriteUVContent();
+  void send();
+};
+
+class PADDLE_API UVWriter {
+  std::vector<uint8_t> data;
+  std::shared_ptr<LibUVHandle> handle;
+  void* operator new(size_t);
+
+ public:
+  explicit UVWriter(std::shared_ptr<LibUVHandle> handle)
+      : handle(std::move(handle)) {}
+  template <typename T>
+  void writeValue(T val);
+  void writeVector(const std::vector<uint8_t>& val);
+  void writeString(const std::string& val);
+  void send();
+};
+
+class PADDLE_API LibUVClient : public LibUVTCPSocket {
+  SegmentedDataStream stream;
+  LibUVMasterDaemon* store;
+  std::string _address{"null"};
+  const std::string& address() const { return _address; }
+  static void allocBuffer(uv_handle_t* handle, size_t buf_size, uv_buf_t* buf);
+  static void readCallback(uv_stream_t* client,
+                           ssize_t nread,
+                           const uv_buf_t* buf);
+
+ protected:
+  void doProcess(const uv_buf_t* buf, size_t nread) override;
+  bool doSetCommand();
+  bool doGetCommand();
+  bool doAddCommand();
+  bool doCheckCommand();
+  bool doWaitCommand();
+  void onClose() override;
+
+ public:
+  explicit LibUVClient(uv_loop_t* loop, LibUVMasterDaemon* store)
+      : LibUVTCPSocket(loop), store(store) {}
+  void readStart();
+  static std::shared_ptr<LibUVClient> make(uv_loop_t* loop,
+                                           LibUVMasterDaemon* store);
+  std::shared_ptr<LibUVClient> ptr();
+};
+}  // namespace phi::distributed::detail