[Pten] Support inplace and intermediate in C++ API

.gitignore

@@ -2,16 +2,17 @@ paddle/fluid/operators/distributed/send_recv.proto
 paddle/fluid/API.spec
 paddle/fluid/API_DEV.spec
 paddle/fluid/API_PR.spec
+paddle/fluid/eager/api/generated/*
 paddle/fluid/op_use_default_grad_maker_DEV.spec
 paddle/fluid/op_use_default_grad_maker_PR.spec
+paddle/pten/api/backward/backward_api.h
 paddle/pten/api/include/api.h
 paddle/pten/api/lib/api.cc
-paddle/pten/api/backward/backward_api.h
+paddle/pten/api/lib/dygraph_api.*
 paddle/pten/api/lib/backward_api.cc
+paddle/pten/extension.h
 paddle/pten/include/*
 paddle/pten/infermeta/generated.*
-paddle/pten/extension.h
-paddle/fluid/eager/api/generated/*
 
 *.DS_Store
 *.vs

paddle/pten/api/lib/CMakeLists.txt

@@ -23,8 +23,12 @@ set(api_gen_file ${CMAKE_SOURCE_DIR}/python/paddle/utils/code_gen/api_gen.py)
 set(api_yaml_file ${CMAKE_SOURCE_DIR}/python/paddle/utils/code_gen/api.yaml)
 set(api_header_file ${CMAKE_SOURCE_DIR}/paddle/pten/api/include/api.h)
 set(api_source_file ${CMAKE_SOURCE_DIR}/paddle/pten/api/lib/api.cc)
+set(dygraph_api_header_file ${CMAKE_SOURCE_DIR}/paddle/pten/api/lib/dygraph_api.h)
+set(dygraph_api_source_file ${CMAKE_SOURCE_DIR}/paddle/pten/api/lib/dygraph_api.cc)
 set(api_header_file_tmp ${api_header_file}.tmp)
 set(api_source_file_tmp ${api_source_file}.tmp)
+set(dygraph_api_header_file_tmp ${dygraph_api_header_file}.tmp)
+set(dygraph_api_source_file_tmp ${dygraph_api_source_file}.tmp)
 
 # backward api file
 set(bw_api_gen_file ${CMAKE_SOURCE_DIR}/python/paddle/utils/code_gen/backward_api_gen.py)
@@ -46,14 +50,18 @@ endif()
 
 # generate forward api
 add_custom_command(
-  OUTPUT ${api_header_file} ${api_source_file}
+  OUTPUT ${api_header_file} ${api_source_file} ${dygraph_api_header_file} ${dygraph_api_source_file}
   COMMAND ${PYTHON_EXECUTABLE} -m pip install pyyaml
   COMMAND ${PYTHON_EXECUTABLE} ${api_gen_file}
                  --api_yaml_path ${api_yaml_file}
                  --api_header_path ${api_header_file_tmp}
                  --api_source_path ${api_source_file_tmp}
+                 --dygraph_api_header_path ${dygraph_api_header_file_tmp}
+                 --dygraph_api_source_path ${dygraph_api_source_file_tmp}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${api_header_file_tmp} ${api_header_file}
   COMMAND ${CMAKE_COMMAND} -E copy_if_different ${api_source_file_tmp} ${api_source_file}
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${dygraph_api_header_file_tmp} ${dygraph_api_header_file}
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${dygraph_api_source_file_tmp} ${dygraph_api_source_file}
   COMMENT "copy_if_different ${api_header_file} ${api_source_file}"
   DEPENDS ${api_yaml_file} ${api_gen_file} ${api_gen_base}
   VERBATIM)
@@ -85,5 +93,6 @@ cc_library(pten_data_transform SRCS data_transform.cc DEPS pten_tensor transfer_
 cc_library(manual_api SRCS manual_api.cc DEPS pten_tensor pten kernel_dispatch)
 cc_library(sparse_api SRCS sparse_api.cc DEPS pten_tensor pten kernel_dispatch pten_data_transform)
 cc_library(pten_function_api SRCS ${api_source_file} DEPS pten_tensor pten kernel_dispatch pten_data_transform)
+cc_library(pten_dygraph_api SRCS ${dygraph_api_source_file} DEPS pten_tensor pten kernel_dispatch pten_data_transform)
 cc_library(pten_bw_function_api SRCS ${bw_api_source_file} DEPS pten_tensor pten kernel_dispatch backward_infermeta pten_data_transform pten_function_api)
 cc_library(wrapped_infermeta SRCS ${wrapped_infermeta_source_file} DEPS pten)

paddle/pten/api/lib/api_utils.h

@@ -72,11 +72,14 @@ inline pten::MetaTensor MakeMetaTensor(const pten::SelectedRows& tensor) {
 /* ------------------ for output ----------------------- */
 
 inline pten::DenseTensor* SetKernelOutput(Backend backend, Tensor* out) {
-  auto dense_tensor = std::make_shared<pten::DenseTensor>(
-      pten::make_intrusive<SharedStorage>(pten::TransToPtenPlace(backend)),
-      pten::DenseTensorMeta());
-  out->set_impl(dense_tensor);
-  return dense_tensor.get();
+  if (!out->initialized()) {
+    auto dense_tensor = std::make_shared<pten::DenseTensor>(
+        pten::make_intrusive<SharedStorage>(pten::TransToPtenPlace(backend)),
+        pten::DenseTensorMeta());
+    out->set_impl(dense_tensor);
+    return dense_tensor.get();
+  }
+  return static_cast<pten::DenseTensor*>(out->impl().get());
 }
 
 inline std::vector<pten::DenseTensor*> SetKernelOutput(
@@ -96,9 +99,12 @@ inline std::vector<pten::DenseTensor*> SetKernelOutput(
 
 inline pten::SelectedRows* SetSelectedRowsKernelOutput(Backend backend,
                                                        Tensor* out) {
-  auto select_rows = std::make_shared<pten::SelectedRows>();
-  out->set_impl(select_rows);
-  return select_rows.get();
+  if (!out->initialized()) {
+    auto select_rows = std::make_shared<pten::SelectedRows>();
+    out->set_impl(select_rows);
+    return select_rows.get();
+  }
+  return static_cast<pten::SelectedRows*>(out->impl().get());
 }
 
 }  // namespace experimental

paddle/pten/api/lib/tensor.cc

@@ -252,10 +252,13 @@ Tensor::data<pten::dtype::bfloat16>() const;
 
 template <typename T>
 T *Tensor::data() {
-  PADDLE_THROW(pten::errors::Unimplemented(
-      "It is not currently supported to directly obtain the modifiable data "
-      "address through the tensor::data<T>() method, please use the "
-      "tensor::mutable_data<T>() method."));
+  if (is_dense_tensor()) {
+    return std::dynamic_pointer_cast<pten::DenseTensor>(impl_)->data<T>();
+  } else if (pten::SelectedRows::classof(impl_.get())) {
+    return std::dynamic_pointer_cast<pten::SelectedRows>(impl_)
+        ->mutable_value()
+        ->data<T>();
+  }
   return nullptr;
 }
 

paddle/pten/tests/api/test_reshape_api.cc

@@ -67,6 +67,25 @@ TEST(API, reshape) {
   ASSERT_EQ(value_equal, true);
 }
 
+TEST(API, reshape_) {
+  // 1. create tensor
+  auto x = paddle::experimental::full(
+      {3, 2, 2, 3}, 1.0, experimental::DataType::FLOAT32);
+
+  // 2. test API
+  paddle::experimental::Tensor out = paddle::experimental::reshape_(x, {12, 3});
+  // 3. check result
+  std::vector<int64_t> expect_shape = {12, 3};
+  ASSERT_EQ(out.shape()[0], expect_shape[0]);
+  ASSERT_EQ(out.shape()[1], expect_shape[1]);
+  ASSERT_EQ(out.numel(), 36);
+  ASSERT_EQ(out.is_cpu(), true);
+  ASSERT_EQ(out.type(), pten::DataType::FLOAT32);
+  ASSERT_EQ(out.layout(), pten::DataLayout::NCHW);
+  ASSERT_EQ(out.initialized(), true);
+  ASSERT_EQ(out.data<float>(), x.data<float>());
+}
+
 TEST(Tensor, old_reshape) {
   paddle::experimental::Tensor x(paddle::PlaceType::kCPU);
   x.reshape({3, 4});

paddle/pten/tests/api/test_scale_api.cc

@@ -62,7 +62,7 @@ TEST(API, scale_sr) {
       experimental::full({3, 4}, 1.0, pten::DataType::FLOAT32).impl());
   *(selected_rows->mutable_value()) = *dense_tensor;
   experimental::Tensor x(selected_rows);
-  const auto out = experimental::scale(x, 2.0, 1.0, true);
+  auto out = experimental::scale(x, 2.0, 1.0, true);
 
   ASSERT_EQ(out.dims().size(), 2);
   ASSERT_EQ(out.dims()[0], 3);

python/paddle/utils/code_gen/api.yaml

@@ -142,11 +142,12 @@
 
 - api : reshape
   args : (Tensor x, ScalarArray shape)
-  output : Tensor
+  output : Tensor(out)
   infer_meta :
     func : ReshapeInferMeta
   kernel :
     func : reshape
+  inplace : (x -> out)
 
 - api : scale
   args : (Tensor x, Scalar scale, float bias, bool bias_after_scale)

python/paddle/utils/code_gen/api_base.py

@@ -48,10 +48,14 @@ def __init__(self, api_item_yaml):
             self.support_selected_rows_kernel = False if len(self.kernel[
                 'func']) == 1 else True
             self.data_transform = self.parse_data_transform(api_item_yaml)
+            self.inplace_map = self.parse_inplace(api_item_yaml)
 
     def get_api_name(self, api_item_yaml):
         return api_item_yaml['api']
 
+    def get_api_func_name(self):
+        return self.api
+
     def parse_args(self, api_name, api_item_yaml):
         inputs, attrs, args_str = self.parse_input_and_attr(
             api_name, api_item_yaml['args'])
@@ -225,13 +229,37 @@ def parse_data_transform(self, api_item_yaml):
 
         return data_transform
 
+    def parse_inplace(self, api_item_yaml):
+        if 'inplace' in api_item_yaml:
+            inplace_map = {}
+            inplace_list = api_item_yaml['inplace'].split(',')
+            for item in inplace_list:
+                result = re.search(r"(?P<in>\w+)\s*->\s(?P<out>\w+)", item)
+                in_val = result.group('in')
+                out_val = result.group('out')
+                assert in_val in self.inputs['names'], \
+                    f"{self.api} : Inplace input error: the input var name('{in_val}') is not found in the input args of {self.api}."
+                assert out_val in self.outputs['names'], \
+                    f"{self.api} : Inplace output error: the output var name('{out_val}') is not found in the output args of {self.api}."
+
+                inplace_map[out_val] = in_val
+
+            return inplace_map
+        else:
+            return None
+
     # Override by child class
     def get_return_type(self, out_type_list):
         return None
 
     def gene_api_declaration(self):
         api_declaration = f"""
-PADDLE_API {self.outputs['return_type']} {self.api}({self.args_str['args_declare']});
+PADDLE_API {self.outputs['return_type']} {self.get_api_func_name()}({self.args_str['args_declare']});
+"""
+
+        if self.is_base_api and self.inplace_map is not None:
+            api_declaration = api_declaration + f"""
+PADDLE_API {self.outputs['return_type']} {self.get_api_func_name() + '_'}({self.args_str['args_declare']});
 """
 
         return api_declaration
@@ -527,14 +555,18 @@ def get_selected_rows_kernel_args(self, code_indent):
         return input_tensor_code, kernel_args[:-2], kernel_signature
 
     # Override by child class
-    def gene_output(self, output_type_list, set_out_func, code_indent):
+    def gene_output(self,
+                    output_type_list,
+                    set_out_func,
+                    code_indent,
+                    inplace_flag=False):
         return None, None, None
 
-    def gen_dense_tensor_kernel_code(self, code_indent):
+    def gen_dense_tensor_kernel_code(self, code_indent, inplace_flag=False):
         input_tensors, kernel_args, kernel_signature = self.get_kernel_args(
             code_indent)
         outputs_args, kernel_output_names, output_create = self.gene_output(
-            self.outputs['types'], 'SetKernelOutput', code_indent)
+            self.outputs['types'], 'SetKernelOutput', code_indent, inplace_flag)
         return f"""
 {code_indent}  auto kernel = pten::KernelFactory::Instance().SelectKernelOrThrowError(
 {code_indent}      "{self.kernel['func'][0]}", {{kernel_backend, kernel_layout, kernel_data_type}});
@@ -552,11 +584,12 @@ def gen_dense_tensor_kernel_code(self, code_indent):
 
 {code_indent}  return out;"""
 
-    def gen_selected_rows_kernel_code(self, code_indent):
+    def gen_selected_rows_kernel_code(self, code_indent, inplace_flag=False):
         input_tensors, kernel_args, kernel_signature = self.get_selected_rows_kernel_args(
             code_indent)
         outputs_args, kernel_output_names, output_create = self.gene_output(
-            self.outputs['types'], 'SetSelectedRowsKernelOutput', code_indent)
+            self.outputs['types'], 'SetSelectedRowsKernelOutput', code_indent,
+            inplace_flag)
         return f"""
 {code_indent}  auto kernel = pten::KernelFactory::Instance().SelectKernelOrThrowError(
 {code_indent}      "{self.kernel['func'][1]}", {{kernel_backend, kernel_layout, kernel_data_type}});
@@ -574,32 +607,38 @@ def gen_selected_rows_kernel_code(self, code_indent):
 
 {code_indent}  return out;"""
 
