【Pten】Auto-Generate InterMeta register

paddle/pten/core/infermeta_utils.cc

@@ -66,6 +66,19 @@ const MetaConfig& InferMetaContext::GetMetaConfig() const { return config_; }
 const MetaTensor& InferMetaContext::InputAt(size_t idx) const {
   return *inputs_.at(idx);
 }
+
+std::vector<MetaTensor> InferMetaContext::InputsBetween(size_t start,
+                                                        size_t end) const {
+  std::vector<MetaTensor> result;
+  result.reserve(end - start);
+
+  for (size_t i = start; i < end; ++i) {
+    result.emplace_back(*inputs_.at(i));
+  }
+
+  return result;
+}
+
 MetaTensor* InferMetaContext::MutableOutputAt(size_t idx) {
   return outputs_.at(idx).get();
 }

paddle/pten/core/infermeta_utils.h

@@ -17,6 +17,8 @@ limitations under the License. */
 #include <string>
 #include <utility>
 
+#include "paddle/pten/common/scalar.h"
+#include "paddle/pten/common/scalar_array.h"
 #include "paddle/pten/core/enforce.h"
 #include "paddle/pten/core/macros.h"
 #include "paddle/pten/core/meta_tensor.h"
@@ -46,6 +48,7 @@ class InferMetaContext {
 
   const MetaConfig& GetMetaConfig() const;
   const MetaTensor& InputAt(size_t idx) const;
+  std::vector<MetaTensor> InputsBetween(size_t start, size_t end) const;
   MetaTensor* MutableOutputAt(size_t idx);
 
   template <typename AttrType>
@@ -85,7 +88,8 @@ class InferMetaContext {
                     "InferMeta's Attributes should appear before Outputs.");   \
       attr_type arg = ctx->AttrAt<attr_type>(attr_idx);                        \
       InferMetaFnCallHelper<                                                   \
-          Tail...>::template Call<in_idx, attr_idx + 1, out_idx>(pargs...,     \
+          Tail...>::template Call<in_idx, attr_idx + 1, out_idx>(ctx,          \
+                                                                 pargs...,     \
                                                                  arg);         \
     }                                                                          \
   }
@@ -124,6 +128,35 @@ struct InferMetaFnImpl<Return (*)(Args...), infer_meta_fn> {
     }
   };
 
+  template <typename... Tail>
+  struct InferMetaFnCallHelper<const std::vector<MetaTensor>&, Tail...> {
+    template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
+    static void Call(InferMetaContext* ctx, PreviousArgs&... pargs) {
+      static_assert(attr_idx == 0,
+                    "InferMeta's Input should appear before Attributes.");
+      static_assert(out_idx == 0,
+                    "InferMeta's Input should appear before Outputs.");
+      const std::pair<int, int> range = ctx->InputRangeAt(in_idx);
+      std::vector<MetaTensor> arg =
+          ctx->InputsBetween(range.first, range.second);
+      InferMetaFnCallHelper<
+          Tail...>::template Call<in_idx + 1, attr_idx, out_idx>(ctx,
+                                                                 pargs...,
+                                                                 arg);
+    }
+  };
+
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(bool);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(int);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(int64_t);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const std::vector<int>&);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(
+      const std::vector<int64_t>&);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataType);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(DataLayout);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const Scalar&);
+  PT_SPECIALIZE_InferMetaFnCallHelper_FOR_ATTRIBUTE(const ScalarArray&);
+
   // TODO(chenweihang): support vector<MetaTensor> input later
 
   template <typename... Tail>
@@ -227,7 +260,6 @@ struct InferMetaFnRegistrar {
       "PT_REGISTER_INFER_META_FN must be called in global namespace.");       \
   static const ::pten::InferMetaFnRegistrar                                   \
       __registrar_arg_map_fn_for_##kernel_name_prefix(                        \
-          #kernel_name_prefix, PT_INFER_META(variadic_infer_meta_fn));        \
-  int TouchInferMetaFnSymbol_##op_type() { return 0; }
+          #kernel_name_prefix, PT_INFER_META(variadic_infer_meta_fn))
 
 }  // namespace pten

paddle/pten/infermeta/nullary.cc

