From d9acc87e7ece15c4f10634d01cb91c58218d3794 Mon Sep 17 00:00:00 2001
From: pangyoki <pangyoki@126.com>
Date: Wed, 26 Jan 2022 15:18:45 +0800
Subject: [PATCH] add profile record (infer_shape, compute) for dygraph
 (#39023)

* add profile record for dygraph

* add op type in record

* fix little bug

* solve conflict
---
 paddle/fluid/imperative/prepared_operator.cc | 47 ++++++++++++++------
 paddle/fluid/pybind/op_function_generator.cc |  6 ++-
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc
index 2caa5319094fe..d9a21c9247b93 100644
--- a/paddle/fluid/imperative/prepared_operator.cc
+++ b/paddle/fluid/imperative/prepared_operator.cc
@@ -25,6 +25,7 @@
 #include "paddle/fluid/platform/device/xpu/xpu_op_list.h"
 #endif
 #include "paddle/fluid/platform/device/gpu/gpu_info.h"
+#include "paddle/fluid/platform/profiler.h"
 
 DECLARE_bool(check_nan_inf);
 DECLARE_bool(run_pten_kernel);
@@ -505,12 +506,21 @@ static void PreparedOpRunImpl(
   // TODO(zjl): remove scope in dygraph
   framework::Scope scope;
 
-  DygraphInferShapeContext<VarType> infer_shape_ctx(
-      &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
-  op.Info().infer_shape_(&infer_shape_ctx);
+  {
+    platform::RecordEvent record_event(op.Type() + " infer_shape",
+                                       platform::EventRole::kInnerOp);
+    DygraphInferShapeContext<VarType> infer_shape_ctx(
+        &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
+    op.Info().infer_shape_(&infer_shape_ctx);
+  }
+
+  {
+    platform::RecordEvent record_event(op.Type() + " compute",
+                                       platform::EventRole::kInnerOp);
 
-  func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
-                                        attrs, default_attrs));
+    func(DygraphExecutionContext<VarType>(op, scope, *dev_ctx, ctx, ins, outs,
+                                          attrs, default_attrs));
+  }
 
   if (FLAGS_check_nan_inf) {
     framework::details::CheckOpHasNanOrInfInDygraph<VarType>(
@@ -551,18 +561,27 @@ static void PreparedOpRunPtImpl(
     const NameVarMap<VarType>& ins, const NameVarMap<VarType>& outs,
     const framework::AttributeMap& attrs,
     const framework::AttributeMap& default_attrs) {
-  DygraphInferShapeContext<VarType> infer_shape_ctx(
-      &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
-  op.Info().infer_shape_(&infer_shape_ctx);
+  {
+    platform::RecordEvent record_event(op.Type() + " infer_shape",
+                                       platform::EventRole::kInnerOp);
+    DygraphInferShapeContext<VarType> infer_shape_ctx(
+        &ins, &outs, &attrs, &default_attrs, op.Type(), &kernel_type);
+    op.Info().infer_shape_(&infer_shape_ctx);
+  }
+
+  {
+    platform::RecordEvent record_event(op.Type() + " compute",
+                                       platform::EventRole::kInnerOp);
 
-  PreparePtenData<VarType>(pt_kernel, pt_kernel_signature, ins);
+    PreparePtenData<VarType>(pt_kernel, pt_kernel_signature, ins);
 
-  pten::KernelContext pt_kernel_context;
-  BuildDygraphPtenKernelContext<VarType>(pt_kernel_signature, pt_kernel, ins,
-                                         outs, attrs, default_attrs, dev_ctx,
-                                         &pt_kernel_context);
+    pten::KernelContext pt_kernel_context;
+    BuildDygraphPtenKernelContext<VarType>(pt_kernel_signature, pt_kernel, ins,
+                                           outs, attrs, default_attrs, dev_ctx,
+                                           &pt_kernel_context);
 
-  pt_kernel(&pt_kernel_context);
+    pt_kernel(&pt_kernel_context);
+  }
 
   if (FLAGS_benchmark) {
     dev_ctx->Wait();
diff --git a/paddle/fluid/pybind/op_function_generator.cc b/paddle/fluid/pybind/op_function_generator.cc
index 5587952facc53..957c0b0ee6d1d 100644
--- a/paddle/fluid/pybind/op_function_generator.cc
+++ b/paddle/fluid/pybind/op_function_generator.cc
@@ -123,6 +123,7 @@ static PyObject * %s(PyObject *self, PyObject *args, PyObject *kwargs)
   PyThreadState *tstate = nullptr;
   try
   {
+    platform::RecordEvent op_type_record_event("%s pybind_imperative_func");
     %s
     framework::AttributeMap attrs;
     ConstructAttrMapFromPyArgs("%s", args, %d, PyTuple_GET_SIZE(args) , attrs);
@@ -371,8 +372,8 @@ std::string GenerateOpFunctionsBody(
 
   // generate op funtcion body
   auto op_function_str = paddle::string::Sprintf(
-      OP_FUNCTION_TEMPLATE, func_name, ins_cast_str, op_type, input_args_num,
-      inplace_strategy_str, outs_initializer, ins_initializer,
+      OP_FUNCTION_TEMPLATE, func_name, op_type, ins_cast_str, op_type,
+      input_args_num, inplace_strategy_str, outs_initializer, ins_initializer,
       ins_initializer_with_null + outs_initializer_with_null +
           view_strategy_str,
       op_type, inplace_mapping_str, return_str);
@@ -461,6 +462,7 @@ int main(int argc, char* argv[]) {
 #endif
 
   std::vector<std::string> headers{"\"paddle/fluid/imperative/tracer.h\"",
+                                   "\"paddle/fluid/platform/profiler.h\"",
                                    "\"pybind11/detail/common.h\"",
                                    "<Python.h>"};