chore: import upstream snapshot with attribution

paddle/phi/core/framework/CMakeLists.txt

@@ -0,0 +1,12 @@
+proto_library(framework_proto SRCS framework.proto)
+proto_library(pass_desc_proto SRCS pass_desc.proto DEPS framework_proto)
+proto_library(data_feed_proto SRCS data_feed.proto)
+proto_library(trainer_desc_proto SRCS trainer_desc.proto DEPS framework_proto
+              data_feed_proto)
+proto_library(heter_service_proto SRCS heter_service.proto)
+
+file(
+  GLOB framework_cc
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "*.cc")
+collect_srcs(core_srcs SRCS ${framework_cc})

paddle/phi/core/framework/convert_utils.h

@@ -0,0 +1,28 @@
+// 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/phi/core/framework/var_type_helper.h"
+#include "paddle/phi/core/utils/data_type.h"
+
+namespace phi {
+
+inline paddle::framework::proto::VarType::Type TransToProtoVarTypeReturnType(
+    const DataType& dtype) {
+  return static_cast<paddle::framework::proto::VarType::Type>(
+      phi::TransToProtoVarType(dtype));
+}
+
+}  // namespace phi

paddle/phi/core/framework/data_feed.proto

@@ -0,0 +1,67 @@
+// 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.
+
+syntax = "proto2";
+package paddle.framework;
+
+message Slot {
+  required string name = 1;
+  required string type = 2;
+  optional bool is_dense = 3 [ default = false ];
+  optional bool is_used = 4 [ default = false ];
+  repeated int32 shape = 5; // we can define N-D phi::DenseTensor
+}
+
+message MultiSlotDesc {
+  repeated Slot slots = 1;
+  optional string uid_slot = 2;
+}
+
+message GraphConfig {
+  optional int32 walk_degree = 1 [ default = 1 ];
+  optional int32 walk_len = 2 [ default = 20 ];
+  optional int32 window = 3 [ default = 5 ];
+  optional int32 once_sample_startid_len = 4 [ default = 8000 ];
+  optional int32 sample_times_one_chunk = 5 [ default = 10 ];
+  optional int32 batch_size = 6 [ default = 1 ];
+  optional int32 debug_mode = 7 [ default = 0 ];
+  optional string first_node_type = 8;
+  optional string meta_path = 9;
+  optional bool gpu_graph_training = 10 [ default = true ];
+  optional bool sage_mode = 11 [ default = false ];
+  optional string samples = 12;
+  optional int64 train_table_cap = 13 [ default = 80000 ];
+  optional int64 infer_table_cap = 14 [ default = 80000 ];
+  optional string excluded_train_pair = 15;
+  optional string infer_node_type = 16;
+  optional bool get_degree = 17 [ default = false ];
+  optional bool weighted_sample = 18 [ default = false ];
+  optional bool return_weight = 19 [ default = false ];
+  optional string pair_label = 20;
+  optional bool is_thread_sharding = 21 [ default = false ];
+  optional int32 accumulate_num = 22 [ default = 1 ];
+}
+
+message DataFeedDesc {
+  optional string name = 1;
+  optional int32 batch_size = 2 [ default = 32 ];
+  optional MultiSlotDesc multi_slot_desc = 3;
+  optional string pipe_command = 4;
+  optional int32 thread_num = 5;
+  optional string rank_offset = 6;
+  optional int32 pv_batch_size = 7 [ default = 32 ];
+  optional int32 input_type = 8 [ default = 0 ];
+  optional string so_parser_name = 9;
+  optional GraphConfig graph_config = 10;
+}

paddle/phi/core/framework/data_type_transform.cc

@@ -0,0 +1,253 @@
+// 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/framework/data_type_transform.h"
+
+#include "paddle/phi/common/transform.h"
+#include "paddle/phi/core/framework/convert_utils.h"
+#include "paddle/phi/core/framework/selected_rows_serialize.h"
+#include "paddle/phi/core/kernel_factory.h"
+
+#if defined(PADDLE_WITH_XPU)
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+#endif
+
+namespace proto = paddle::framework::proto;
+
+namespace phi {
+
+template <typename InType, typename OutType>
+struct CastDataTypeFunctor {
+  HOSTDEVICE inline OutType operator()(InType in) const {
+#if defined(_MSC_VER)
+    // Avoid unsupported convert of float8/bfloat16/float16 -> complex
+    if constexpr (
+        (std::is_same_v<OutType, phi::complex64> ||
+         std::is_same_v<
+             OutType,
+             phi::complex128>)&&(std::is_same_v<InType, phi::float8_e4m3fn> ||
+                                 std::is_same_v<InType, phi::float8_e5m2> ||
+                                 std::is_same_v<InType, phi::bfloat16> ||
+                                 std::is_same_v<InType, phi::float16>)) {
+      // default value，only to avoid compile error
+      return OutType(0);
+    } else {
+      return static_cast<OutType>(in);
+    }
+#else
+    return static_cast<OutType>(in);
+#endif
+  }
+};
+
+#if defined(PADDLE_WITH_XPU)
+
+template <typename InType, typename OutType>
+static void XPUCastData(const DenseTensor& in,
+                        DenseTensor* out,
+                        const phi::XPUContext* dev_ctx) {
+  using XPUInTDType = typename XPUTypeTrait<InType>::Type;
+  using XPUOutTDType = typename XPUTypeTrait<OutType>::Type;
+  int r = xpu::cast<XPUInTDType, XPUOutTDType>(
+      dev_ctx->x_context(),
+      reinterpret_cast<const XPUInTDType*>(in.data<InType>()),
+      reinterpret_cast<XPUOutTDType*>(dev_ctx->Alloc<OutType>(out)),
+      in.numel());
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+  dev_ctx->Wait();
+}
+
+template <typename InType>
+static void XPUTransDataType(
+    const DenseTensor& in,
+    DenseTensor* out,
+    const paddle::framework::proto::VarType::Type& dst_type,
+    const phi::DeviceContext* ctx) {
+  auto* context = static_cast<const phi::XPUContext*>(ctx);
+
+#define XPUCastCallback(cpp_type, proto_type)          \
+  do {                                                 \
+    if (dst_type == proto_type) {                      \
+      XPUCastData<InType, cpp_type>(in, out, context); \
+    }                                                  \
+  } while (0)
+
+  if (dst_type == proto::VarType::FP32 || dst_type == proto::VarType::FP16 ||
+      dst_type == proto::VarType::BOOL || dst_type == proto::VarType::INT16 ||
+      dst_type == proto::VarType::INT32 || dst_type == proto::VarType::INT64 ||
+      dst_type == proto::VarType::FP64) {
+    _ForEachDataTypeForXPU_(XPUCastCallback);
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Data type (%s) is not supported in XPU when casting data type.",
+        VarDataTypeToString(dst_type)));
+  }
+}
+
+#endif
+
+template <typename InType>
+struct CastDataType {
+  CastDataType(const DenseTensor& in,
+               DenseTensor* out,
+               const phi::DeviceContext* ctx)
+      : in_(in), out_(out), ctx_(ctx) {}
+  const DenseTensor in_;
+  DenseTensor* out_;
+  const phi::DeviceContext* ctx_;
+
+  template <typename OutType>
+  void apply() {
+    auto* in_begin = in_.data<InType>();
+    auto* in_end = in_begin + in_.numel();
+    auto* out_begin = ctx_->Alloc<OutType>(out_);
+
+    if (phi::is_cpu_place(in_.place())) {
+      phi::Transform<phi::CPUContext> trans;
+      auto* context = static_cast<const phi::CPUContext*>(ctx_);
+      trans(*context,
+            in_begin,
+            in_end,
+            out_begin,
+            CastDataTypeFunctor<InType, OutType>());
+#if defined(__NVCC__) || defined(__HIPCC__)
+    } else if (phi::is_gpu_place(in_.place())) {
+      phi::Transform<phi::GPUContext> trans;
+      auto* context = static_cast<const phi::GPUContext*>(ctx_);
+      trans(*context,
+            in_begin,
+            in_end,
+            out_begin,
+            CastDataTypeFunctor<InType, OutType>());
+      context->Wait();
+#endif
+#if defined(PADDLE_WITH_IPU)
+    } else if (phi::is_ipu_place(in_.place())) {
+      phi::Transform<phi::CPUContext> trans;
+      auto* context = static_cast<const phi::CPUContext*>(ctx_);
+      trans(*context,
+            in_begin,
+            in_end,
+            out_begin,
+            CastDataTypeFunctor<InType, OutType>());
+#endif
+    } else {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Place type is not supported when casting data type."));
+    }
+  }
+};
+
+void TransDataType(const phi::KernelKey& kernel_type_for_var,
+                   const phi::KernelKey& expected_kernel_type,
+                   const DenseTensor& in,
+                   DenseTensor* out) {
+  PADDLE_ENFORCE_EQ(in.dtype(),
+                    kernel_type_for_var.dtype(),
+                    common::errors::InvalidArgument(
+                        "The src dtype(%s) of input tensor and kernel_type(%s) "
+                        "are not consistent.",
+                        DataTypeToString(in.dtype()),
+                        DataTypeToString(kernel_type_for_var.dtype())));
+  auto dst_type =
+      phi::TransToProtoVarTypeReturnType(expected_kernel_type.dtype());
+  TransDataType(in, dst_type, out);
+}
+
+void TransDataType(const DenseTensor& in,
+                   const paddle::framework::proto::VarType::Type& type,
+                   DenseTensor* out) {
+  phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
+
+  out->Resize(in.dims());
+  auto src_type = phi::TransToProtoVarTypeReturnType(in.dtype());
+  auto dst_type = type;
+  auto ctx = pool.Get(in.place());
+
+#if defined(PADDLE_WITH_XPU)
+  if (phi::is_xpu_place(in.place())) {
+    switch (src_type) {
+      case proto::VarType::FP16:
+        XPUTransDataType<phi::float16>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::FP32:
+        XPUTransDataType<float>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::FP64:
+        XPUTransDataType<double>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::BOOL:
+        XPUTransDataType<bool>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::INT16:
+        XPUTransDataType<int16_t>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::INT32:
+        XPUTransDataType<int>(in, out, dst_type, ctx);
+        break;
+      case proto::VarType::INT64:
+        XPUTransDataType<int64_t>(in, out, dst_type, ctx);
+        break;
+      default:
+        PADDLE_THROW(common::errors::Unimplemented(
+            "Data type (%s) is not supported in XPU when casting data type.",
+            VarDataTypeToString(src_type)));
+    }
+    return;
+  }
+#endif
+  switch (src_type) {
+    case proto::VarType::FP16:
+      phi::VisitDataType(dst_type, CastDataType<phi::float16>(in, out, ctx));
+      break;
+    case proto::VarType::BF16:
+      phi::VisitDataType(dst_type, CastDataType<phi::bfloat16>(in, out, ctx));
+      break;
+    case proto::VarType::FP8_E4M3FN:
+      phi::VisitDataType(dst_type,
+                         CastDataType<phi::float8_e4m3fn>(in, out, ctx));
+      break;
+    case proto::VarType::FP8_E5M2:
+      phi::VisitDataType(dst_type,
+                         CastDataType<phi::float8_e5m2>(in, out, ctx));
+      break;
+    case proto::VarType::FP32:
+      phi::VisitDataType(dst_type, CastDataType<float>(in, out, ctx));
+      break;
+    case proto::VarType::FP64:
+      phi::VisitDataType(dst_type, CastDataType<double>(in, out, ctx));
+      break;
+    case proto::VarType::INT32:
+      phi::VisitDataType(dst_type, CastDataType<int>(in, out, ctx));
+      break;
+    case proto::VarType::INT64:
+      phi::VisitDataType(dst_type, CastDataType<int64_t>(in, out, ctx));
+      break;
+    case proto::VarType::BOOL:
+      phi::VisitDataType(dst_type, CastDataType<bool>(in, out, ctx));
+      break;
+    case proto::VarType::INT16:
+      phi::VisitDataType(dst_type, CastDataType<int16_t>(in, out, ctx));
+      break;
+    case proto::VarType::UINT8:
+      phi::VisitDataType(dst_type, CastDataType<uint8_t>(in, out, ctx));
+      break;
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Data type (%s) is not supported when casting data type.",
+          VarDataTypeToString(src_type)));
+  }
+}
+
+}  // namespace phi