-    def gene_api_code(self):
-        if self.is_base_api:
-            api_code = f"""
-PADDLE_API {self.outputs['return_type']} {self.api}({self.args_str["args_define"]}) {{
+    def gene_base_api_code(self, inplace_flag=False):
+        api_func_name = self.get_api_func_name() + ('_' if inplace_flag else '')
+        api_code = f"""
+PADDLE_API {self.outputs['return_type']} {api_func_name}({self.args_str["args_define"]}) {{
 {self.gene_kernel_select()}
 """
 
-            if self.support_selected_rows_kernel:
-                code_indent = '  '
-                api_code = api_code + f"""
+        if self.support_selected_rows_kernel:
+            code_indent = '  '
+            return api_code + f"""
   if(kernel_type == KernelType::DENSE_TENSOR_KENREL){{
-{self.gen_dense_tensor_kernel_code(code_indent)}
+{self.gen_dense_tensor_kernel_code(code_indent, inplace_flag)}
   }} else {{
-{self.gen_selected_rows_kernel_code(code_indent)}      
+{self.gen_selected_rows_kernel_code(code_indent, inplace_flag)}
   }}
 }}
 """
 
-                return api_code
-            else:
-                code_indent = ''
-                return api_code + self.gen_dense_tensor_kernel_code(
-                    code_indent) + """
+        else:
+            code_indent = ''
+            return api_code + self.gen_dense_tensor_kernel_code(
+                code_indent, inplace_flag) + """
 }
 """
 
+    def gene_api_code(self):
+        if self.is_base_api:
+            api_code = self.gene_base_api_code()
+            if self.inplace_map is not None:
+                api_code = api_code + self.gene_base_api_code(inplace_flag=True)
+            return api_code
+
         else:
             inveke_func_name = self.invoke.split('(')[0].strip()
             if inveke_func_name in self.attrs['names']: