From 4d413d02cc708ed472428527a35c0e036ddb8300 Mon Sep 17 00:00:00 2001
From: Chen Weihang <chenweihang@baidu.com>
Date: Thu, 20 Jan 2022 09:48:01 +0800
Subject: [PATCH] revert cached kernel context removing (#39055)

---
 paddle/fluid/imperative/layer.cc  | 2 ++
 paddle/fluid/imperative/op_base.h | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/paddle/fluid/imperative/layer.cc b/paddle/fluid/imperative/layer.cc
index d2c63d5b21..f47b024973 100644
--- a/paddle/fluid/imperative/layer.cc
+++ b/paddle/fluid/imperative/layer.cc
@@ -409,6 +409,8 @@ void VarBase::_CopyGradientFrom(const VarBase& src) {
   }
 }
 
+pten::KernelContext OpBase::pt_kernel_context_;
+
 void OpBase::SetType(const std::string& type) {
   op_ = framework::OpRegistry::CreateOp(type, {}, {}, {}, false);
 }
diff --git a/paddle/fluid/imperative/op_base.h b/paddle/fluid/imperative/op_base.h
index 3d08476055..8d27e4f42a 100644
--- a/paddle/fluid/imperative/op_base.h
+++ b/paddle/fluid/imperative/op_base.h
@@ -210,6 +210,9 @@ class OpBase {
   std::unique_ptr<framework::OperatorBase> op_;
   platform::Place place_;
   size_t id_{-1UL};
+  // In order to reduce the compatibility phase
+  // performance overhead, temporarily cache KernelContext
+  static pten::KernelContext pt_kernel_context_;
   std::vector<std::shared_ptr<std::function<void()>>> void_function_post_hooks_;
 };
 
-- 
GitLab