paddle/phi/core/framework/data_type_transform.h

@@ -0,0 +1,32 @@
+// 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 <utility>
+
+#include "paddle/phi/core/framework/var_type_helper.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+void TransDataType(const KernelKey& kernel_type_for_var,
+                   const KernelKey& expected_kernel_type,
+                   const DenseTensor& in,
+                   DenseTensor* out);
+
+void TransDataType(const DenseTensor& in,
+                   const paddle::framework::proto::VarType::Type& type,
+                   DenseTensor* out);
+}  // namespace phi

paddle/phi/core/framework/dense_tensor_serialize.cc

@@ -0,0 +1,130 @@
+// 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/framework/dense_tensor_serialize.h"
+#include <cstdint>
+#include "paddle/phi/core/framework/convert_utils.h"
+
+namespace phi {
+
+void SerializeToStream(std::ostream &os,
+                       const DenseTensor &tensor,
+                       const DeviceContext &dev_ctx) {
+  constexpr uint32_t kCurTensorVersion = 0;
+  {  // the 1st field, uint32_t version for DenseTensor
+    os.write(reinterpret_cast<const char *>(&kCurTensorVersion),
+             sizeof(kCurTensorVersion));
+  }
+  {
+    // the 2nd field, LoD information
+    // uint64_t lod_level
+    // uint64_t lod_level_1 size in byte.
+    // int*     lod_level_1 data
+    // ...
+    auto lod = tensor.lod();
+    uint64_t size = lod.size();
+    os.write(reinterpret_cast<const char *>(&size), sizeof(size));
+
+    for (auto &each : lod) {
+      size = each.size() * sizeof(LegacyLoD::value_type::value_type);
+      os.write(reinterpret_cast<const char *>(&size), sizeof(size));
+      os.write(reinterpret_cast<const char *>(each.data()),
+               static_cast<std::streamsize>(size));
+    }
+  }
+  // the 3rd field, Tensor
+  TensorToStream(os, static_cast<DenseTensor>(tensor), dev_ctx);
+}
+
+void SerializeToStream(std::ostream &os, const DenseTensor &tensor) {
+  DeviceContextPool &pool = DeviceContextPool::Instance();
+  const DeviceContext *dev_ctx = nullptr;
+  auto place = tensor.place();
+  dev_ctx = pool.Get(place);
+  SerializeToStream(os, tensor, *dev_ctx);
+}
+
+void DeserializeFromStream(std::istream &os, DenseTensor *tensor) {
+  DeviceContextPool &pool = DeviceContextPool::Instance();
+  const DeviceContext *dev_ctx = nullptr;
+  dev_ctx = pool.Get(CPUPlace());
+  DeserializeFromStream(os, tensor, *dev_ctx);
+}
+
+void DeserializeFromStream(std::istream &is,
+                           DenseTensor *tensor,
+                           const DeviceContext &dev_ctx,
+                           const size_t &seek,
+                           const std::vector<int64_t> &shape) {
+  {
+    // the 1st field, unit32_t version for DenseTensor
+    uint32_t version = 0;
+    is.read(reinterpret_cast<char *>(&version), sizeof(version));
+
+    PADDLE_ENFORCE_EQ(
+        version,
+        0U,
+        common::errors::InvalidArgument(
+            "Deserialize to tensor failed, maybe the loaded file is "
+            "not a paddle model(expected file format: 0, but %u found).",
+            version));
+  }
+  {
+    // the 2nd field, LoD information
+    uint64_t lod_level = 0;
+    is.read(reinterpret_cast<char *>(&lod_level), sizeof(lod_level));
+    auto &lod = *tensor->mutable_lod();
+    lod.resize(lod_level);
+  }
+  // the 3rd field, Tensor
+  TensorFromStream(
+      is, static_cast<DenseTensor *>(tensor), dev_ctx, seek, shape);
+}
+
+void DeserializeFromStream(std::istream &is,
+                           DenseTensor *tensor,
+                           const DeviceContext &dev_ctx) {
+  {
+    // the 1st field, unit32_t version for DenseTensor
+    uint32_t version = 0;
+    is.read(reinterpret_cast<char *>(&version), sizeof(version));
+
+    PADDLE_ENFORCE_EQ(
+        version,
+        0U,
+        common::errors::InvalidArgument(
+            "Deserialize to tensor failed, maybe the loaded file is "
+            "not a paddle model(expected file format: 0, but %u found).",
+            version));
+  }
+  {
+    // the 2nd field, LoD information
+    uint64_t lod_level = 0;
+    is.read(reinterpret_cast<char *>(&lod_level), sizeof(lod_level));
+    auto &lod = *tensor->mutable_lod();
+    lod.resize(lod_level);
+    for (uint64_t i = 0; i < lod_level; ++i) {
+      uint64_t size = 0;
+      is.read(reinterpret_cast<char *>(&size), sizeof(size));
+      std::vector<size_t> tmp(size / sizeof(size_t));
+      is.read(reinterpret_cast<char *>(tmp.data()),
+              static_cast<std::streamsize>(size));
+      lod[i] = tmp;
+    }
+  }
+  // the 3rd field, Tensor
+  TensorFromStream(is, static_cast<DenseTensor *>(tensor), dev_ctx);
+}
+
+}  // namespace phi

paddle/phi/core/framework/dense_tensor_serialize.h

@@ -0,0 +1,50 @@
+// 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 <memory>
+#include <utility>
+#include "paddle/common/ddim.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/framework/dense_tensor_tostream.h"
+#include "paddle/phi/core/mixed_vector.h"
+#include "paddle/utils/test_macros.h"
+
+namespace phi {
+
+/*
+ * Serialize/Deserialize DenseTensor to std::ostream
+ * You can pass ofstream or ostringstream to serialize to file
+ * or to a in memory string. GPU tensor will be copied to CPU.
+ */
+PADDLE_API void SerializeToStream(std::ostream& os,
+                                  const DenseTensor& tensor,
+                                  const DeviceContext& dev_ctx);
+PADDLE_API void DeserializeFromStream(std::istream& is,
+                                      DenseTensor* tensor,
+                                      const DeviceContext& dev_ctx);
+PADDLE_API void DeserializeFromStream(std::istream& is,
+                                      DenseTensor* tensor,
+                                      const DeviceContext& dev_ctx,
+                                      const size_t& seek,
+                                      const std::vector<int64_t>& shape);
+
+PADDLE_API void SerializeToStream(std::ostream& os, const DenseTensor& tensor);
+
+PADDLE_API void DeserializeFromStream(std::istream& os, DenseTensor* tensor);
+
+}  // namespace phi

paddle/phi/core/framework/dense_tensor_tostream.cc

