[PHI]Add new Tensor type and migrate save_combine kernel

paddle/fluid/framework/CMakeLists.txt

@@ -237,7 +237,8 @@ cc_test(
 cc_library(
   var_type_traits
   SRCS var_type_traits.cc
-  DEPS framework_proto scope tensor_array sparse_coo_tensor sparse_csr_tensor)
+  DEPS framework_proto scope tensor_array sparse_coo_tensor sparse_csr_tensor
+       extended_tensor)
 if(WITH_GPU)
   target_link_libraries(var_type_traits dynload_cuda)
 endif()

paddle/fluid/framework/operator.cc

@@ -2937,6 +2937,9 @@ void OperatorWithKernel::BuildPhiKernelContext(
         need_prepare_phi_data_ = true;
         tensor_in = &(var->Get<framework::LoDTensorArray>());
         phi_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
+      } else if (var->IsType<framework::Vocab>()) {
+        tensor_in = &(var->Get<framework::Vocab>());
+        phi_kernel_context->EmplaceBackInputWithoutSetRange(tensor_in);
       } else {
         PADDLE_THROW(platform::errors::Unimplemented(
             "Unsupported input `%s` type when call pt kernel.",
@@ -3085,6 +3088,20 @@ void OperatorWithKernel::BuildPhiKernelContext(
           }
         }
         break;
+      case phi::AttributeType::STRING_PTR: {
+        // std::string* is ragard as attr in PHI,
+        // but it is regard as output in Fluid.
+        auto it = ctx.outputs.find(attr_names[i]);
+        if (it == ctx.outputs.end() || it->second.empty()) {
+          phi_kernel_context->EmplaceBackAttr(nullptr);
+          continue;
+        }
+        auto* var = (it->second)[0];
+        if (var->template IsType<std::string>()) {
+          phi_kernel_context->EmplaceBackAttr(
+              var->template GetMutable<std::string>());
+        }
+      } break;
       case phi::AttributeType::SCALARS: {
         PADDLE_ENFORCE_NE(
             attr_iter,

paddle/fluid/framework/string_array.h

@@ -20,13 +20,14 @@ limitations under the License. */
 #include <string>
 #include <unordered_map>
 #include <vector>
+#include "paddle/phi/core/vocab.h"
 
 namespace paddle {
 namespace framework {
 
 using String = std::string;
 using Strings = std::vector<std::string>;
-using Vocab = std::unordered_map<std::wstring, std::int32_t>;
+using Vocab = phi::Vocab;
 
 // Convert the std::string type to the std::string type.
 bool ConvertStrToWstr(const std::string& src, std::wstring* res);

paddle/fluid/imperative/CMakeLists.txt

@@ -5,7 +5,7 @@ cc_library(
 cc_library(
   var_helper
   SRCS var_helper.cc
-  DEPS tensor selected_rows)
+  DEPS tensor selected_rows extended_tensor)
 if(WITH_XPU)
   cc_library(
     prepared_operator

paddle/fluid/operators/save_combine_op.cc

@@ -101,11 +101,3 @@ REGISTER_OPERATOR(save_combine,
                   ops::SaveCombineOp,
                   ops::SaveCombineOpProtoMaker,
                   ops::SaveCombineOpInferVarType);
-
-REGISTER_OP_CPU_KERNEL(
-    save_combine,
-    ops::SaveCombineOpKernel<phi::CPUContext, float>,
-    ops::SaveCombineOpKernel<phi::CPUContext, double>,
-    ops::SaveCombineOpKernel<phi::CPUContext, paddle::platform::bfloat16>,
-    ops::SaveCombineOpKernel<phi::CPUContext, int>,
-    ops::SaveCombineOpKernel<phi::CPUContext, int64_t>);

paddle/fluid/operators/save_load_combine_op_test.cc

@@ -20,8 +20,10 @@ limitations under the License. */
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/platform/bfloat16.h"
 #include "paddle/fluid/platform/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
 
-USE_CPU_ONLY_OP(save_combine);
+USE_OP_ITSELF(save_combine);
+PD_DECLARE_KERNEL(save_combine_tensor, CPU, ALL_LAYOUT);
 USE_CPU_ONLY_OP(load_combine);
 
 template <typename T, typename U>

paddle/fluid/pybind/eager_method.cc

@@ -54,6 +54,7 @@ typedef SSIZE_T ssize_t;
 #include "paddle/fluid/memory/allocation/mmap_allocator.h"
 #include "paddle/fluid/pybind/tensor_py.h"
 #include "paddle/phi/core/ddim.h"
+#include "paddle/phi/core/vocab.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace paddle {
@@ -1448,7 +1449,7 @@ static PyObject* tensor_method_set_vocab(TensorObject* self,
                                          PyObject* args,
                                          PyObject* kwargs) {
   EAGER_TRY
-  using Vocab = std::unordered_map<std::wstring, int>;
+  using Vocab = phi::Vocab;
   auto vocab = CastPyArg2Vocab(PyTuple_GET_ITEM(args, 0), 0);
   auto var_tensor = std::make_shared<egr::VariableCompatTensor>();
   *var_tensor->GetMutable<Vocab>() = vocab;
@@ -1479,7 +1480,7 @@ static PyObject* tensor_method_get_map_tensor(TensorObject* self,
       true,
       paddle::platform::errors::Fatal(
           "this method is only effective for VariableCompatTensor"));
-  using Vocab = std::unordered_map<std::wstring, int>;
+  using Vocab = phi::Vocab;
   auto* var_tensor =
       static_cast<const egr::VariableCompatTensor*>(self->tensor.impl().get());
   return ToPyObject(var_tensor->Get<Vocab>());

paddle/fluid/pybind/eager_utils.cc

@@ -514,11 +514,9 @@ paddle::framework::proto::VarType::Type CastPyArg2ProtoType(PyObject* obj,
   return dtype;
 }
 
-std::unordered_map<std::wstring, int> CastPyArg2Vocab(PyObject* obj,
-                                                      ssize_t arg_pos) {
+phi::Vocab CastPyArg2Vocab(PyObject* obj, ssize_t arg_pos) {
   if (PyDict_Check(obj)) {
-    return ::pybind11::handle(obj)
-        .cast<std::unordered_map<std::wstring, int>>();
+    return ::pybind11::handle(obj).cast<phi::Vocab>();
   } else {
     PADDLE_THROW(platform::errors::InvalidArgument(
         "argument (position %d) must be dict, but got %s",
@@ -807,7 +805,7 @@ PyObject* ToPyObject(
   return dict;
 }
 
-PyObject* ToPyObject(const std::unordered_map<std::wstring, int>& value) {
+PyObject* ToPyObject(const phi::Vocab& value) {
   PyObject* dict = PyDict_New();
   for (const auto& map_iter : value) {
     // Convert Key

paddle/fluid/pybind/eager_utils.h

@@ -29,6 +29,7 @@ typedef SSIZE_T ssize_t;
 #include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/core/vocab.h"
 #include "pybind11/pybind11.h"
 #include "pybind11/stl.h"
 namespace paddle {
@@ -73,8 +74,7 @@ std::vector<std::vector<size_t>> CastPyArg2VectorOfVectorOfSize_t(
     PyObject* obj, size_t arg_pos);
 framework::proto::VarType::Type CastPyArg2ProtoType(PyObject* obj,
                                                     ssize_t arg_pos);
-std::unordered_map<std::wstring, int> CastPyArg2Vocab(PyObject* obj,
-                                                      ssize_t arg_pos);
+phi::Vocab CastPyArg2Vocab(PyObject* obj, ssize_t arg_pos);
 std::vector<std::string> CastPyArg2Strings(PyObject* obj, ssize_t arg_pos);
 std::shared_ptr<jit::Function> CastPyArg2JitFunction(PyObject* obj,
                                                      ssize_t arg_pos);
@@ -113,7 +113,7 @@ PyObject* ToPyObject(const paddle::framework::proto::VarType& type);
 PyObject* ToPyObject(const void* value);
 PyObject* ToPyObject(
     const std::unordered_map<std::string, std::vector<std::string>>& value);
-PyObject* ToPyObject(const std::unordered_map<std::wstring, int>& value);
+PyObject* ToPyObject(const phi::Vocab& value);
 
 class PyTensorHook : public egr::TensorHook {
  public:

paddle/phi/CMakeLists.txt

@@ -41,7 +41,8 @@ set(PHI_DEPS
     sparse_coo_tensor
     string_tensor
     api_scalar
-    api_int_array)
+    api_int_array
+    extended_tensor)
 
 get_property(phi_kernels GLOBAL PROPERTY PHI_KERNELS)
 set(PHI_DEPS ${PHI_DEPS} ${phi_kernels})

paddle/phi/core/CMakeLists.txt

@@ -69,6 +69,11 @@ cc_library(
   SRCS tensor_array.cc
   DEPS dense_tensor tensor_base)
 
+cc_library(
+  extended_tensor
+  SRCS extended_tensor.cc
+  DEPS tensor_base)
+
 cc_library(
   meta_tensor
   SRCS meta_tensor.cc

paddle/phi/core/attribute.h

@@ -46,7 +46,8 @@ using Attribute = paddle::variant<bool,
                                   IntArray,
                                   DataType,
                                   DataLayout,
-                                  Place>;
+                                  Place,
+                                  std::string*>;
 
 using AttributeMap = paddle::flat_hash_map<std::string, Attribute>;
 

paddle/phi/core/extended_tensor.cc

@@ -0,0 +1,61 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/extended_tensor.h"
+
+namespace phi {
+
+int64_t ExtendedTensor::numel() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `numel` method."));
+}
+
+const DDim& ExtendedTensor::dims() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `dims` method."));
+}
+
+const Place& ExtendedTensor::place() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `place` method."));
+}
+
+DataType ExtendedTensor::dtype() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `dtype` method."));
+}
+
+DataLayout ExtendedTensor::layout() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `dtype` method."));
+}
+
+bool ExtendedTensor::valid() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `valid` method."));
+}
+
+bool ExtendedTensor::initialized() const {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `initialized` method."));
+}
+
+void* ExtendedTensor::AllocateFrom(Allocator* allocator,
+                                   DataType dtype,
+                                   size_t requested_size) {
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ExtendedTensor does not support `AllocateFrom` method."));
+}
+
+}  // namespace phi

paddle/phi/core/extended_tensor.h

@@ -0,0 +1,56 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/phi/core/allocator.h"
+#include "paddle/phi/core/tensor_base.h"
+#include "paddle/phi/core/tensor_meta.h"
+
+namespace phi {
+
+/// \brief The ExtendedTensor is a interface for custom designed class.
+/// If you want to pass some self-designed data as input/output to kernels,
+/// you can inherit from this class to store your self-designed data.
+class ExtendedTensor : public TensorBase,
+                       public TypeInfoTraits<TensorBase, ExtendedTensor> {
+ public:
+  ExtendedTensor() = default;
+  virtual ~ExtendedTensor() = default;
+
+ public:
+  /// \brief Returns the name of the class for type traits.
+  /// \return The name of the class.
+  static const char* name() { return "ExtendedTensor"; }
+
+  int64_t numel() const override;
+
+  const DDim& dims() const override;
+
+  const Place& place() const override;
+
+  DataType dtype() const override;
+
+  DataLayout layout() const override;
+
+  bool valid() const override;
+
+  bool initialized() const override;
+
+  void* AllocateFrom(Allocator* allocator,
+                     DataType dtype,
+                     size_t requested_size = 0) override;
+};
+
+}  // namespace phi

paddle/phi/core/kernel_context.cc

@@ -123,6 +123,8 @@ const AttrType& KernelContext::AttrAt(size_t idx) const {
   }
 }
 
+using string_ptr = std::string*;
+
 template const bool& KernelContext::AttrAt(size_t idx) const;
 template const int& KernelContext::AttrAt(size_t idx) const;
 template const int64_t& KernelContext::AttrAt(size_t idx) const;
@@ -142,5 +144,6 @@ template const IntArray& KernelContext::AttrAt(size_t idx) const;
 template const DataType& KernelContext::AttrAt(size_t idx) const;
 template const DataLayout& KernelContext::AttrAt(size_t idx) const;
 template const Place& KernelContext::AttrAt(size_t idx) const;
+template const string_ptr& KernelContext::AttrAt(size_t idx) const;
 
 }  // namespace phi

paddle/phi/core/kernel_factory.h

@@ -142,6 +142,7 @@ enum class AttributeType {
   DATA_TYPE,
   DATA_LAYOUT,
   PLACE,
+  STRING_PTR,
 };
 
 struct AttributeArgDef {

paddle/phi/core/kernel_registry.h

@@ -27,6 +27,7 @@
 #include "paddle/phi/core/kernel_utils.h"
 #include "paddle/phi/core/macros.h"
 #include "paddle/phi/core/type_defs.h"
+#include "paddle/phi/core/vocab.h"
 
 namespace phi {
 
@@ -100,6 +101,12 @@ struct KernelArgsParseFunctor<Return_ (*)(Args_...)> {
                               default_tensor_layout,
                               default_key.dtype(),
                               arg_type);
+      } else if (arg_type ==
+                 std::type_index(typeid(const std::vector<const Vocab*>&))) {
+        args_def->AppendInput(default_key.backend(),
+                              default_tensor_layout,
+                              default_key.dtype(),
+                              arg_type);
       } else if (arg_type == std::type_index(typeid(
                                  const std::vector<const SelectedRows*>&))) {
         args_def->AppendInput(default_key.backend(),
@@ -203,6 +210,8 @@ struct KernelArgsParseFunctor<Return_ (*)(Args_...)> {
         args_def->AppendAttribute(AttributeType::FLOAT64);
       } else if (arg_type == std::type_index(typeid(std::string))) {
         args_def->AppendAttribute(AttributeType::STRING);
+      } else if (arg_type == std::type_index(typeid(std::string*))) {
+        args_def->AppendAttribute(AttributeType::STRING_PTR);
       } else if (arg_type ==
                  std::type_index(typeid(const std::vector<bool>&))) {
         args_def->AppendAttribute(AttributeType::BOOLS);