From 112f16143b756cc55ebba3962ebf6bf154087183 Mon Sep 17 00:00:00 2001
From: Yiqun Liu <liuyiqun01@baidu.com>
Date: Thu, 11 Apr 2019 00:15:18 +0800
Subject: [PATCH] Add an option to enable the cache of expected kernel in train
 phase. (#16724)

* Add an option to enable the cache of expected kernel in train phase.
test=develop

* Change the default value of cache_expected_kernel to true.
---
 paddle/fluid/framework/details/build_strategy.cc        | 6 ++++++
 paddle/fluid/framework/details/build_strategy.h         | 1 +
 paddle/fluid/framework/ir/expected_kernel_cache_pass.cc | 2 +-
 paddle/fluid/pybind/pybind.cc                           | 4 ++++
 4 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/framework/details/build_strategy.cc b/paddle/fluid/framework/details/build_strategy.cc
index afe5078bf..20cfa7529 100644
--- a/paddle/fluid/framework/details/build_strategy.cc
+++ b/paddle/fluid/framework/details/build_strategy.cc
@@ -150,6 +150,11 @@ class ParallelExecutorPassBuilder : public ir::PassBuilder {
       AppendPass("runtime_context_cache_pass");
     }
 
+    if (strategy_.cache_expected_kernel_) {
+      VLOG(10) << "Add expected_kernel_cache_pass";
+      AppendPass("expected_kernel_cache_pass");
+    }
+
     AppendMultiDevPass(strategy_);
 
     if (strategy_.fuse_all_reduce_ops_) {
@@ -337,3 +342,4 @@ USE_PASS(fuse_adam_op_pass);
 USE_PASS(fuse_sgd_op_pass);
 USE_PASS(fuse_all_reduce_op_pass);
 USE_PASS(runtime_context_cache_pass);
+USE_PASS(expected_kernel_cache_pass);
diff --git a/paddle/fluid/framework/details/build_strategy.h b/paddle/fluid/framework/details/build_strategy.h
index 121d4a27c..b1601cfbc 100644
--- a/paddle/fluid/framework/details/build_strategy.h
+++ b/paddle/fluid/framework/details/build_strategy.h
@@ -108,6 +108,7 @@ struct BuildStrategy {
   bool remove_unnecessary_lock_{true};
 
   bool cache_runtime_context_{false};
+  bool cache_expected_kernel_{true};
 
   // NOTE:
   // Before you add new options, think if it's a general strategy that works
diff --git a/paddle/fluid/framework/ir/expected_kernel_cache_pass.cc b/paddle/fluid/framework/ir/expected_kernel_cache_pass.cc
index ee67af0af..4a99d4c1a 100644
--- a/paddle/fluid/framework/ir/expected_kernel_cache_pass.cc
+++ b/paddle/fluid/framework/ir/expected_kernel_cache_pass.cc
@@ -23,7 +23,7 @@ namespace ir {
 void ExpectedKernelCachePass::ApplyImpl(ir::Graph* graph) const {
   VLOG(3) << "Applies Expected Kernel Cache strategy.";
   for (const Node* n : graph->Nodes()) {
-    if (n->IsOp()) {
+    if (n->IsOp() && n->Op()) {
       n->Op()->SetAttr(kEnableCacheExpectedKernel, true);
     }
   }
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index f0ea6d9b0..a8a2a94d4 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -1366,6 +1366,10 @@ All parameter, weight, gradient are variables in Paddle.
           "cache_runtime_context",
           [](const BuildStrategy &self) { return self.cache_runtime_context_; },
           [](BuildStrategy &self, bool b) { self.cache_runtime_context_ = b; })
+      .def_property(
+          "cache_expected_kernel",
+          [](const BuildStrategy &self) { return self.cache_expected_kernel_; },
+          [](BuildStrategy &self, bool b) { self.cache_expected_kernel_ = b; })
       .def("_finalize_strategy_and_create_passes",
            [](BuildStrategy &self) -> std::shared_ptr<ir::PassBuilder> {
              return self.CreatePassesFromStrategy(true);
-- 
GitLab