@@ -0,0 +1,370 @@
+// 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/framework/dense_tensor_tostream.h"
+
+#include <algorithm>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/compat/convert_utils.h"
+#include "paddle/phi/core/framework/convert_utils.h"
+#include "paddle/phi/core/kernel_factory.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/contiguous_kernel.h"
+
+namespace phi {
+
+namespace proto = paddle::framework::proto;
+
+template <typename Context>
+DenseTensor InnerTensorContiguous(const Context& dev_ctx,
+                                  const DenseTensor& tensor) {
+  DenseTensor dense_out;
+  MetaTensor meta_input(tensor);
+  MetaTensor meta_out(&dense_out);
+  UnchangedInferMeta(meta_input, &meta_out);
+
+  PD_VISIT_ALL_TYPES(tensor.dtype(), "InnerTensorContiguous", ([&] {
+                       ContiguousKernel<data_t, Context>(
+                           dev_ctx, tensor, &dense_out);
+                     }));
+  return dense_out;
+}
+
+DenseTensor InnerTensorContiguous(const DenseTensor& tensor) {
+  auto& pool = DeviceContextPool::Instance();
+
+  if (tensor.place().GetType() == AllocationType::CPU) {
+    auto* dev_ctx = static_cast<CPUContext*>(pool.Get(tensor.place()));
+    return InnerTensorContiguous<CPUContext>(*dev_ctx, tensor);
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  } else if (tensor.place().GetType() == AllocationType::GPU) {
+    auto* dev_ctx = static_cast<GPUContext*>(pool.Get(tensor.place()));
+    return InnerTensorContiguous<GPUContext>(*dev_ctx, tensor);
+#endif
+#ifdef PADDLE_WITH_XPU
+  } else if (tensor.place().GetType() == AllocationType::XPU) {
+    auto* dev_ctx = static_cast<XPUContext*>(pool.Get(tensor.place()));
+    return InnerTensorContiguous<XPUContext>(*dev_ctx, tensor);
+#endif
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  } else if (tensor.place().GetType() == AllocationType::CUSTOM) {
+    auto* dev_ctx = static_cast<CustomContext*>(pool.Get(tensor.place()));
+    DenseTensor dense_out;
+    MetaTensor meta_input(tensor);
+    MetaTensor meta_out(&dense_out);
+    UnchangedInferMeta(meta_input, &meta_out);
+    const KernelKey& kernel_key = {TransToPhiBackend(tensor.place()),
+                                   DataLayout::ALL_LAYOUT,
+                                   tensor.dtype()};
+    using kernel_signature =
+        void (*)(const DeviceContext&, const DenseTensor&, DenseTensor*);
+    PD_VISIT_KERNEL("contiguous",
+                    kernel_key,
+                    kernel_signature,
+                    false,
+                    *dev_ctx,
+                    tensor,
+                    &dense_out);
+    return dense_out;
+#endif
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Place type is not supported when casting data type."));
+  }
+
+  return tensor;
+}
+
+void TensorToStream(std::ostream& os,
+                    const DenseTensor& tensor,
+                    const DeviceContext& dev_ctx) {
+  const auto ensure_contiguous = [](const DenseTensor& tensor) {
+    if (tensor.meta().is_contiguous()) {
+      return tensor;
+    }
+    return InnerTensorContiguous(tensor);
+  };
+  const DenseTensor& contiguous_tensor = ensure_contiguous(tensor);
+  {  // the 1st field, uint32_t version
+    constexpr uint32_t version = 0;
+    os.write(reinterpret_cast<const char*>(&version), sizeof(version));
+  }
+  {  // the 2nd field, tensor description
+     // int32_t  size
+     // void*    protobuf message
+    proto::VarType::TensorDesc desc;
+    desc.set_data_type(
+        TransToProtoVarTypeReturnType(contiguous_tensor.dtype()));
+    auto dims = common::vectorize(contiguous_tensor.dims());
+    auto* pb_dims = desc.mutable_dims();
+    pb_dims->Resize(static_cast<int>(dims.size()), 0);
+    std::copy(dims.begin(), dims.end(), pb_dims->begin());
+    int32_t size = static_cast<int32_t>(desc.ByteSizeLong());
+    os.write(reinterpret_cast<const char*>(&size), sizeof(size));
+    auto out = desc.SerializeAsString();
+    os.write(out.data(), size);
+  }
+  {  // the 3rd field, tensor data
+    uint64_t size =
+        contiguous_tensor.numel() * SizeOf(contiguous_tensor.dtype());
+
+    auto* data_ptr = contiguous_tensor.data();
+    PADDLE_ENFORCE_LT(size,
+                      (std::numeric_limits<std::streamsize>::max)(),
+                      common::errors::ResourceExhausted(
+                          "tensor size %d overflow when writing tensor", size));
+    if (is_gpu_place(contiguous_tensor.place())) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
+      std::unique_ptr<char[]> buf(new char[kBufSize]);
+      auto& gpu_dev_ctx = static_cast<const GPUContext&>(dev_ctx);
+      CPUPlace cpu;
+      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
+      while (size != 0) {
+        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
+        memory_utils::Copy(cpu,
+                           buf.get(),
+                           contiguous_tensor.place(),
+                           reinterpret_cast<const void*>(data),  // NOLINT
+                           size_to_write,
+                           gpu_dev_ctx.stream());
+        gpu_dev_ctx.Wait();
+        os.write(buf.get(), size_to_write);
+        data += size_to_write;
+        size -= size_to_write;
+      }
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "CUDAPlace is not supported when not compiled with CUDA"));
+#endif
+    } else if (is_xpu_place(contiguous_tensor.place())) {
+#ifdef PADDLE_WITH_XPU
+      constexpr size_t kBufSize = 1024 * 1024 * 64;  // 64MB
+      std::unique_ptr<char[]> buf(new char[kBufSize]);
+      auto& xpu_dev_ctx = static_cast<const XPUContext&>(dev_ctx);
+      CPUPlace cpu;
+      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
+      while (size != 0) {
+        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
+        memory_utils::Copy(cpu,
+                           buf.get(),
+                           contiguous_tensor.place(),
+                           reinterpret_cast<const void*>(data),
+                           size_to_write);
+        xpu_dev_ctx.Wait();
+        os.write(buf.get(), size_to_write);
+        data += size_to_write;
+        size -= size_to_write;
+      }
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "XPUPlace is not supported when not compiled with XPU"));
+#endif
+    } else if (is_custom_place(contiguous_tensor.place())) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+      constexpr size_t kBufSize = 1024 * 1024 * 64;     // 64MB
+      std::unique_ptr<char[]> buf(new char[kBufSize]);  // NOLINT
+      auto& custom_device_context = static_cast<const CustomContext&>(dev_ctx);
+      CPUPlace cpu;
+      uintptr_t data = reinterpret_cast<uintptr_t>(data_ptr);
+      while (size != 0) {
+        size_t size_to_write = std::min(kBufSize, static_cast<size_t>(size));
+        memory_utils::Copy(cpu,
+                           buf.get(),
+                           contiguous_tensor.place(),
+                           reinterpret_cast<const void*>(data),
+                           size_to_write,
+                           custom_device_context.stream());
+        custom_device_context.Wait();
+        os.write(buf.get(), size_to_write);
+        data += size_to_write;
+        size -= size_to_write;
+      }
+#else
+      PADDLE_THROW(common::errors::Unimplemented(
+          "CustomPlace is not supported when not compiled with "
+          "CustomDevice"));
+#endif
+    } else {
+      os.write(static_cast<const char*>(data_ptr),
+               static_cast<std::streamsize>(size));
+    }
+  }
+}
+
+struct DeserializedDataFunctor {
+  DeserializedDataFunctor(void** buf, DenseTensor* tensor, const Place& place)
+      : buf_(buf), tensor_(tensor), place_(place) {}
+
+  template <typename T>
+  void apply() {
+    auto& pool = DeviceContextPool::Instance();
+    auto* dev_ctx = pool.Get(place_);
+    *buf_ = dev_ctx->Alloc<T>(tensor_);
+  }
+
+  void** buf_;
+  DenseTensor* tensor_;
+  Place place_;
+};
+
+void TensorFromStream(std::istream& is,
+                      DenseTensor* tensor,
+                      const DeviceContext& dev_ctx,
+                      const size_t& seek,
+                      const std::vector<int64_t>& shape) {
+  uint32_t version = 0;
+  is.read(reinterpret_cast<char*>(&version), sizeof(version));
+
+  PADDLE_ENFORCE_EQ(
+      version,
+      0U,
+      common::errors::InvalidArgument(
+          "tensor version %u is not supported, Only version 0 is supported",
+          version));
+
+  proto::VarType::TensorDesc desc;
+  {  // int32_t size
+    // proto buffer
+    int32_t size = 0;
+    is.read(reinterpret_cast<char*>(&size), sizeof(size));
+    std::unique_ptr<char[]> buf(new char[size]);  // NOLINT
+    is.read(reinterpret_cast<char*>(buf.get()), size);
+    PADDLE_ENFORCE_EQ(
+        desc.ParseFromArray(buf.get(), size),
+        true,
+        common::errors::InvalidArgument("Cannot parse tensor desc"));
+  }
+  {  // read tensor
+    tensor->Resize(common::make_ddim(shape));
+    size_t seekg = seek * SizeOfType(desc.data_type());
+    is.seekg(seekg, is.cur);  // NOLINT
+
+    void* buf = nullptr;
+    CPUContext ctx;
+    size_t size = tensor->numel() * SizeOfType(desc.data_type());
+    if (is_gpu_place(dev_ctx.GetPlace()) || is_xpu_place(dev_ctx.GetPlace()) ||
+        is_custom_place(dev_ctx.GetPlace())) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_XPU) || defined(PADDLE_WITH_CUSTOM_DEVICE)
+      DenseTensor cpu_tensor;
+      cpu_tensor.Resize(common::make_ddim(shape));
+      VisitDataType(desc.data_type(),
+                    DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), size);  // NOLINT
+      auto dst_place = dev_ctx.GetPlace();
+      Copy(dev_ctx, cpu_tensor, dst_place, false, tensor);
+      if (is_custom_place(dev_ctx.GetPlace())) {
+        dev_ctx.Wait();
+      }
+#else
+      if (is_gpu_place(dev_ctx.GetPlace())) {
+        PADDLE_THROW(common::errors::Unimplemented(
+            "CUDAPlace is not supported when not compiled with CUDA"));
+      } else if (is_xpu_place(dev_ctx.GetPlace())) {
+        PADDLE_THROW(common::errors::Unimplemented(
+            "XPUPlace is not supported when not compiled with XPU"));
+      }
+#endif
+    } else {
+      VisitDataType(desc.data_type(),
+                    DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), size);  // NOLINT
+    }
+  }
+}
+
+void TensorFromStream(std::istream& is,
+                      DenseTensor* tensor,
+                      const DeviceContext& dev_ctx) {
+  uint32_t version = 0;
+  is.read(reinterpret_cast<char*>(&version), sizeof(version));
+  PADDLE_ENFORCE_EQ(
+      version,
+      0U,
+      common::errors::InvalidArgument(
+          "tensor version %u is not supported, Only version 0 is supported",
+          version));
+  proto::VarType::TensorDesc desc;
+  {  // int32_t size
+     // proto buffer
+    int32_t size = -1;
+    is.read(reinterpret_cast<char*>(&size), sizeof(size));
+    PADDLE_ENFORCE_EQ(
+        is.good(),
+        true,
+        common::errors::Unavailable("Cannot read tensor desc size"));
+    PADDLE_ENFORCE_GE(
+        size,
+        0,
+        common::errors::InvalidArgument("DenseTensor desc size should >= 0"));
+    std::unique_ptr<char[]> buf(new char[size]);  // NOLINT
+    is.read(reinterpret_cast<char*>(buf.get()), size);
+    PADDLE_ENFORCE_EQ(
+        desc.ParseFromArray(buf.get(), size),
+        true,
+        common::errors::InvalidArgument("Cannot parse tensor desc"));
+  }
+  {  // read tensor
+    std::vector<int64_t> dims;
+    dims.reserve(static_cast<size_t>(desc.dims().size()));
+    std::copy(desc.dims().begin(), desc.dims().end(), std::back_inserter(dims));
+    tensor->Resize(common::make_ddim(dims));
+    void* buf = nullptr;
+    CPUContext ctx;
+    size_t size = tensor->numel() * SizeOfType(desc.data_type());
+    if (is_gpu_place(dev_ctx.GetPlace()) || is_xpu_place(dev_ctx.GetPlace()) ||
+        is_custom_place(dev_ctx.GetPlace())) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_XPU) || defined(PADDLE_WITH_CUSTOM_DEVICE)
+      DenseTensor cpu_tensor;
+      cpu_tensor.Resize(common::make_ddim(dims));
+      VisitDataType(desc.data_type(),
+                    DeserializedDataFunctor(&buf, &cpu_tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), size);  // NOLINT
+      auto dst_place = dev_ctx.GetPlace();
+      Copy(dev_ctx, cpu_tensor, dst_place, false, tensor);
+      if (is_custom_place(dev_ctx.GetPlace())) {
+        dev_ctx.Wait();
+      }
+#else
+      if (is_gpu_place(dev_ctx.GetPlace())) {
+        PADDLE_THROW(common::errors::Unimplemented(
+            "CUDAPlace is not supported when not compiled with CUDA"));
+      } else if (is_xpu_place(dev_ctx.GetPlace())) {
+        PADDLE_THROW(common::errors::Unimplemented(
+            "XPUPlace is not supported when not compiled with XPU"));
+      } else {
+        PADDLE_THROW(
+            common::errors::Unimplemented("CustomPlace is not supported when "
+                                          "not compiled with CustomDevice"));
+      }
+#endif
+    } else {
+      VisitDataType(desc.data_type(),
+                    DeserializedDataFunctor(&buf, tensor, ctx.GetPlace()));
+      is.read(static_cast<char*>(buf), size);  // NOLINT
+    }
+  }
+}
+
+}  // namespace phi

paddle/phi/core/framework/dense_tensor_tostream.h

@@ -0,0 +1,42 @@
+// 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 <algorithm>
+#include <codecvt>
+#include <locale>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/framework/var_type_helper.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+namespace phi {
+
+TEST_API void TensorToStream(std::ostream& os,
+                             const DenseTensor& tensor,
+                             const DeviceContext& dev_ctx);
+TEST_API void TensorFromStream(std::istream& is,
+                               DenseTensor* tensor,
+                               const DeviceContext& dev_ctx);
+void TensorFromStream(std::istream& is,
+                      DenseTensor* tensor,
+                      const DeviceContext& dev_ctx,
+                      const size_t& seek,
+                      const std::vector<int64_t>& shape);
+
+}  // namespace phi

paddle/phi/core/framework/feed_fetch_type.h

@@ -0,0 +1,54 @@
+// 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 <vector>
+#include "paddle/phi/core/extended_tensor.h"
+#include "paddle/phi/core/tensor_array.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/vocab/string_array.h"
+
+namespace phi {
+using FeedType = phi::DenseTensor;
+using FetchType = paddle::variant<phi::DenseTensor,
+                                  phi::TensorArray,
+                                  phi::Vocab,
+                                  phi::SparseCooTensor>;
+
+template <>
+struct PhiVectorType<FeedType> {
+  const char *type_name = "PhiVectorFeedType";
+};
+
+template <>
+struct PhiVectorType<FetchType> {
+  const char *type_name = "PhiVectorFetchType";
+};
+
+using FeedList = PhiVector<FeedType>;
+using FetchList = PhiVector<FetchType>;
+using FetchUnmergedList = std::vector<std::vector<FetchType>>;
+
+}  // namespace phi
+
+namespace paddle {
+namespace framework {
+using FeedType = phi::FeedType;
+using FetchType = phi::FetchType;
+using FeedList = phi::FeedList;
+using FetchList = phi::FetchList;
+using FetchUnmergedList = phi::FetchUnmergedList;
+}  // namespace framework
+}  // namespace paddle

paddle/phi/core/framework/framework.proto

@@ -0,0 +1,273 @@
+// 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.
+
+syntax = "proto2";
+package paddle.framework.proto;
+
+// Any incompatible changes to ProgramDesc and its dependencies should
+// raise the version defined version.h.
+//
+// Serialization and Deserialization codes should be modified in a way
+// that supports old versions following the version and compatibility policy.
+message Version { optional int64 version = 1 [ default = 0 ]; }
+
+enum AttrType {
+  INT = 0;
+  FLOAT = 1;
+  STRING = 2;
+  INTS = 3;
+  FLOATS = 4;
+  STRINGS = 5;
+  BOOLEAN = 6;
+  BOOLEANS = 7;
+  BLOCK = 8;
+  LONG = 9;
+  BLOCKS = 10;
+  LONGS = 11;
+  FLOAT64S = 12;
+  VAR = 13;
+  VARS = 14;
+  FLOAT64 = 15;
+  SCALAR = 16;
+  SCALARS = 17;
+}
+
+
+message Complex {
+  required double r = 1;
+  required double i = 2;
+};
+
+message Scalar {
+  enum Type {
+    BOOLEAN = 1;
+    LONG = 2;
+    FLOAT64 = 3;
+    COMPLEX128 = 4;
+  }
+  required Type type = 1;
+
+  optional bool b = 2;
+  optional int64 i = 3;
+  optional double r = 4;
+  optional Complex c = 5;
+};
+
+// OpDesc describes an instance of a C++ framework::OperatorBase
+// derived class type.
+message OpDesc {
+
+  message Attr {
+    required string name = 1;
+    required AttrType type = 2;
+    optional int32 i = 3;
+    optional float f = 4;
+    optional string s = 5;
+    repeated int32 ints = 6;
+    repeated float floats = 7;
+    repeated string strings = 8;
+    optional bool b = 10;
+    repeated bool bools = 11;
+    optional int32 block_idx = 12;
+    optional int64 l = 13;
+    repeated int32 blocks_idx = 14;
+    repeated int64 longs = 15;
+    repeated double float64s = 16;
+    optional string var_name = 17;
+    repeated string vars_name = 18;
+    optional double float64 = 19;
+    optional Scalar scalar = 20;
+    repeated Scalar scalars = 21;
+  };
+
+  message Var {
+    required string parameter = 1;
+    repeated string arguments = 2;
+  };
+
+  required string type = 3;
+  repeated Var inputs = 1;
+  repeated Var outputs = 2;
+  repeated Attr attrs = 4;
+  optional bool is_target = 5 [ default = false ];
+};
+
+// OpProto describes a C++ framework::OperatorBase derived class.
+message OpProto {
+
+  // VarProto describes the C++ type framework::Variable.
+  message Var {
+    required string name = 1;
+    required string comment = 2;
+
+    optional bool duplicable = 3 [ default = false ];
+    optional bool intermediate = 4 [ default = false ];
+    optional bool dispensable = 5 [ default = false ];
+    optional bool extra = 6 [ default = false ];
+    optional bool quant = 7 [ default = false ];
+  }
+
+  // AttrProto describes the C++ type Attribute.
+  message Attr {
+    required string name = 1;
+    required AttrType type = 2;
+    required string comment = 3;
+    // If that attribute is generated, it means the Paddle third
+    // language binding has responsibility to fill that
+    // attribute. End-User should not set that attribute.
+    optional bool generated = 4 [ default = false ];
+    optional bool extra = 5 [ default = false ];
+    optional bool quant = 6 [ default = false ];
+    optional bool support_tensor = 7 [ default = false];
+  }
+
+  required string type = 1;
+  repeated Var inputs = 2;
+  repeated Var outputs = 3;
+  repeated Attr attrs = 4;
+  required string comment = 5;
+}
+
+message VarType {
+  enum Type {
+    // Pod Types
+    BOOL = 0;
+    INT16 = 1;
+    INT32 = 2;
+    INT64 = 3;
+    FP16 = 4;
+    FP32 = 5;
+    FP64 = 6;
+    // phi::DenseTensor<size_t> is used in C++.
+    SIZE_T = 19;
+    UINT8 = 20;
+    INT8 = 21;
+    BF16 = 22;
+    COMPLEX64 = 23;
+    COMPLEX128 = 24;
+    FP8_E4M3FN = 32;
+    FP8_E5M2 = 33;
+    UINT16 = 36;
+    UINT32 = 37;
+    UINT64 = 38;
+    // Other types that may need additional descriptions
+    DENSE_TENSOR = 7;
+    SELECTED_ROWS = 8;
+    FEED_MINIBATCH = 9;
+    FETCH_LIST = 10;
+    STEP_SCOPES = 11;
+    LOD_RANK_TABLE = 12;
+    DENSE_TENSOR_ARRAY = 13;
+    PLACE_LIST = 14;
+    READER = 15;
+    // Any runtime decided variable type is raw
+    // raw variables should manage their own allocations
+    // in operators like nccl_op
+    RAW = 17;
+    TUPLE = 18;
+
+    STRING = 25;
+    STRINGS = 26;
+    VOCAB = 27;
+    FEED_LIST = 28;
+    // The data type of phi::StringTensor
+    PSTRING = 29;
+    // the data type of phi::SparseCooTensor
+    SPARSE_COO = 30;
+    // the data type of phi::SparseCsrTensor
+    SPARSE_CSR = 31;
+  }
+
+  required Type type = 1;
+
+  message TensorDesc {
+    // Should only be PODType. Is enforced in C++
+    required Type data_type = 1;
+    repeated int64 dims = 2; // [UNK, 640, 480] is saved as [-1, 640, 480]
+  }
+  optional TensorDesc selected_rows = 2;
+
+  message DenseTensorDesc {
+    required TensorDesc tensor = 1;
+    optional int32 legacy_lod_level = 2 [ default = 0 ];
+  }
+  optional DenseTensorDesc dense_tensor = 3;
+
+  message DenseTensorArrayDesc {
+    required TensorDesc tensor = 1;
+    optional int32 legacy_lod_level = 2 [ default = 0 ];
+  }
+  optional DenseTensorArrayDesc tensor_array = 4;
+
+  message ReaderDesc { repeated DenseTensorDesc dense_tensor = 1; }
+  optional ReaderDesc reader = 5;
+
+  message Tuple { repeated Type element_type = 1; }
+  optional Tuple tuple = 7;
+
+  optional TensorDesc string = 8;
+  optional TensorDesc strings = 9;
+  optional TensorDesc vocab = 10;
+  optional TensorDesc sparse_coo = 11;
+  optional TensorDesc sparse_csr = 12;
+}
+
+message VarDesc {
+
+  message Attr {
+    required string name = 1;
+    required AttrType type = 2;
+    optional int32 i = 3;
+    optional string s = 4;
+    repeated int32 ints = 5;
+  };
+
+  required string name = 1;
+  required VarType type = 2;
+  optional bool persistable = 3 [ default = false ];
+  // True if the variable is an input data and
+  // have to check the feed data shape and dtype
+  optional bool need_check_feed = 4 [ default = false ];
+  optional bool is_parameter = 5 [ default = false ];
+  optional bool stop_gradient = 6 [ default = false ];
+  repeated Attr attrs = 7;
+}
+
+message BlockDesc {
+  required int32 idx = 1;
+  required int32 parent_idx = 2;
+  repeated VarDesc vars = 3;
+  repeated OpDesc ops = 4;
+  optional int32 forward_block_idx = 5 [ default = -1 ];
+}
+
+// In some cases, Paddle may perform operator definition iterations,
+// and the operator uses OpVersionMap for compatibility testing.
+message OpVersion { required int32 version = 1; }
+message OpVersionMap {
+  message OpVersionPair {
+    required string op_name = 1;
+    required OpVersion op_version = 2;
+  }
+  repeated OpVersionPair pair = 1;
+}
+
+// TODO(panyx0718): A model can have multiple programs. Need a
+// way to distinguish them. Maybe ID or name?
+message ProgramDesc {
+  reserved 2, 3; // For backward compatibility.
+  repeated BlockDesc blocks = 1;
+  optional Version version = 4;
+  optional OpVersionMap op_version_map = 5;
+}

paddle/phi/core/framework/heter_service.proto

@@ -0,0 +1,70 @@
+// 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.
+
+syntax = "proto2";
+package paddle.framework;
+option cc_generic_services = true;
+
+// It can be: phi::DenseTensor、SelectedRows or NCCL_ID
+enum VarType {
+  DENSE_TENSOR = 0;
+  SELECTED_ROWS = 1;
+  NCCL_ID = 2;
+}
+
+// VariableMessage is serialized paddle variable message.
+// NOTICE(gongwb):don't modify this proto if you are not
+//   not familiar with how we serialize in sendrecvop_utils.h
+//   and deserialize it in  variable_response.h.
+message VariableMessage {
+  enum Type {
+    // Pod Types
+    BOOL = 0;
+    INT16 = 1;
+    INT32 = 2;
+    INT64 = 3;
+    FP16 = 4;
+    FP32 = 5;
+    FP64 = 6;
+  }
+
+  message LodData { repeated int64 lod_data = 1; }
+  optional string varname = 1;
+  // TODO(Yancey1989): reference framework::proto::VarDesc::VarType
+  optional VarType type = 2;
+  // bool persistable is not needed for sending.
+  // tensor info:
+  optional Type data_type = 3;
+  repeated int64 dims = 4;
+
+  // lod details:
+  optional int64 lod_level = 5;
+  repeated LodData lod = 6;
+  // selected_rows height, aka. original dim0
+  optional int64 slr_height = 7;
+  // tensor data
+  optional bytes data = 8;
+}
+message HeterRequest {
+  required int32 cmd = 1;
+  optional int32 cur_batch = 2;
+  repeated VariableMessage vars = 3;
+};
+
+message HeterResponse {
+  // optional VariableMessage vars = 1;
+  repeated VariableMessage vars = 1;
+};
+
+service HeterService { rpc service(HeterRequest) returns (HeterResponse); };

paddle/phi/core/framework/pass_desc.proto

@@ -0,0 +1,91 @@
+// 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.
+
+syntax = "proto2";
+
+import "paddle/phi/core/framework/framework.proto";
+package paddle.framework.proto;
+
+// Describes one substitute subgraph.
+message PassDesc {
+  enum RoleType {
+    kVariable = 0;
+    kOperator = 1;
+  }
+  enum OperationType {
+    kAdd = 0;
+    kSub = 1;
+    kMul = 2;
+    kDiv = 3;
+    kSize = 4;
+    kMod = 5;
+  }
+  enum ConditionType {
+    kEQ = 0;
+    kNE = 1;
+    kGT = 2;
+    kGE = 3;
+    kLT = 4;
+    kLE = 5;
+  }
+  // Representation of attr in var or operator.
+  message Attr {
+    required RoleType role = 1;
+    optional string var_name = 2;
+    optional int32 op_index = 3;
+    required string name = 4;
+    optional string element_name = 5;
+    optional int32 element_index = 6;
+    optional OperationType operation = 7;
+  }
+  // The operation to be performed.
+  message Operation {
+    required OperationType type = 1;
+    optional Attr attr = 2;
+    optional OpDesc.Attr value = 3;
+  }
+  message VarMap {
+    required string pattern_var = 1;
+    required string replace_var = 2;
+  }
+  message AttrMap {
+    required Attr pattern_attr = 1;
+    required Attr replace_attr = 2;
+    optional Operation operation = 3;
+  }
+  message AttrCondition {
+    required Attr attr = 1;
+    required ConditionType type = 2;
+    optional Attr condition_attr = 3;
+    optional OpDesc.Attr condition_value = 4;
+    optional Operation operation = 5;
+  }
+  // A pair of subgraphs for matching and rewriting.
+  repeated OpDesc pattern = 1;
+  repeated OpDesc replace = 2;
+  // Mapping vars between pattern and replace subgraphs.
+  repeated VarMap var_maps = 3;
+  // Mapping attrs of vars and ops between pattern and replace subgraphs.
+  repeated AttrMap var_attr_maps = 4;
+  repeated AttrMap op_attr_maps = 5;
+  // Limit the attrs of vars and ops in pattern subgraph.
+  repeated AttrCondition var_attr_conditions = 6;
+  repeated AttrCondition op_attr_conditions = 7;
+}
+
+// A series of PassDesc.
+message MultiPassDesc {
+  optional string pass_type = 1;
+  repeated PassDesc pass_descs = 2;
+}

paddle/phi/core/framework/reader.cc

@@ -0,0 +1,80 @@
+// 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/framework/reader.h"
+#include <deque>
+#include "glog/logging.h"
+
+namespace paddle::framework {
+
+void ReaderBase::ReadNext(phi::TensorArray *out) {
+  std::lock_guard<std::mutex> lock(mu_);
+  PADDLE_ENFORCE_EQ(status_,
+                    ReaderStatus::kRunning,
+                    common::errors::Unavailable(
+                        "The current reader has stopped running and cannot "
+                        "continue to read the next batch of data."));
+  ReadNextImpl(out);
+}
+
+void ReaderBase::InsertDecoratedReader(
+    const std::shared_ptr<ReaderBase> &decorated_reader) {
+  std::lock_guard<std::mutex> guard(mu_);
+  decorated_readers_.emplace_back(decorated_reader);
+}
+
+std::unordered_set<ReaderBase *> ReaderBase::GetEndPoints() {
+  std::unordered_set<ReaderBase *> result;
+  std::deque<ReaderBase *> queue;
+  queue.emplace_back(this);
+  while (!queue.empty()) {  // BFS search
+    auto *front = queue.front();
+    queue.pop_front();
+    if (front->decorated_readers_.empty()) {
+      result.emplace(front);
+    } else {
+      for (auto &reader : front->decorated_readers_) {
+        if (auto *reader_ptr = reader.lock().get()) {
+          queue.emplace_back(reader_ptr);
+        }
+      }
+    }
+  }
+
+  return result;
+}
+
+void ReaderBase::Shutdown() {
+  std::lock_guard<std::mutex> lock(mu_);
+  if (status_ != ReaderStatus::kStopped) {
+    ShutdownImpl();
+    status_ = ReaderStatus::kStopped;
+  }
+}
+
+void ReaderBase::Start() {
+  std::lock_guard<std::mutex> lock(mu_);
+  if (status_ != ReaderStatus::kRunning) {
+    StartImpl();
+    status_ = ReaderStatus::kRunning;
+  }
+}
+
+ReaderBase::~ReaderBase() = default;
+
+DecoratedReader::~DecoratedReader() {
+  VLOG(1) << "~DecoratedReader";
+  reader_->Shutdown();
+}
+}  // namespace paddle::framework

paddle/phi/core/framework/reader.h

@@ -0,0 +1,224 @@
+// 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 <memory>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/framework/framework.pb.h"
+#include "paddle/phi/core/tensor_array.h"
+
+namespace paddle {
+namespace framework {
+class ReaderBase {
+ public:
+  explicit ReaderBase(const std::vector<DDim>& shapes,
+                      const std::vector<proto::VarType::Type>& var_types,
+                      const std::vector<bool>& need_check_feed)
+      : shapes_(shapes),
+        var_types_(var_types),
+        need_check_feed_(need_check_feed) {
+    PADDLE_ENFORCE_EQ(
+        shapes_.size(),
+        need_check_feed_.size(),
+        common::errors::InvalidArgument(
+            "Construct ReaderBase with mismatched sizes of shapes "
+            "and need_check_feed"));
+    PADDLE_ENFORCE_EQ(
+        var_types_.size(),
+        need_check_feed_.size(),
+        common::errors::InvalidArgument(
+            "Construct ReaderBase with mismatched sizes of var_types "
+            "and need_check_feed"));
+  }
+
+  PADDLE_API virtual void ReadNext(phi::TensorArray* out);
+
+  PADDLE_API virtual void Shutdown();
+
+  PADDLE_API virtual void Start();
+
+  // Return the readers which are the end of decorating chain. Basically
+  // they are readers just before read op.
+  PADDLE_API std::unordered_set<ReaderBase*> GetEndPoints();
+
+  // Returns the shapes of the fed variables
+  const std::vector<DDim>& Shapes() const { return shapes_; }
+
+  // Returns the dtypes of the fed variables
+  const std::vector<proto::VarType::Type>& VarTypes() const {
+    return var_types_;
+  }
+
+  // For Backward compatibility, old fluid.layers.data doesn't check shape.
+  // This function returns whether you have the check shape for this Reader.
+  const std::vector<bool>& NeedCheckFeed() const { return need_check_feed_; }
+
+  PADDLE_API virtual ~ReaderBase();
+
+ protected:
+  virtual void ReadNextImpl(phi::TensorArray* out UNUSED) {}
+
+  virtual void ShutdownImpl() {}
+
+  virtual void StartImpl() {}
+
+  enum ReaderStatus { kRunning, kStopped };
+
+  ReaderStatus status_{kRunning};
+
+  mutable std::mutex mu_;
+
+  // The shapes of the fed variables.
+  std::vector<DDim> shapes_;
+
+  // The dtypes of the fed variables.
+  std::vector<proto::VarType::Type> var_types_;
+
+  // Whether to check the shape and dtype of fed variables.
+  std::vector<bool> need_check_feed_;
+
+ private:
+  friend class DecoratedReader;
+  // These methods can be only invoked inside DecoratedReader to record the
+  // decorating chain.
+  PADDLE_API void InsertDecoratedReader(
+      const std::shared_ptr<ReaderBase>& decorated_reader);
+  // A set of which readers that decorated this reader.
+  std::vector<std::weak_ptr<ReaderBase>> decorated_readers_;
+};
+
+class DecoratedReader : public ReaderBase,
+                        public std::enable_shared_from_this<DecoratedReader> {
+ public:
+  explicit DecoratedReader(const std::shared_ptr<ReaderBase>& reader)
+      : ReaderBase(
+            reader->Shapes(), reader->VarTypes(), reader->NeedCheckFeed()),
+        reader_(reader) {
+    PADDLE_ENFORCE_NOT_NULL(
+        reader_,
+        common::errors::InvalidArgument(
+            "The underlying reader of DecoratedReader should not be null"));
+  }
+
+  void RegisterDecorateChain() {
+    reader_->InsertDecoratedReader(shared_from_this());
+  }
+
+  PADDLE_API ~DecoratedReader();
+
+  const std::shared_ptr<ReaderBase>& UnderlyingReader() const {
+    return reader_;
+  }
+
+ protected:
+  void ShutdownImpl() override { reader_->Shutdown(); }
+
+  void StartImpl() override { reader_->Start(); }
+
+  std::shared_ptr<ReaderBase> reader_;
+};
+
+// FileReader is just a conceptual class.
+class FileReader : public ReaderBase {
+ public:
+  explicit FileReader(const std::vector<DDim>& shapes,
+                      const std::vector<proto::VarType::Type>& var_types,
+                      const std::vector<bool>& need_check_feed)
+      : ReaderBase(shapes, var_types, need_check_feed) {}
+};
+
+// The ReaderHolder is used as reader' unified wrapper,
+// making it easier to access different type reader in Variables.
+class ReaderHolder {
+ public:
+  template <typename T>
+  void Reset(const std::shared_ptr<T>& reader) {
+    auto reader_base = std::dynamic_pointer_cast<ReaderBase>(reader);
+    PADDLE_ENFORCE_NOT_NULL(
+        reader_base,
+        common::errors::InvalidArgument(
+            "The underlying reader of ReaderHolder should not be null"));
+    reader_ = reader_base;
+  }
+
+  ~ReaderHolder() {}
+
+  const std::shared_ptr<ReaderBase>& Get() const { return reader_; }
+
+  void ReadNext(phi::TensorArray* out) {
+    PADDLE_ENFORCE_NOT_NULL(
+        reader_,
+        common::errors::InvalidArgument(
+            "The underlying reader of ReaderHolder should not be null"));
+    reader_->ReadNext(out);
+  }
+
+  void ResetAll() {
+    auto end_readers = reader_->GetEndPoints();
+    for (auto* reader : end_readers) {
+      reader->Shutdown();
+    }
+    for (auto* reader : end_readers) {
+      reader->Start();
+    }
+  }
+
+  void Shutdown() {
+    PADDLE_ENFORCE_NOT_NULL(
+        reader_,
+        common::errors::InvalidArgument(
+            "The underlying reader of ReaderHolder should not be null"));
+    reader_->Shutdown();
+  }
+
+  void Start() {
+    PADDLE_ENFORCE_NOT_NULL(
+        reader_,
+        common::errors::InvalidArgument(
+            "The underlying reader of ReaderHolder should not be null"));
+    reader_->Start();
+  }
+
+  const std::vector<DDim>& Shapes() const { return reader_->Shapes(); }
+
+  const std::vector<proto::VarType::Type>& VarTypes() const {
+    return reader_->VarTypes();
+  }
+
+  const std::vector<bool>& NeedCheckFeed() const {
+    return reader_->NeedCheckFeed();
+  }
+
+  void Clear() { reader_.reset(); }
+
+  operator const std::shared_ptr<ReaderBase>&() const { return this->reader_; }
+
+ private:
+  std::shared_ptr<ReaderBase> reader_;
+};
+
+template <typename T, typename... ARGS>
+inline std::shared_ptr<DecoratedReader> MakeDecoratedReader(ARGS&&... args) {
+  std::shared_ptr<DecoratedReader> reader(new T(std::forward<ARGS>(args)...));
+  reader->RegisterDecorateChain();
+  return reader;
+}
+}  // namespace framework
+}  // namespace paddle