@@ -16,10 +16,10 @@ limitations under the License. */
 
 namespace pten {
 
-void CreateInferMeta(const std::vector<int64_t>& shape,
-                     DataType dtype,
-                     DataLayout layout,
-                     MetaTensor* out) {
+void CreateInferMetaBase(const std::vector<int64_t>& shape,
+                         DataType dtype,
+                         DataLayout layout,
+                         MetaTensor* out) {
   auto out_dims = pten::framework::make_ddim(shape);
   out->set_dims(out_dims);
   out->set_dtype(dtype);
@@ -30,7 +30,7 @@ void CreateInferMeta(const ScalarArray& shape,
                      DataType dtype,
                      DataLayout layout,
                      MetaTensor* out) {
-  CreateInferMeta(shape.GetData(), dtype, layout, out);
+  CreateInferMetaBase(shape.GetData(), dtype, layout, out);
 }
 
 }  // namespace pten

paddle/pten/infermeta/nullary.h

@@ -28,10 +28,10 @@ namespace pten {
 //   Because functions in this file not only can infer shape, but also need
 //   infer lod or other useful data.
 
-void CreateInferMeta(const std::vector<int64_t>& shape,
-                     DataType dtype,
-                     DataLayout layout,
-                     MetaTensor* out);
+void CreateInferMetaBase(const std::vector<int64_t>& shape,
+                         DataType dtype,
+                         DataLayout layout,
+                         MetaTensor* out);
 
 void CreateInferMeta(const ScalarArray& shape,
                      DataType dtype,

paddle/pten/infermeta/unary.cc

@@ -242,14 +242,14 @@ void SumInferMeta(const MetaTensor& x,
                   DataType dtype,
                   bool keep_dim,
                   MetaTensor* out) {
-  ReduceInferMeta(x, axis, keep_dim, dtype, std::move(out));
+  ReduceInferMetaBase(x, axis, keep_dim, dtype, out);
 }
 
-void ReduceInferMeta(const MetaTensor& x,
-                     const std::vector<int64_t>& axis,
-                     bool keep_dim,
-                     DataType dtype,
-                     MetaTensor* out) {
+void ReduceInferMetaBase(const MetaTensor& x,
+                         const std::vector<int64_t>& axis,
+                         bool keep_dim,
+                         DataType dtype,
+                         MetaTensor* out) {
   bool reduce_all = true;
   std::set<int64_t> dims_set(axis.begin(), axis.end());
   for (int64_t i = 0; i < x.dims().size(); ++i) {
@@ -304,7 +304,7 @@ void ReduceInferMeta(const MetaTensor& x,
                      const std::vector<int64_t>& axis,
                      bool keep_dim,
                      MetaTensor* out) {
-  ReduceInferMeta(x, axis, keep_dim, DataType::UNDEFINED, out);
+  ReduceInferMetaBase(x, axis, keep_dim, DataType::UNDEFINED, out);
 }
 
 void TransferLayoutInferMeta(const MetaTensor& x,
@@ -316,5 +316,3 @@ void TransferLayoutInferMeta(const MetaTensor& x,
 }
 
 }  // namespace pten
-
-PT_REGISTER_INFER_META_FN(sign, pten::UnchangedInferMeta);

paddle/pten/infermeta/unary.h

@@ -53,11 +53,11 @@ void ReshapeInferMeta(const MetaTensor& x,
                       const ScalarArray& shape,
                       MetaTensor* out);
 
-void ReduceInferMeta(const MetaTensor& x,
-                     const std::vector<int64_t>& axis,
-                     bool keep_dim,
-                     DataType dtype,
-                     MetaTensor* out);
+void ReduceInferMetaBase(const MetaTensor& x,
+                         const std::vector<int64_t>& axis,
+                         bool keep_dim,
+                         DataType dtype,
+                         MetaTensor* out);
 
 void ReduceInferMeta(const MetaTensor& x,
                      const std::vector<int64_t>& axis,

paddle/pten/kernels/math_kernel.h

@@ -156,7 +156,7 @@ DenseTensor Mean(const Context& dev_ctx,
                  bool keep_dim) {
   auto dense_out = pten::Empty<T, Context>(dev_ctx);
   MetaTensor meta_out(&dense_out);
-  ReduceInferMeta(x, axis, keep_dim, x.dtype(), &meta_out);
+  ReduceInferMetaBase(x, axis, keep_dim, x.dtype(), &meta_out);
   MeanKernel<T, Context>(dev_ctx, x, axis, keep_dim, &dense_out);
   return dense_out;
 }

python/paddle/utils/code_gen/api.yaml

@@ -161,6 +161,14 @@
   kernel :
     func : scale
 
+- api : sign
+  args : (const Tensor& x)
+  output : Tensor
+  infer_meta :
+    func : UnchangedInferMeta
+  kernel :
+    func : sign
+
 - api : subtract
   args : (const Tensor& x, const Tensor& y)
   output : Tensor
@@ -173,10 +181,10 @@
 - api : sum
   args : (const Tensor& x, const std::vector<int64_t>& axis={}, DataType dtype=DataType::UNDEFINED, bool keep_dim=false)
   output : Tensor
-  infer_meta : 
+  infer_meta :
     func : SumInferMeta
     param: [x, axis, dtype, keep_dim]
-  kernel : 
+  kernel :
     func : sum
     param : [x, axis, dtype, keep_dim]
     data_type : x

python/paddle/utils/code_gen/api_base.py

@@ -379,14 +379,7 @@ def get_kernel_args(self):
         input_infos = self.inputs['input_info']
         kernel_args_type_list = ['const platform::DeviceContext&']
 
-        input_tensor_code = ""
-        for input_name in input_names:
-            # set input code
-            input_tensor_code = input_tensor_code + f"""
-      auto {PREFIX_TENSOR_NAME}{input_name} = TensorToDenseTensor({input_name});"""
-
         attr_names = self.attrs['names']
-
         kernel_param = self.kernel['param']
         if kernel_param is None:
             kernel_param = input_names + attr_names
@@ -401,11 +394,11 @@ def get_kernel_args(self):
                 elif input_name in self.data_transform['support_trans_dtype']:
                     trans_flag = "{false, true}"
                 input_tensor_code = input_tensor_code + f"""
-      auto {PREFIX_TENSOR_NAME}{input_name} = PrepareData({input_name}, kernel.InputAt({i}), {trans_flag});"""
+  auto {PREFIX_TENSOR_NAME}{input_name} = PrepareData({input_name}, kernel.InputAt({i}), {trans_flag});"""
 
             else:
                 input_tensor_code = input_tensor_code + f"""
-      auto {PREFIX_TENSOR_NAME}{input_name} = TensorToDenseTensor({input_name});"""
+  auto {PREFIX_TENSOR_NAME}{input_name} = TensorToDenseTensor({input_name});"""
 
         kernel_args = "*dev_ctx, "
         for param in kernel_param:

python/paddle/utils/code_gen/api_gen.py

@@ -60,6 +60,14 @@ def gene_output(self, output_type_list):
 
         return kernel_output, output_names, output_create
 
+    def gene_infer_meta_register(self):
+        if self.is_base_api:
+            return f"""
+PT_REGISTER_INFER_META_FN({self.kernel['func']}, pten::{self.infer_meta['func']});"""
+
+        else:
+            return ''
+
 
 def header_include():
     return """
@@ -83,6 +91,7 @@ def source_include(header_file_path):
 #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/core/infermeta_utils.h"
 #include "paddle/pten/core/kernel_registry.h"
 #include "paddle/pten/infermeta/binary.h"
 #include "paddle/pten/infermeta/multiary.h"
@@ -127,15 +136,21 @@ def generate_api(api_yaml_path, header_file_path, source_file_path):
     source_file.write(source_include(include_header_file))
     source_file.write(namespace[0])
 
+    infer_meta_register_code = ''
+
     for api in apis:
         api_code = ForwardAPI(api)
         print(api_code.gene_api_declaration())
         header_file.write(api_code.gene_api_declaration())
         source_file.write(api_code.gene_api_code())
+        infer_meta_register_code = infer_meta_register_code + api_code.gene_infer_meta_register(
+        )
 
     header_file.write(namespace[1])
     source_file.write(namespace[1])
+
     source_file.write(api_register())
+    source_file.write(infer_meta_register_code)
 
     header_file.close()
     source_file.close()