Support custom implement for C++ API

paddle/pten/api/include/manual_api.h

@@ -30,11 +30,5 @@ namespace experimental {
 // TODO(chenweihang): Replace backend by place when place is ready
 PADDLE_API Tensor copy_to(const Tensor& x, Backend backend, bool blocking);
 
-// TODO(chentianyu03): Split API has extra logic to calculate the outputs size,
-// api_gen do not support
-PADDLE_API std::vector<Tensor> split(const Tensor& x,
-                                     const ScalarArray& num_or_sections,
-                                     const Scalar& axis);
-
 }  // namespace experimental
 }  // namespace paddle

paddle/pten/api/lib/CMakeLists.txt

@@ -82,8 +82,9 @@ add_custom_command(
   VERBATIM)
 
 cc_library(pten_data_transform SRCS data_transform.cc DEPS pten_tensor transfer_layout_kernel cast_kernel data_device_transform)
+cc_library(api_custom_impl SRCS api_custom_impl.cc DEPS pten_tensor pten kernel_dispatch pten_data_transform)
 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_bw_function_api SRCS ${bw_api_source_file} DEPS pten_tensor pten kernel_dispatch backward_infermeta pten_data_transform pten_function_api)
+cc_library(pten_function_api SRCS ${api_source_file} DEPS pten_tensor pten kernel_dispatch pten_data_transform api_custom_impl)
+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 api_custom_impl)
 cc_library(wrapped_infermeta SRCS ${wrapped_infermeta_source_file} DEPS pten)

paddle/pten/api/lib/api_custom_impl.cc