paddle/phi/core/framework/selected_rows_serialize.cc

@@ -0,0 +1,95 @@
+// 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/framework/selected_rows_serialize.h"
+
+namespace phi {
+
+void SerializeToStream(std::ostream& os,
+                       const SelectedRows& selected_rows,
+                       const DeviceContext& dev_ctx) {
+  {  // the 1st field, uint32_t version
+    constexpr uint32_t version = 0;
+    os.write(reinterpret_cast<const char*>(&version), sizeof(version));
+  }
+  {
+    // the 2nd field, rows information
+    auto& rows = selected_rows.rows();
+    uint64_t size = rows.size();
+    os.write(reinterpret_cast<const char*>(&size), sizeof(size));
+    for (uint64_t i = 0; i < size; ++i) {
+      os.write(reinterpret_cast<const char*>(&rows[i]), sizeof(rows[i]));
+    }
+  }
+  {
+    // the 3rd field, the height of SelectedRows
+    int64_t height = selected_rows.height();
+    os.write(reinterpret_cast<const char*>(&height), sizeof(height));
+  }
+  // the 4st field, Tensor data
+  TensorToStream(os, selected_rows.value(), dev_ctx);
+}
+
+void SerializeToStream(std::ostream& os, const SelectedRows& selected_rows) {
+  DeviceContextPool& pool = DeviceContextPool::Instance();
+  const DeviceContext* dev_ctx = nullptr;
+  auto place = selected_rows.place();
+  dev_ctx = pool.Get(place);
+  SerializeToStream(os, selected_rows, *dev_ctx);
+}
+
+void DeserializeFromStream(std::istream& is, SelectedRows* selected_rows) {
+  DeviceContextPool& pool = DeviceContextPool::Instance();
+  const DeviceContext* dev_ctx = nullptr;
+  dev_ctx = pool.Get(CPUPlace());
+  DeserializeFromStream(is, selected_rows, *dev_ctx);
+}
+
+void DeserializeFromStream(std::istream& is,
+                           SelectedRows* selected_rows,
+                           const DeviceContext& dev_ctx) {
+  {
+    // the 1st field, unit32_t version for SelectedRows
+    uint32_t version = 0;
+    is.read(reinterpret_cast<char*>(&version), sizeof(version));
+    PADDLE_ENFORCE_EQ(version,
+                      0U,
+                      common::errors::InvalidArgument(
+                          "Only version 0 SelectedRows is supported."));
+  }
+  {
+    // the 2nd field, rows information
+    uint64_t size = 0;
+    is.read(reinterpret_cast<char*>(&size), sizeof(size));
+    PADDLE_ENFORCE_EQ(
+        is.good(),
+        true,
+        common::errors::Unavailable("Cannot read the number of rows."));
+    auto& rows = *selected_rows->mutable_rows();
+    rows.resize(size);
+    for (uint64_t i = 0; i < size; ++i) {
+      is.read(reinterpret_cast<char*>(&rows[i]), sizeof(int64_t));
+    }
+  }
+  {
+    // the 3rd field, the height of the SelectedRows
+    int64_t height = 0;
+    is.read(reinterpret_cast<char*>(&height), sizeof(int64_t));
+    selected_rows->set_height(height);
+  }
+  // the 4st field, tensor which contains the data
+  TensorFromStream(is, selected_rows->mutable_value(), dev_ctx);
+}
+
+}  // namespace phi

paddle/phi/core/framework/selected_rows_serialize.h

@@ -0,0 +1,47 @@
+// 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 <algorithm>
+#include <memory>
+#include <mutex>  // NOLINT
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "paddle/phi/core/framework/dense_tensor_tostream.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+/*
+ * Serialize/Deserialize SelectedRows to std::ostream
+ * You can pass ofstream or ostringstream to serialize to file
+ * or to a in memory string. GPU tensor will be copied to CPU.
+ */
+PADDLE_API void SerializeToStream(std::ostream& os,
+                                  const SelectedRows& selected_rows,
+                                  const DeviceContext& dev_ctx);
+PADDLE_API void DeserializeFromStream(std::istream& is,
+                                      SelectedRows* selected_rows,
+                                      const DeviceContext& dev_ctx);
+
+PADDLE_API void SerializeToStream(std::ostream& os,
+                                  const SelectedRows& selected_rows);
+
+PADDLE_API void DeserializeFromStream(std::istream& is,
+                                      SelectedRows* selected_rows);
+
+}  // namespace phi

paddle/phi/core/framework/trainer_desc.proto

