From d350cdbc973a81d3a1e1b3cd90aaaddf2bdfac2f Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 9 Oct 2017 19:59:38 -0700
Subject: [PATCH] Polish CompileTime InferShape

1. InferShape should be a method for `OpDescBind` not `Operator`, since
  `OpDescBind` is a compile-time concept.
2. Pre-create operators and store them into a map. Make InferShape
   faster
---
 paddle/framework/op_desc.cc                   | 35 +++++++++++++++++++
 paddle/framework/op_desc.h                    |  2 ++
 paddle/pybind/protobuf.cc                     |  3 +-
 paddle/pybind/pybind.cc                       | 15 --------
 .../v2/framework/tests/test_infer_shape.py    |  6 ++--
 5 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/paddle/framework/op_desc.cc b/paddle/framework/op_desc.cc
index 02aa74a842..6ce453cc99 100644
--- a/paddle/framework/op_desc.cc
+++ b/paddle/framework/op_desc.cc
@@ -13,7 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/framework/op_desc.h"
+#include <functional>
+#include <unordered_map>
 #include "paddle/framework/block_desc.h"
+#include "paddle/framework/operator.h"
 
 namespace paddle {
 namespace framework {
@@ -184,5 +187,37 @@ void OpDescBind::Sync() {
     need_update_ = false;
   }
 }
+
+using InferShapeFuncMap =
+    std::unordered_map<std::string /*op_type*/,
+                       std::function<void(InferShapeContext *)>>;
+
+static InferShapeFuncMap &InferShapeFuncs() {
+  static InferShapeFuncMap *g_map = nullptr;
+  if (g_map == nullptr) {
+    g_map = new InferShapeFuncMap();
+    auto &info_map = OpInfoMap::Instance();
+    // all registered kernels
+    for (auto &pair : OperatorWithKernel::AllOpKernels()) {
+      auto &info = info_map.Get(pair.first);
+      auto op =
+          static_cast<OperatorWithKernel *>(info.Creator()("", {}, {}, {}));
+      g_map->insert(
+          {pair.first, [op](InferShapeContext *ctx) { op->InferShape(ctx); }});
+    }
+  }
+  return *g_map;
+}
+
+void OpDescBind::InferShape(const BlockDescBind &block) const {
+  auto &funcs = InferShapeFuncs();
+  auto it = funcs.find(this->Type());
+  if (it == funcs.end()) {
+    PADDLE_THROW("Operator %s has not been registered", this->Type());
+  }
+  CompileTimeInferShapeContext ctx(*this, block);
+  it->second(&ctx);
+}
+
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/framework/op_desc.h b/paddle/framework/op_desc.h
index d0c314771c..81c4225041 100644
--- a/paddle/framework/op_desc.h
+++ b/paddle/framework/op_desc.h
@@ -100,6 +100,8 @@ class OpDescBind {
     return &this->attrs_;
   }
 
+  void InferShape(const BlockDescBind &block) const;
+
  private:
   template <typename MapType>
   static std::vector<typename MapType::key_type> MapKeys(const MapType &map) {
diff --git a/paddle/pybind/protobuf.cc b/paddle/pybind/protobuf.cc
index 47bd7bc3bb..6333cc332e 100644
--- a/paddle/pybind/protobuf.cc
+++ b/paddle/pybind/protobuf.cc
@@ -196,7 +196,8 @@ void BindOpDesc(py::module &m) {
       .def("set_attr", &OpDescBind::SetAttr)
       .def("attr", &OpDescBind::GetAttr)
       .def("set_block_attr", &OpDescBind::SetBlockAttr)
-      .def("get_block_attr", &OpDescBind::GetBlockAttr);
+      .def("get_block_attr", &OpDescBind::GetBlockAttr)
+      .def("infer_shape", &OpDescBind::InferShape);
 }
 
 }  // namespace pybind
diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc
index 356c4986e2..0f6e3101e2 100644
--- a/paddle/pybind/pybind.cc
+++ b/paddle/pybind/pybind.cc
@@ -231,21 +231,6 @@ All parameter, weight, gradient are variables in Paddle.
                                    desc.InitializationErrorString());
                     return OpRegistry::CreateOp(desc);
                   })
-      .def_static("infer_shape",
-                  [](OpDescBind &op_desc, BlockDescBind &block) {
-                    auto op = OpRegistry::CreateOp(*op_desc.Proto());
-                    auto *op_with_kernel =
-                        dynamic_cast<OperatorWithKernel *>(op.get());
-                    if (op_with_kernel != nullptr) {
-                      auto ctx = CompileTimeInferShapeContext(op_desc, block);
-                      op_with_kernel->InferShape(&ctx);
-                    } else {
-                      PADDLE_THROW(
-                          "OP(%s) is not type of OperatorWithKernel, "
-                          "should not call this function",
-                          op_desc.Type());
-                    }
-                  })
       .def("backward",
            [](const OperatorBase &forwardOp,
               const std::unordered_set<std::string> &no_grad_vars) {
diff --git a/python/paddle/v2/framework/tests/test_infer_shape.py b/python/paddle/v2/framework/tests/test_infer_shape.py
index b38ec9c037..99562890fd 100644
--- a/python/paddle/v2/framework/tests/test_infer_shape.py
+++ b/python/paddle/v2/framework/tests/test_infer_shape.py
@@ -1,6 +1,6 @@
 import unittest
+
 import paddle.v2.framework.core as core
-from paddle.v2.framework.op import Operator
 
 
 class TestInferShape(unittest.TestCase):
@@ -26,7 +26,7 @@ class TestInferShape(unittest.TestCase):
         sum_op_desc.set_input("X", ["x1", "x2"])
         sum_op_desc.set_output("Out", ["out"])
 
-        core.Operator.infer_shape(sum_op_desc, block)
+        sum_op_desc.infer_shape(block)
         self.assertEqual(out.shape(), shape)
 
     def test_mul_op(self):
@@ -55,7 +55,7 @@ class TestInferShape(unittest.TestCase):
         mul_op_desc.set_attr("x_num_col_dims", 1)
         mul_op_desc.set_attr("y_num_col_dims", 1)
 
-        core.Operator.infer_shape(mul_op_desc, block)
+        mul_op_desc.infer_shape(block)
         self.assertEqual(out.shape(), [x_shape[0], y_shape[1]])
 
 
-- 
GitLab