@@ -0,0 +1,102 @@
+/* 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/pten/api/lib/api_custom_impl.h"
+
+#include "paddle/pten/api/lib/api_registry.h"
+#include "paddle/pten/api/lib/api_utils.h"
+#include "paddle/pten/api/lib/data_transform.h"
+#include "paddle/pten/api/lib/kernel_dispatch.h"
+#include "paddle/pten/api/lib/utils/storage.h"
+#include "paddle/pten/common/backend.h"
+#include "paddle/pten/core/kernel_registry.h"
+#include "paddle/pten/core/meta_tensor.h"
+#include "paddle/pten/infermeta/binary.h"
+#include "paddle/pten/infermeta/multiary.h"
+#include "paddle/pten/infermeta/nullary.h"
+#include "paddle/pten/infermeta/unary.h"
+
+#include "glog/logging.h"
+
+namespace paddle {
+namespace experimental {
+
+PADDLE_API std::vector<Tensor> split_impl(const Tensor& x,
+                                          const ScalarArray& num_or_sections,
+                                          const Scalar& axis) {
+  Backend kernel_backend = Backend::UNDEFINED;
+  DataLayout kernel_layout = DataLayout::UNDEFINED;
+  DataType kernel_data_type = DataType::UNDEFINED;
+
+  if (kernel_backend == Backend::UNDEFINED ||
+      kernel_layout == DataLayout::UNDEFINED ||
+      kernel_data_type == DataType::UNDEFINED) {
+    auto kernel_key_set = ParseKernelKeyByInputArgs(x);
+    auto kernel_key = kernel_key_set.GetHigestPriorityKernelKey();
+    if (kernel_backend == Backend::UNDEFINED) {
+      kernel_backend = kernel_key.backend();
+    }
+    if (kernel_layout == DataLayout::UNDEFINED) {
+      kernel_layout = kernel_key.layout();
+    }
+    if (kernel_data_type == DataType::UNDEFINED) {
+      kernel_data_type = kernel_key.dtype();
+    }
+  }
+
+  auto kernel = pten::KernelFactory::Instance().SelectKernelOrThrowError(
+      "split", {kernel_backend, kernel_layout, kernel_data_type});
+  VLOG(6) << "split API kernel key: [" << kernel_backend << ", "
+          << kernel_layout << ", " << kernel_data_type << "]";
+  VLOG(6) << "split API kernel: " << kernel;
+
+  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
+
+  auto dense_x = PrepareData(x, kernel.InputAt(0), {});
+
+  // Calculate the number of out tensors
+  size_t out_number;
+  if (num_or_sections.GetData().size() == 1) {
+    out_number = num_or_sections.GetData()[0];
+  } else {
+    out_number = num_or_sections.GetData().size();
+  }
+
+  std::vector<Tensor> out;
+  auto dense_outs = SetKernelOutput(out_number, kernel_backend, &out);
+  std::vector<pten::MetaTensor> meta_outs;
+  for (size_t i = 0; i < out_number; ++i) {
+    meta_outs.push_back(dense_outs[i]);
+  }
+
+  pten::SplitInferMeta(
+      MakeMetaTensor(*dense_x), num_or_sections, axis, &meta_outs);
+
+  using kernel_signature = void (*)(const platform::DeviceContext&,
+                                    const pten::DenseTensor&,
+                                    const pten::ScalarArray&,
+                                    const pten::Scalar&,
+                                    std::vector<pten::DenseTensor*>&);
+  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
+  (*kernel_fn)(*dev_ctx,
+               *dense_x,
+               pten::ScalarArray(num_or_sections),
+               pten::Scalar(axis),
+               dense_outs);
+
+  return out;
+}
+
+}  // namespace experimental
+}  // namespace paddle

paddle/pten/api/lib/api_custom_impl.h

@@ -0,0 +1,29 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/pten/api/include/tensor.h"
+#include "paddle/pten/common/scalar.h"
+#include "paddle/pten/common/scalar_array.h"
+
+namespace paddle {
+namespace experimental {
+
+PADDLE_API std::vector<Tensor> split_impl(const Tensor& x,
+                                          const ScalarArray& num_or_sections,
+                                          const Scalar& axis);
+
+}  // namespace experimental
+}  // namespace paddle

paddle/pten/api/lib/manual_api.cc

@@ -78,71 +78,6 @@ PADDLE_API Tensor copy_to(const Tensor& x, Backend backend, bool blocking) {
   return out;
 }
 
-PADDLE_API std::vector<Tensor> split(const Tensor& x,
-                                     const ScalarArray& num_or_sections,
-                                     const Scalar& axis) {
-  Backend kernel_backend = Backend::UNDEFINED;
-  DataLayout kernel_layout = DataLayout::UNDEFINED;
-  DataType kernel_data_type = DataType::UNDEFINED;
-
-  if (kernel_backend == Backend::UNDEFINED ||
-      kernel_layout == DataLayout::UNDEFINED ||
-      kernel_data_type == DataType::UNDEFINED) {
-    auto kernel_key_set = ParseKernelKeyByInputArgs(x);
-    auto kernel_key = kernel_key_set.GetHigestPriorityKernelKey();
-    if (kernel_backend == Backend::UNDEFINED) {
-      kernel_backend = kernel_key.backend();
-    }
-    if (kernel_layout == DataLayout::UNDEFINED) {
-      kernel_layout = kernel_key.layout();
-    }
-    if (kernel_data_type == DataType::UNDEFINED) {
-      kernel_data_type = kernel_key.dtype();
-    }
-  }
-
-  auto kernel = pten::KernelFactory::Instance().SelectKernelOrThrowError(
-      "split", {kernel_backend, kernel_layout, kernel_data_type});
-  VLOG(6) << "split API kernel key: [" << kernel_backend << ", "
-          << kernel_layout << ", " << kernel_data_type << "]";
-  VLOG(6) << "split API kernel: " << kernel;
-
-  auto* dev_ctx = GetDeviceContextByBackend(kernel_backend);
-
-  auto dense_x = PrepareData(x, kernel.InputAt(0), {});
-
-  // Calculate the number of out tensors
-  size_t out_number;
-  if (num_or_sections.GetData().size() == 1) {
-    out_number = num_or_sections.GetData()[0];
-  } else {
-    out_number = num_or_sections.GetData().size();
-  }
-
-  std::vector<Tensor> out;
-  auto dense_outs = SetKernelOutput(out_number, kernel_backend, &out);
-  std::vector<pten::MetaTensor> meta_outs;
-  for (size_t i = 0; i < out_number; ++i) {
-    meta_outs.push_back(dense_outs[i]);
-  }
-
-  pten::SplitInferMeta(
-      MakeMetaTensor(*dense_x), num_or_sections, axis, &meta_outs);
-
-  using kernel_signature = void (*)(const platform::DeviceContext&,
-                                    const pten::DenseTensor&,
-                                    const pten::ScalarArray&,
-                                    const pten::Scalar&,
-                                    std::vector<pten::DenseTensor*>&);
-  auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
-  (*kernel_fn)(*dev_ctx,
-               *dense_x,
-               pten::ScalarArray(num_or_sections),
-               pten::Scalar(axis),
-               dense_outs);
-
-  return out;
-}
 }  // namespace experimental
 }  // namespace paddle
 

paddle/pten/tests/api/test_split_api.cc

@@ -17,7 +17,6 @@
 
 #include "paddle/pten/api/include/api.h"
 
-#include "paddle/pten/api/include/manual_api.h"
 #include "paddle/pten/api/lib/utils/allocator.h"
 #include "paddle/pten/core/dense_tensor.h"
 #include "paddle/pten/core/kernel_registry.h"

python/paddle/utils/code_gen/api.yaml

@@ -165,6 +165,11 @@
   kernel :
     func : sign
 
+- api : split
+  args : (const Tensor& x, const ScalarArray& num_or_sections, const Scalar& axis)
+  output : std::vector<Tensor>
+  invoke : split_impl(x, num_or_sections, axis)
+
 - api : subtract
   args : (const Tensor& x, const Tensor& y)
   output : Tensor

python/paddle/utils/code_gen/api_gen.py

@@ -78,6 +78,7 @@ def source_include(header_file_path):
 
 #include "glog/logging.h"
 
+#include "paddle/pten/api/lib/api_custom_impl.h"
 #include "paddle/pten/api/lib/api_registry.h"
 #include "paddle/pten/api/lib/api_utils.h"
 #include "paddle/pten/api/lib/data_transform.h"

python/paddle/utils/code_gen/backward_api_gen.py

@@ -124,6 +124,7 @@ def source_include(header_file_path):
 
 #include "glog/logging.h"
 
+#include "paddle/pten/api/lib/api_custom_impl.h"
 #include "paddle/pten/api/lib/api_registry.h"
 #include "paddle/pten/api/lib/api_utils.h"
 #include "paddle/pten/api/lib/data_transform.h"