@@ -0,0 +1,220 @@
+// 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.
+
+syntax = "proto2";
+option optimize_for = LITE_RUNTIME;
+import "paddle/phi/core/framework/data_feed.proto";
+import "paddle/phi/core/framework/framework.proto";
+package paddle.framework;
+
+message TrainerDesc {
+  // class name for create trainer desc
+  // the matchness of trainer name and device worker name
+  // will be checked in python API
+  optional string class_name = 1;
+  // class name for creating device worker
+  optional string device_worker_name = 2;
+  // thread number
+  optional int32 thread_num = 3;
+  // if we need to binding cpu
+  optional bool binding_cpu = 4 [ default = false ];
+  repeated string filelist = 5;
+  optional bool debug = 6 [ default = false ];
+  optional FetchConfig fetch_config = 7;
+  optional bool use_cvm = 8 [ default = false ];
+  optional bool dump_slot = 9 [ default = false ];
+  optional float scale_datanorm = 10 [ default = -1 ];
+  optional int32 mpi_rank = 11 [ default = -1 ];
+  optional string dump_fields_path = 12;
+  repeated string dump_fields = 13;
+  optional string dump_converter = 14;
+  repeated string dump_param = 15;
+  optional int32 mpi_size = 16 [ default = -1 ];
+  optional int32 dump_file_num = 17 [ default = 16 ];
+  repeated string check_nan_var_names = 18;
+  optional CopyTableConfig copy_table_config = 19;
+  // adjust ins weight
+  optional AdjustInsWeightConfig adjust_ins_weight_config = 20;
+  optional bool no_cvm = 21 [ default = false ];
+  optional bool thread_barrier = 22;
+  repeated string loss_names = 23;
+  optional bool enable_random_dump = 24 [ default = false ];
+  optional bool random_with_lineid = 25 [ default = false ];
+  optional int32 dump_interval = 26 [ default = 10000 ];
+  repeated int32 worker_places = 27;
+
+  repeated string xpu_send_list = 28;
+  repeated string xpu_recv_list = 29;
+  optional int32 xpu_start_idx = 30;
+  optional int32 xpu_end_idx = 31;
+
+  optional bool use_ps_gpu = 32 [ default = false ];
+  optional string user_define_dump_filename = 33;
+  optional bool scale_sparse_gradient_with_batch_size = 34 [ default = true ];
+
+  repeated int32 trainers = 35;
+  optional int32 trainer_id = 36;
+
+  // add for gpu
+  optional string fleet_desc = 37;
+  optional bool is_dump_in_simple_mode = 38 [ default = false ];
+  optional string dump_fields_mode = 39 [ default = "w" ];
+  optional int32 dump_num_decimals = 40 [ default = 9 ];
+  optional bool use_gpu_graph = 41 [ default = false ];
+  // device worker parameters
+  optional HogwildWorkerParameter hogwild_param = 101;
+  optional DownpourWorkerParameter downpour_param = 103;
+  optional PullDenseWorkerParameter pull_dense_param = 102;
+  optional SectionWorkerParameter section_param = 104;
+  optional HeterSectionWorkerParameter heter_section_param = 105;
+  // datafeed desc
+  optional DataFeedDesc data_desc = 201;
+}
+
+message HogwildWorkerParameter {
+  repeated string skip_ops = 1;
+  repeated string stat_var_names = 2;
+}
+
+message DownpourWorkerParameter {
+  repeated TableParameter sparse_table = 1;
+  repeated TableParameter dense_table = 2;
+  repeated string skip_ops = 3;
+  repeated ProgramConfig program_config = 4;
+  optional bool push_sparse = 5 [ default = true ];
+  optional bool push_dense = 6 [ default = true ];
+  repeated string stat_var_names = 7;
+}
+
+message SectionWorkerParameter {
+  optional SectionConfig section_config = 1;
+  optional int32 queue_size = 2 [ default = 1 ];
+  optional int64 sync_steps = 3 [ default = 1 ];
+  optional int32 start_cpu_core_id = 4 [ default = 1 ];
+  repeated string param_need_sync = 5;
+  optional int32 num_microbatches = 6;
+  optional int32 num_pipeline_stages = 7 [ default = 1 ];
+  optional int32 pipeline_stage = 8 [ default = 1 ];
+  optional int32 schedule_mode = 9 [ default = 0 ];
+}
+
+message HeterSectionWorkerParameter {
+  optional SectionConfig section_config = 1;
+  optional int32 queue_size = 2 [ default = 1 ];
+  optional int64 sync_steps = 3 [ default = 1 ];
+  optional int32 start_cpu_core_id = 4 [ default = 1 ];
+  repeated string param_need_sync = 5;
+  optional int32 num_microbatches = 6;
+  optional int32 num_pipeline_stages = 7 [ default = 1 ];
+  optional int32 pipeline_stage = 8 [ default = 1 ];
+}
+
+message SectionConfig {
+  enum Place {
+    CPUPlace = 0;
+    CUDAPlace = 1;
+    CUDAPinnedPlace = 2;
+  }
+
+  // FIXME: How to use proto::ProgramDesc
+  // required string program_desc_str = 1;
+  optional proto.ProgramDesc program_desc = 1;
+  optional Place place = 2;
+  optional int32 concurrency = 3 [ default = 1 ];
+  repeated string section_in_var_names = 4;
+  repeated string section_out_var_names = 5;
+  optional int32 place_id = 6 [ default = -1 ];
+}
+
+message FetchConfig {
+  enum Method { PRINT = 0; }
+  repeated string fetch_var_names = 1;
+  repeated string fetch_var_str_format = 2;
+  optional int32 print_period = 3 [ default = 100 ];
+  optional Method method = 4 [ default = PRINT ];
+}
+
+message AdjustInsWeightConfig {
+  optional bool need_adjust = 1 [ default = false ];
+  optional string nid_slot = 2 [ default = "" ];
+  optional float nid_adjw_threshold = 3 [ default = 0.0 ];
+  optional float nid_adjw_ratio = 4 [ default = 0.0 ];
+  optional string ins_weight_slot = 5 [ default = "" ];
+}
+
+message TableDependencyMap {
+  required int32 key = 1;
+  repeated int32 values = 2;
+}
+
+message CopyTableConfig {
+  optional bool need_copy = 1 [ default = false ];
+  optional int32 batch_num = 2 [ default = 100 ];
+  repeated int32 src_sparse_tables = 3;
+  repeated int32 dest_sparse_tables = 4;
+  repeated int32 src_dense_tables = 5;
+  repeated int32 dest_dense_tables = 6;
+  repeated string src_var_list = 7;
+  repeated string dest_var_list = 8;
+  // when dest dense table has no grad, should pull explicitly
+  optional bool dense_pull_after_copy = 9 [ default = false ];
+  // copy feasigns or copy the whole table
+  optional bool sparse_copy_by_feasign = 10 [ default = true ];
+  // table dependency for pull/push
+  optional bool enable_dependency = 11 [ default = false ];
+  repeated TableDependencyMap table_dependency_map = 12;
+}
+
+message CondTableMap {
+  required int32 key = 1;
+  required int32 value = 2;
+}
+message ProgramConfig {
+  required string program_id = 1;
+  repeated int32 push_sparse_table_id = 2;
+  repeated int32 push_dense_table_id = 3;
+  repeated int32 pull_sparse_table_id = 4;
+  repeated int32 pull_dense_table_id = 5;
+  repeated CondTableMap partial_pushdense_condtable_map = 10;
+}
+
+message PullDenseWorkerParameter {
+  // dense table only and specialized usage
+  optional int32 threshold = 1 [ default = 1 ];
+  optional int32 device_num = 2;
+  optional int32 sleep_time_ms = 3 [ default = 2 ];
+  repeated TableParameter dense_table = 4;
+}
+
+message TableParameter {
+  // dense table only
+  optional uint64 table_id = 1;
+  repeated string dense_value_name = 2;
+  repeated string dense_grad_name = 3;
+  repeated int32 push_dense_wait_times = 5;
+  // sparse table only
+  repeated string sparse_key_name = 6;
+  repeated string sparse_value_name = 7;
+  repeated string sparse_grad_name = 8;
+  repeated int32 push_sparse_wait_times = 9;
+  // sparse table only and specialized usage
+  optional int32 emb_dim = 10;
+  optional int32 fea_dim = 11;
+  optional string label_var_name = 12;
+  // if table will pull sparse to local first
+  optional bool is_local = 13 [ default = false ];
+  // if table will pull sparse asynchronously in worker
+  optional bool is_async = 14 [ default = false ];
+  optional string async_wait_op_name = 15;
+}

paddle/phi/core/framework/var_type_helper.cc

