chore: import upstream snapshot with attribution

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