[Cherry-pick2.3] Optimize dygraph performance part3 (#42256)

chenwhql · web-flow · commit 9495708ac40b · 2022-04-27T08:57:37.000+08:00
* Change small vector size (#42202) * change samll vector size * Update type_defs.h * Optimize dygraph InferShape perf (#42155) * init commit * remove two hash impl * fix bug * polish details * fix compile failed * fix compile failed * fix compile failed * add default kernel sig cache * fix get kernel arg defs error * remove kernel arg defs cache * fix origin op execute
diff --git a/paddle/fluid/framework/infershape_utils.cc b/paddle/fluid/framework/infershape_utils.cc
@@ -402,11 +402,11 @@ std::vector<phi::MetaTensor*> CompatInferMetaContext::MutableOutputBetween(
 CompatInferMetaContext BuildInferMetaContext(InferShapeContext* ctx,
                                              const std::string& op_type) {
   // 1. get kernel args
-  auto* arg_map_fn = phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_type);
+  auto* arg_map_fn = ctx->GetPhiArgumentMappingFn();
   InferShapeArgumentMappingContext arg_map_context(*ctx);
-  KernelSignature signature =
-      arg_map_fn ? (*arg_map_fn)(arg_map_context)
-                 : phi::DefaultKernelSignatureMap::Instance().Get(op_type);
+  phi::KernelSignature signature = arg_map_fn
+                                       ? (*arg_map_fn)(arg_map_context)
+                                       : *ctx->GetPhiDefaultKernelSignature();
   VLOG(3) << "BuildInferMetaContext: op kernel signature - " << signature;
 
   // 2. build infermeta context
diff --git a/paddle/fluid/framework/new_executor/new_executor_defs.cc b/paddle/fluid/framework/new_executor/new_executor_defs.cc
@@ -393,6 +393,16 @@ void InterpretercoreInferShapeContext::SetOutputsDim(
   SetDims(vars, dims);
 }
 
+const phi::ArgumentMappingFn*
+InterpretercoreInferShapeContext::GetPhiArgumentMappingFn() const {
+  return phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_.Type());
+}
+
+const phi::KernelSignature*
+InterpretercoreInferShapeContext::GetPhiDefaultKernelSignature() const {
+  return &phi::DefaultKernelSignatureMap::Instance().Get(op_.Type());
+}
+
 void InterpretercoreInferShapeContext::SetSkipLoD(bool skip) {
   can_skip_lod_ = skip;
 }
diff --git a/paddle/fluid/framework/new_executor/new_executor_defs.h b/paddle/fluid/framework/new_executor/new_executor_defs.h
@@ -111,6 +111,10 @@ class InterpretercoreInferShapeContext : public InferShapeContext {
   void SetOutputsDim(const std::string& name,
                      const std::vector<DDim>& dims) override;
 
+  const phi::ArgumentMappingFn* GetPhiArgumentMappingFn() const override;
+
+  const phi::KernelSignature* GetPhiDefaultKernelSignature() const override;
+
   void SetSkipLoD(bool skip);
 
  protected:
diff --git a/paddle/fluid/framework/op_desc.cc b/paddle/fluid/framework/op_desc.cc
@@ -271,6 +271,14 @@ class CompileTimeInferShapeContext : public InferShapeContext {
     SetDims(names, dims);
   }
 
+  const phi::ArgumentMappingFn *GetPhiArgumentMappingFn() const override {
+    return phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_.Type());
+  }
+
+  const phi::KernelSignature *GetPhiDefaultKernelSignature() const override {
+    return &phi::DefaultKernelSignatureMap::Instance().Get(op_.Type());
+  }
+
  protected:
   std::vector<proto::VarType::Type> GetVarTypes(
       const std::vector<std::string> &names) const {
diff --git a/paddle/fluid/framework/operator.cc b/paddle/fluid/framework/operator.cc
@@ -1007,6 +1007,14 @@ class RuntimeInferShapeContext : public InferShapeContext {
     SetDims(vars, dims);
   }
 
+  const phi::ArgumentMappingFn* GetPhiArgumentMappingFn() const override {
+    return phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_.Type());
+  }
+
+  const phi::KernelSignature* GetPhiDefaultKernelSignature() const override {
+    return &phi::DefaultKernelSignatureMap::Instance().Get(op_.Type());
+  }
+
  protected:
   DDim GetDim(Variable* var) const {
     PADDLE_ENFORCE_NOT_NULL(
@@ -1279,16 +1287,16 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
   phi::KernelKey pt_kernel_key;
   std::string pt_kernel_name;
   if (phi::KernelFactory::Instance().HasCompatiblePhiKernel(type_)) {
-    if (pt_kernel_signature_ == nullptr || pt_kernel_ == nullptr) {
-      pt_kernel_signature_.reset(
-          new KernelSignature(std::move(GetExpectedPhiKernelArgs(exe_ctx))));
-      VLOG(6) << *pt_kernel_signature_.get();
+    if (kernel_signature_ == nullptr || pt_kernel_ == nullptr) {
+      kernel_signature_.reset(new phi::KernelSignature(
+          std::move(GetExpectedPhiKernelArgs(exe_ctx))));
+      VLOG(6) << *kernel_signature_.get();
 
       kernel_type_.reset(
           new OpKernelType(std::move(InnerGetExpectedKernelType(exe_ctx))));
       dev_ctx = pool.Get(kernel_type_->place_);
 
-      pt_kernel_name = pt_kernel_signature_->name;
+      pt_kernel_name = kernel_signature_->name;
       pt_kernel_key = TransOpKernelTypeToPhiKernelKey(*kernel_type_.get());
       pt_kernel_.reset(
           new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
@@ -1303,7 +1311,7 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
                 << "` not found.";
       }
     } else {
-      pt_kernel_name = pt_kernel_signature_->name;
+      pt_kernel_name = kernel_signature_->name;
 // NOTE(Liu-xiandong): The register kernel used KP have library_type[KP],
 // But the default library_type is Plain, so we need to modify the
 // library_type here, otherwise it can't work.
@@ -1449,8 +1457,7 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
       phi::KernelContext pt_kernel_context;
       // Do data transform before building KernelContext
       // TODO(zhiqiu): support TransferInplaceVarsBack
-      PreparePhiData(exec_scope, *pt_kernel_, *pt_kernel_signature_,
-                     runtime_ctx);
+      PreparePhiData(exec_scope, *pt_kernel_, *kernel_signature_, runtime_ctx);
       BuildPhiKernelContext(*runtime_ctx, dev_ctx, &pt_kernel_context);
       (*pt_kernel_)(&pt_kernel_context);
     } else {
@@ -1545,14 +1552,14 @@ OpKernelType OperatorWithKernel::InnerGetExpectedKernelType(
 
 phi::KernelKey OperatorWithKernel::ChoosePhiKernel(
     const ExecutionContext& ctx) const {
-  pt_kernel_signature_.reset(
-      new KernelSignature(std::move(GetExpectedPhiKernelArgs(ctx))));
-  VLOG(6) << *pt_kernel_signature_.get();
+  kernel_signature_.reset(
+      new phi::KernelSignature(std::move(GetExpectedPhiKernelArgs(ctx))));
+  VLOG(6) << *kernel_signature_.get();
 
   kernel_type_.reset(
       new OpKernelType(std::move(InnerGetExpectedKernelType(ctx))));
 
-  auto pt_kernel_name = pt_kernel_signature_->name;
+  auto pt_kernel_name = kernel_signature_->name;
   auto pt_kernel_key = TransOpKernelTypeToPhiKernelKey(*kernel_type_.get());
   pt_kernel_.reset(new phi::Kernel(phi::KernelFactory::Instance().SelectKernel(
       pt_kernel_name, pt_kernel_key)));
@@ -2153,16 +2160,16 @@ OpKernelType OperatorWithKernel::GetKernelTypeForVar(
                       tensor.layout());
 }
 
-KernelSignature OperatorWithKernel::GetExpectedPhiKernelArgs(
+phi::KernelSignature OperatorWithKernel::GetExpectedPhiKernelArgs(
     const ExecutionContext& ctx) const {
   ExecutionArgumentMappingContext arg_mapping_ctx(ctx);
   if (arg_map_fn_ == nullptr) {
     auto* arg_map_fn = phi::OpUtilsMap::Instance().GetArgumentMappingFn(type_);
     if (arg_map_fn) {
       arg_map_fn_.reset(new phi::ArgumentMappingFn(*arg_map_fn));
     } else {
-      auto func =
-          [this](const phi::ArgumentMappingContext& ctx) -> KernelSignature {
+      auto func = [this](
+          const phi::ArgumentMappingContext& ctx) -> phi::KernelSignature {
         return phi::DefaultKernelSignatureMap::Instance().Get(type_);
       };
       arg_map_fn_.reset(new phi::ArgumentMappingFn(func));
@@ -2173,7 +2180,8 @@ KernelSignature OperatorWithKernel::GetExpectedPhiKernelArgs(
 
 Scope* OperatorWithKernel::PreparePhiData(
     const Scope& scope, const phi::Kernel& pt_kernel,
-    const KernelSignature& pt_kernel_signature, RuntimeContext* ctx) const {
+    const phi::KernelSignature& pt_kernel_signature,
+    RuntimeContext* ctx) const {
   const auto& input_names = pt_kernel_signature.input_names;
   auto input_defs = pt_kernel.args_def().input_defs();
   PADDLE_ENFORCE_EQ(input_names.size(), input_defs.size(),
@@ -2271,9 +2279,9 @@ void OperatorWithKernel::BuildPhiKernelContext(
     phi::KernelContext* pt_kernel_context) const {
   pt_kernel_context->SetDeviceContext(dev_ctx);
 
-  auto& input_names = pt_kernel_signature_->input_names;
-  auto& attr_names = pt_kernel_signature_->attr_names;
-  auto& output_names = pt_kernel_signature_->output_names;
+  auto& input_names = kernel_signature_->input_names;
+  auto& attr_names = kernel_signature_->attr_names;
+  auto& output_names = kernel_signature_->output_names;
 
   auto input_defs = pt_kernel_->args_def().input_defs();
   auto attr_defs = pt_kernel_->args_def().attribute_defs();
diff --git a/paddle/fluid/framework/operator.h b/paddle/fluid/framework/operator.h
@@ -632,7 +632,7 @@ class OperatorWithKernel : public OperatorBase {
                              phi::KernelContext* pt_kernel_context) const;
 
   phi::KernelSignature* PhiKernelSignature() const {
-    return pt_kernel_signature_.get();
+    return kernel_signature_.get();
   }
 
   phi::Kernel* PhiKernel() const { return pt_kernel_.get(); }
@@ -704,7 +704,7 @@ class OperatorWithKernel : public OperatorBase {
   // we may polish the implementation here
   mutable bool run_phi_kernel_ = false;
   mutable bool run_kp_kernel = false;
-  mutable std::unique_ptr<phi::KernelSignature> pt_kernel_signature_;
+  mutable std::unique_ptr<phi::KernelSignature> kernel_signature_;
   mutable std::unique_ptr<phi::Kernel> pt_kernel_;
   mutable std::unique_ptr<phi::ArgumentMappingFn> arg_map_fn_;
 };
diff --git a/paddle/fluid/framework/phi_utils.cc b/paddle/fluid/framework/phi_utils.cc
@@ -45,7 +45,7 @@ class KernelArgsNameMakerByOpProto : public KernelArgsNameMaker {
   const paddle::SmallVector<const char*>& GetOutputArgsNames() override;
   const paddle::SmallVector<const char*>& GetAttrsArgsNames() override;
 
-  KernelSignature GetKernelSignature();
+  phi::KernelSignature GetKernelSignature();
 
  private:
   DISABLE_COPY_AND_ASSIGN(KernelArgsNameMakerByOpProto);
@@ -221,10 +221,10 @@ KernelArgsNameMakerByOpProto::GetAttrsArgsNames() {
   return attr_names_;
 }
 
-KernelSignature KernelArgsNameMakerByOpProto::GetKernelSignature() {
-  return KernelSignature(phi::TransToPhiKernelName(op_proto_->type()).c_str(),
-                         GetInputArgsNames(), GetAttrsArgsNames(),
-                         GetOutputArgsNames());
+phi::KernelSignature KernelArgsNameMakerByOpProto::GetKernelSignature() {
+  return phi::KernelSignature(
+      phi::TransToPhiKernelName(op_proto_->type()).c_str(), GetInputArgsNames(),
+      GetAttrsArgsNames(), GetOutputArgsNames());
 }
 
 std::once_flag kernel_sig_map_init_flag;
diff --git a/paddle/fluid/framework/phi_utils.h b/paddle/fluid/framework/phi_utils.h
@@ -40,8 +40,6 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
-using KernelSignature = phi::KernelSignature;
-
 /* Kernel Key translate */
 
 OpKernelType TransPhiKernelKeyToOpKernelType(const phi::KernelKey& kernel_key);
diff --git a/paddle/fluid/framework/shape_inference.h b/paddle/fluid/framework/shape_inference.h
@@ -113,6 +113,10 @@ class InferShapeContext {
   virtual paddle::SmallVector<InferShapeVarPtr, phi::kOutputSmallVectorSize>
   GetOutputVarPtrs(const std::string &name) const = 0;
 
+  virtual const phi::ArgumentMappingFn *GetPhiArgumentMappingFn() const = 0;
+
+  virtual const phi::KernelSignature *GetPhiDefaultKernelSignature() const = 0;
+
  protected:
   virtual std::vector<DDim> GetRepeatedDims(const std::string &name) const = 0;
   virtual void SetRepeatedDims(const std::string &name,
diff --git a/paddle/fluid/imperative/infer_shape_context.h b/paddle/fluid/imperative/infer_shape_context.h
@@ -37,13 +37,17 @@ class DygraphInferShapeContext : public framework::InferShapeContext {
       const NameVarMap<VarType>* in, const NameVarMap<VarType>* out,
       const framework::AttributeMap* attr,
       const framework::AttributeMap* default_attr, const std::string op_type,
-      const framework::OpKernelType* op_kernel_type = nullptr)
+      const framework::OpKernelType* op_kernel_type = nullptr,
+      const phi::ArgumentMappingFn* arg_map_fn = nullptr,
+      const phi::KernelSignature* default_kernel_signature = nullptr)
       : var_map_in_(in),
         var_map_out_(out),
         attrs_(attr),
         default_attrs_(default_attr),
         op_type_(op_type),
-        op_kernel_type_(op_kernel_type) {}
+        op_kernel_type_(op_kernel_type),
+        arg_map_fn_(arg_map_fn),
+        default_kernel_signature_(default_kernel_signature) {}
 
   bool HasInput(const std::string& name) const override {
     // has only one input
@@ -377,6 +381,14 @@ class DygraphInferShapeContext : public framework::InferShapeContext {
         "SetLoDLevel function not support in dygraph mode"));
   }
 
+  const phi::ArgumentMappingFn* GetPhiArgumentMappingFn() const override {
+    return arg_map_fn_;
+  }
+
+  const phi::KernelSignature* GetPhiDefaultKernelSignature() const override {
+    return default_kernel_signature_;
+  }
+
  protected:
   DDim GetDim(framework::Variable* var) const {
     PADDLE_ENFORCE_NOT_NULL(var, platform::errors::PreconditionNotMet(
@@ -438,6 +450,9 @@ class DygraphInferShapeContext : public framework::InferShapeContext {
   const framework::AttributeMap* default_attrs_;
   const std::string op_type_;
   const framework::OpKernelType* op_kernel_type_;
+  // arg_map_fn_ and default_kernel_signature_ may be nullptr
+  const phi::ArgumentMappingFn* arg_map_fn_;
+  const phi::KernelSignature* default_kernel_signature_;
 };
 
 }  // namespace imperative
diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc
diff --git a/paddle/fluid/imperative/prepared_operator.h b/paddle/fluid/imperative/prepared_operator.h
diff --git a/paddle/phi/core/type_defs.h b/paddle/phi/core/type_defs.h