@@ -0,0 +1,154 @@
+// 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/framework/var_type_helper.h"
+#include "paddle/phi/common/pstring.h"
+
+namespace phi {
+
+struct DataTypeMap {
+  std::unordered_map<std::type_index, VarType::Type> cpp_to_proto_;
+  std::unordered_map<int, std::type_index> proto_to_cpp_;
+  std::unordered_map<int, std::string> proto_to_str_;
+  std::unordered_map<int, size_t> proto_to_size_;
+};
+
+static DataTypeMap* InitDataTypeMap();
+// C++11 removes the need for manual locking. Concurrent execution shall wait if
+// a static local variable is already being initialized.
+// https://stackoverflow.com/questions/11711920/how-to-implement-multithread-safe-singleton-in-c11-without-using-mutex
+static DataTypeMap& gDataTypeMap() {
+  static DataTypeMap* g_data_type_map_ = InitDataTypeMap();
+  return *g_data_type_map_;
+}
+
+template <typename T>
+static inline void RegisterType(DataTypeMap* map,
+                                VarType::Type proto_type,
+                                const std::string& name) {
+  map->proto_to_cpp_.emplace(static_cast<int>(proto_type), typeid(T));
+  map->cpp_to_proto_.emplace(typeid(T), proto_type);
+  map->proto_to_str_.emplace(static_cast<int>(proto_type), name);
+  map->proto_to_size_.emplace(static_cast<int>(proto_type), sizeof(T));
+}
+
+static DataTypeMap* InitDataTypeMap() {
+  auto retv = new DataTypeMap();
+
+#define RegType(cc_type, proto_type) \
+  RegisterType<cc_type>(retv, proto_type, #cc_type)
+
+  _ForEachDataType_(RegType);
+  // Register pstring individually
+  RegType(pstring, VarType::PSTRING);
+  RegType(::phi::dtype::float8_e5m2, VarType::FP8_E5M2);
+  RegType(::phi::dtype::float8_e4m3fn, VarType::FP8_E4M3FN);
+#undef RegType
+  return retv;
+}
+
+VarType::Type ToDataType(std::type_index type) {
+  auto it = gDataTypeMap().cpp_to_proto_.find(type);
+  if (it != gDataTypeMap().cpp_to_proto_.end()) {
+    return it->second;
+  }
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support %s as tensor data type.", common::demangle(type.name())));
+}
+
+std::type_index ToTypeIndex(VarType::Type type) {
+  auto it = gDataTypeMap().proto_to_cpp_.find(static_cast<int>(type));
+  if (it != gDataTypeMap().proto_to_cpp_.end()) {
+    return it->second;
+  }
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support VarType::Type(%d) as tensor type.", static_cast<int>(type)));
+}
+
+std::string VarDataTypeToString(const VarType::Type type) {
+  auto it = gDataTypeMap().proto_to_str_.find(static_cast<int>(type));
+  if (it != gDataTypeMap().proto_to_str_.end()) {
+    return it->second;
+  }
+  // deal with RAW type
+  if (type == VarType::RAW) {
+    return "RAW(runtime decided type)";
+  }
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not support VarType::Type(%d) as tensor type.", static_cast<int>(type)));
+}
+
+size_t SizeOfType(VarType::Type type) {
+  auto it = gDataTypeMap().proto_to_size_.find(static_cast<int>(type));
+  if (it != gDataTypeMap().proto_to_size_.end()) {
+    return it->second;
+  }
+  PADDLE_THROW(common::errors::Unimplemented("Not support %s as tensor type.",
+                                             VarDataTypeToString(type)));
+}
+
+// Now only supports promotion of complex type
+inline bool NeedPromoteTypes(const VarType::Type& a, const VarType::Type& b) {
+  return (IsComplexType(a) || IsComplexType(b));
+}
+
+int DataTypeNumAlign(const VarType::Type t) {
+  int cast_type_num = -1;
+  if (t == VarType::FP32 || t == VarType::FP64) {
+    cast_type_num = static_cast<int>(t) - 5;
+  } else if (t == VarType::COMPLEX64 || t == VarType::COMPLEX128) {
+    cast_type_num = static_cast<int>(t) - 21;
+  } else {
+    PADDLE_THROW(common::errors::Unavailable(
+        "Only supports to align data type include float32, float64, complex64 "
+        "and complex128, but received data type is %s.",
+        VarDataTypeToString(t)));
+  }
+  return cast_type_num;
+}
+
+// Now only supports promotion of complex type
+VarType::Type PromoteTypesIfComplexExists(const VarType::Type type_a,
+                                          const VarType::Type type_b) {
+  constexpr auto f4 = VarType::FP32;        // 5
+  constexpr auto f8 = VarType::FP64;        // 6
+  constexpr auto c4 = VarType::COMPLEX64;   // 23
+  constexpr auto c8 = VarType::COMPLEX128;  // 24
+
+  if (!NeedPromoteTypes(type_a, type_b)) {
+    // NOTE(chenweihang): keep consistent with rule in original op's impl,
+    // kernel type based on the first input tensor's dtype
+    return type_a;
+  }
+
+  int type_an = DataTypeNumAlign(type_a);
+  int type_bn = DataTypeNumAlign(type_b);
+
+  // Here is a complete rules table, but some rules are not used.
+  // It is still written this way because array accessing is still
+  // more efficient than if-else
+  // NOLINTBEGIN(*-avoid-c-arrays)
+  static constexpr VarType::Type promote_types_table[4][4] = {
+      /*        f4  f8  c4  c8*/
+      /* f4 */ {f4, f8, c4, c8},
+      /* f8 */ {f8, f8, c8, c8},
+      /* c4 */ {c4, c8, c4, c8},
+      /* c8 */ {c8, c8, c8, c8},
+  };
+  // NOLINTEND(*-avoid-c-arrays)
+
+  return promote_types_table[type_an][type_bn];
+}
+
+}  // namespace phi

paddle/phi/core/framework/var_type_helper.h

@@ -0,0 +1,275 @@
+// 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 <iostream>
+#include <string>
+#include <typeindex>
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/framework/framework.pb.h"
+
+namespace phi {
+using VarType = paddle::framework::proto::VarType;
+
+PADDLE_API std::string VarDataTypeToString(const VarType::Type type);
+TEST_API extern size_t SizeOfType(VarType::Type type);
+
+template <typename T>
+struct IsComplex : public std::false_type {};
+
+template <typename T>
+struct IsComplex<phi::dtype::complex<T>> : public std::true_type {};
+
+template <typename T>
+struct DataTypeTrait {};
+
+// Stub handle for void
+template <>
+struct DataTypeTrait<void> {
+  constexpr static VarType::Type DataType() { return VarType::RAW; }
+};
+
+#define _ForEachDataTypeHelper_(callback, cpp_type, proto_type) \
+  callback(cpp_type, ::paddle::framework::proto::VarType::proto_type);
+
+#define _ForEachDataType_(callback)                                           \
+  _ForEachDataTypeHelper_(callback, float, FP32);                             \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float16, FP16);             \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::bfloat16, BF16);            \
+  _ForEachDataTypeHelper_(callback, double, FP64);                            \
+  _ForEachDataTypeHelper_(callback, int, INT32);                              \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);                          \
+  _ForEachDataTypeHelper_(callback, bool, BOOL);                              \
+  _ForEachDataTypeHelper_(callback, uint8_t, UINT8);                          \
+  _ForEachDataTypeHelper_(callback, uint16_t, UINT16);                        \
+  _ForEachDataTypeHelper_(callback, uint32_t, UINT32);                        \
+  _ForEachDataTypeHelper_(callback, uint64_t, UINT64);                        \
+  _ForEachDataTypeHelper_(callback, int16_t, INT16);                          \
+  _ForEachDataTypeHelper_(callback, int8_t, INT8);                            \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::complex<float>, COMPLEX64); \
+  _ForEachDataTypeHelper_(                                                    \
+      callback, ::phi::dtype::complex<double>, COMPLEX128);                   \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float8_e4m3fn, FP8_E4M3FN); \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float8_e5m2, FP8_E5M2);
+
+#define _ForEachIntDataType_(callback)                 \
+  _ForEachDataTypeHelper_(callback, int, INT32);       \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);   \
+  _ForEachDataTypeHelper_(callback, uint8_t, UINT8);   \
+  _ForEachDataTypeHelper_(callback, uint16_t, UINT16); \
+  _ForEachDataTypeHelper_(callback, uint32_t, UINT32); \
+  _ForEachDataTypeHelper_(callback, uint64_t, UINT64); \
+  _ForEachDataTypeHelper_(callback, int16_t, INT16);   \
+  _ForEachDataTypeHelper_(callback, int8_t, INT8);
+
+#define _ForEachDataTypeSmall_(callback)                                      \
+  _ForEachDataTypeHelper_(callback, float, FP32);                             \
+  _ForEachDataTypeHelper_(callback, double, FP64);                            \
+  _ForEachDataTypeHelper_(callback, int, INT32);                              \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);                          \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::complex<float>, COMPLEX64); \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::complex<double>, COMPLEX128);
+
+#define _ForEachDataTypeNormal_(callback)                         \
+  _ForEachDataTypeHelper_(callback, float, FP32);                 \
+  _ForEachDataTypeHelper_(callback, double, FP64);                \
+  _ForEachDataTypeHelper_(callback, int, INT32);                  \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);              \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float16, FP16); \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::bfloat16, BF16);
+
+// For the use of thrust, as index-type elements can be only integers.
+#define _ForEachDataTypeTiny_(callback)          \
+  _ForEachDataTypeHelper_(callback, int, INT32); \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);
+
+// It's only for DataParallel in HIP, bf16 not support in HIP.
+#define _ForEachDataTypeForHIP_(callback)                                     \
+  _ForEachDataTypeHelper_(callback, float, FP32);                             \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float16, FP16);             \
+  _ForEachDataTypeHelper_(callback, double, FP64);                            \
+  _ForEachDataTypeHelper_(callback, int, INT32);                              \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);                          \
+  _ForEachDataTypeHelper_(callback, bool, BOOL);                              \
+  _ForEachDataTypeHelper_(callback, uint8_t, UINT8);                          \
+  _ForEachDataTypeHelper_(callback, uint16_t, UINT16);                        \
+  _ForEachDataTypeHelper_(callback, uint32_t, UINT32);                        \
+  _ForEachDataTypeHelper_(callback, uint64_t, UINT64);                        \
+  _ForEachDataTypeHelper_(callback, int16_t, INT16);                          \
+  _ForEachDataTypeHelper_(callback, int8_t, INT8);                            \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::complex<float>, COMPLEX64); \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::complex<double>, COMPLEX128);
+
+// complex and float8 are not supported on XPU.
+#define _ForEachDataTypeForXPU_(callback)                          \
+  _ForEachDataTypeHelper_(callback, float, FP32);                  \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::float16, FP16);  \
+  _ForEachDataTypeHelper_(callback, ::phi::dtype::bfloat16, BF16); \
+  _ForEachDataTypeHelper_(callback, double, FP64);                 \
+  _ForEachDataTypeHelper_(callback, int, INT32);                   \
+  _ForEachDataTypeHelper_(callback, int64_t, INT64);               \
+  _ForEachDataTypeHelper_(callback, bool, BOOL);                   \
+  _ForEachDataTypeHelper_(callback, uint8_t, UINT8);               \
+  _ForEachDataTypeHelper_(callback, int16_t, INT16);               \
+  _ForEachDataTypeHelper_(callback, int8_t, INT8);
+
+#define DefineDataTypeTrait(cpp_type, proto_type)                         \
+  template <>                                                             \
+  struct DataTypeTrait<cpp_type> {                                        \
+    constexpr static paddle::framework::proto::VarType::Type DataType() { \
+      return proto_type;                                                  \
+    }                                                                     \
+  }
+
+_ForEachDataType_(DefineDataTypeTrait);
+
+#undef DefineDataTypeTrait
+
+TEST_API extern VarType::Type ToDataType(std::type_index type);
+extern std::type_index ToTypeIndex(VarType::Type type);
+
+template <typename Visitor>
+inline void VisitDataType(VarType::Type type, Visitor visitor) {
+#define VisitDataTypeCallback(cpp_type, proto_type) \
+  do {                                              \
+    if (type == proto_type) {                       \
+      visitor.template apply<cpp_type>();           \
+      return;                                       \
+    }                                               \
+  } while (0)
+
+  _ForEachDataType_(VisitDataTypeCallback);
+#undef VisitDataTypeCallback
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Not supported paddle::framework::proto::VarType::Type(%d) as data type.",
+      static_cast<int>(type)));
+}
+
+template <typename Visitor>
+inline void VisitDataTypeSmall(VarType::Type type, Visitor visitor) {
+#define VisitDataTypeCallbackSmall(cpp_type, proto_type) \
+  do {                                                   \
+    if (type == proto_type) {                            \
+      visitor.template apply<cpp_type>();                \
+      return;                                            \
+    }                                                    \
+  } while (0)
+
+  _ForEachDataTypeSmall_(VisitDataTypeCallbackSmall);
+#undef VisitDataTypeCallbackSmall
+}
+
+// for normal dtype, int, int64, float, float64, float16
+template <typename Visitor>
+inline void VisitDataTypeNormal(VarType::Type type, Visitor visitor) {
+#define VisitDataTypeCallbackNormal(cpp_type, proto_type) \
+  do {                                                    \
+    if (type == proto_type) {                             \
+      visitor.template apply<cpp_type>();                 \
+      return;                                             \
+    }                                                     \
+  } while (0)
+
+  _ForEachDataTypeNormal_(VisitDataTypeCallbackNormal);
+#undef VisitDataTypeCallbackNormal
+}
+
+template <typename Visitor>
+inline void VisitIntDataType(VarType::Type type, Visitor visitor) {
+#define VisitIntDataTypeCallback(cpp_type, proto_type) \
+  do {                                                 \
+    if (type == proto_type) {                          \
+      visitor.template apply<cpp_type>();              \
+      return;                                          \
+    }                                                  \
+  } while (0)
+
+  _ForEachIntDataType_(VisitIntDataTypeCallback);
+
+  PADDLE_THROW(common::errors::Unimplemented(
+      "Expected integral data type, but got %s", VarDataTypeToString(type)));
+
+#undef VisitIntDataTypeCallback
+}
+
+template <typename Visitor>
+inline void VisitDataTypeTiny(VarType::Type type, Visitor visitor) {
+#define VisitDataTypeCallbackTiny(cpp_type, proto_type) \
+  do {                                                  \
+    if (type == proto_type) {                           \
+      visitor.template apply<cpp_type>();               \
+      return;                                           \
+    }                                                   \
+  } while (0)
+
+  _ForEachDataTypeTiny_(VisitDataTypeCallbackTiny);
+#undef VisitDataTypeCallbackTiny
+}
+
+template <typename Visitor>
+inline void VisitDataTypeForHIP(VarType::Type type, Visitor visitor) {
+#define VisitDataTypeCallbackHIP(cpp_type, proto_type) \
+  do {                                                 \
+    if (type == proto_type) {                          \
+      visitor.template apply<cpp_type>();              \
+      return;                                          \
+    }                                                  \
+  } while (0)
+
+  _ForEachDataTypeForHIP_(VisitDataTypeCallbackHIP);
+#undef VisitDataTypeCallbackHIP
+}
+
+inline std::ostream& operator<<(std::ostream& out, const VarType::Type& type) {
+  out << VarDataTypeToString(type);
+  return out;
+}
+
+extern inline bool IsComplexType(const VarType::Type& type) {
+  return (type == VarType::COMPLEX64 || type == VarType::COMPLEX128);
+}
+
+extern VarType::Type PromoteTypesIfComplexExists(const VarType::Type type_a,
+                                                 const VarType::Type type_b);
+
+extern inline VarType::Type ToComplexType(VarType::Type t) {
+  switch (t) {
+    case VarType::FP32:
+      return VarType::COMPLEX64;
+    case VarType::FP64:
+      return VarType::COMPLEX128;
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unknown real value data type (%s), now only support float32 and "
+          "float64.",
+          VarDataTypeToString(t)));
+  }
+}
+
+extern inline VarType::Type ToRealType(VarType::Type t) {
+  switch (t) {
+    case VarType::COMPLEX64:
+      return VarType::FP32;
+    case VarType::COMPLEX128:
+      return VarType::FP64;
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unknown complex value data type (%s), now only support complex64 "
+          "and complex128.",
+          VarDataTypeToString(t)));
+  }
+}
+
+}  // namespace phi