Refactoring InferShape

paddle/framework/CMakeLists.txt

@@ -19,13 +19,14 @@ cc_test(scope_test SRCS scope_test.cc DEPS scope)
 proto_library(framework_proto SRCS framework.proto)
 
 cc_library(attribute SRCS attribute.cc DEPS framework_proto)
+cc_library(op_proto_maker SRCS op_proto_maker.cc DEPS framework_proto)
 cc_library(op_info SRCS op_info.cc DEPS attribute framework_proto)
 cc_library(operator SRCS operator.cc DEPS op_info device_context tensor scope)
-cc_test(operator_test SRCS operator_test.cc DEPS operator op_registry)
+cc_test(operator_test SRCS operator_test.cc DEPS operator op_registry op_info)
 
 cc_library(grad_op_builder SRCS grad_op_builder.cc DEPS operator)
-cc_library(op_registry SRCS op_registry.cc DEPS grad_op_builder)
-cc_test(op_registry_test SRCS op_registry_test.cc DEPS op_registry)
+cc_library(op_registry SRCS op_registry.cc DEPS grad_op_builder op_proto_maker)
+cc_test(op_registry_test SRCS op_registry_test.cc DEPS op_registry op_info)
 cc_test(grad_op_builder_test SRCS grad_op_builder_test.cc DEPS grad_op_builder op_registry add_op)
 
 py_proto_compile(framework_py_proto SRCS framework.proto)

paddle/framework/attribute.h

@@ -39,6 +39,21 @@ AttrType AttrTypeID();
 
 Attribute GetAttrValue(const OpDesc::Attr& attr_desc);
 
+class AttrReader {
+ public:
+  explicit AttrReader(const AttributeMap& attrs) : attrs_(attrs) {}
+
+  template <typename T>
+  inline const T& Get(const std::string& name) const {
+    PADDLE_ENFORCE(attrs_.count(name) != 0, "%s should be in AttributeMap",
+                   name);
+    return boost::get<T>(attrs_.at(name));
+  }
+
+ private:
+  const AttributeMap& attrs_;
+};
+
 // check whether a value(attribute) fit a certain limit
 template <typename T>
 class GreaterThanChecker {

paddle/framework/op_info.h

@@ -19,6 +19,7 @@
 #include <unordered_map>
 
 #include "paddle/framework/attribute.h"
+#include "paddle/framework/shape_inference.h"
 
 namespace paddle {
 namespace framework {
@@ -34,6 +35,7 @@ struct OpInfo {
   std::string grad_op_type_;
   OpProto* proto_;
   OpAttrChecker* checker_;
+  ShapeInferenceFn shapeInferFn_;
 
   bool HasOpProtoAndChecker() const {
     return proto_ != nullptr && checker_ != nullptr;

paddle/framework/op_proto_maker.cc

@@ -0,0 +1,59 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/framework/op_proto_maker.h"
+
+namespace paddle {
+namespace framework {
+
+void OpProtoAndCheckerMaker::Validate() {
+  validated_ = true;
+  CheckNoDuplicatedInOutAttrs();
+}
+
+OpProtoAndCheckerMaker::VariableBuilder OpProtoAndCheckerMaker::AddInput(
+    const std::string& name, const std::string& comment) {
+  auto* input = proto_->add_inputs();
+  input->set_name(name);
+  input->set_comment(comment);
+  return OpProtoAndCheckerMaker::VariableBuilder{input};
+}
+
+OpProtoAndCheckerMaker::VariableBuilder OpProtoAndCheckerMaker::AddOutput(
+    const std::string& name, const std::string& comment) {
+  auto* output = proto_->add_outputs();
+  output->set_name(name);
+  output->set_comment(comment);
+  return OpProtoAndCheckerMaker::VariableBuilder{output};
+}
+
+void OpProtoAndCheckerMaker::CheckNoDuplicatedInOutAttrs() {
+  std::unordered_set<std::string> names;
+  auto checker = [&](const std::string& name) {
+    PADDLE_ENFORCE(!names.count(name), "[%s] is duplicated", name);
+    names.insert(name);
+  };
+  for (auto& attr : proto_->attrs()) {
+    checker(attr.name());
+  }
+  for (auto& input : proto_->inputs()) {
+    checker(input.name());
+  }
+  for (auto& output : proto_->outputs()) {
+    checker(output.name());
+  }
+}
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/op_proto_maker.h

@@ -0,0 +1,105 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include "paddle/framework/attribute.h"
+#include "paddle/framework/framework.pb.h"
+#include "paddle/framework/shape_inference.h"
+
+namespace paddle {
+namespace framework {
+
+// this class not only make proto but also init attribute checkers.
+class OpProtoAndCheckerMaker {
+ public:
+  OpProtoAndCheckerMaker(OpProto* proto, OpAttrChecker* op_checker)
+      : proto_(proto), op_checker_(op_checker) {}
+
+  virtual ~OpProtoAndCheckerMaker() {
+    PADDLE_ENFORCE(validated_, "should call Validate after build");
+  }
+
+  void Validate();
+
+ protected:
+  struct VariableBuilder {
+    OpProto::Var* var_;
+
+    VariableBuilder& AsDuplicable() {
+      var_->set_duplicable(true);
+      return *this;
+    }
+
+    VariableBuilder& AsIntermediate() {
+      var_->set_intermediate(true);
+      return *this;
+    }
+
+    VariableBuilder& NotInGradient() {
+      var_->set_not_in_gradient(true);
+      return *this;
+    }
+  };
+
+  VariableBuilder AddInput(const std::string& name, const std::string& comment);
+
+  VariableBuilder AddOutput(const std::string& name,
+                            const std::string& comment);
+
+  template <typename T>
+  TypedAttrChecker<T>& AddAttr(const std::string& name,
+                               const std::string& comment,
+                               bool generated = false) {
+    auto* attr = proto_->add_attrs();
+    attr->set_name(name);
+    attr->set_comment(comment);
+    attr->set_generated(generated);
+    attr->set_type(AttrTypeID<T>());
+    return op_checker_->AddAttrChecker<T>(name);
+  }
+
+  void AddComment(const std::string& comment) { proto_->set_comment(comment); }
+
+  void SetShapeInferenceFn(ShapeInferenceFn fn) { shape_infer_fn_ = fn; }
+
+  void SetGradShapeInferenceFn(ShapeInferenceFn fn) {
+    grad_shape_infer_fn_ = fn;
+  }
+
+ public:
+  const ShapeInferenceFn GetShapeInferenceFn() const { return shape_infer_fn_; }
+
+  const ShapeInferenceFn GetGradShapeInferenceFn() const {
+    return grad_shape_infer_fn_;
+  }
+
+ private:
+  void CheckNoDuplicatedInOutAttrs();
+
+  OpProto* proto_;
+  OpAttrChecker* op_checker_;
+  ShapeInferenceFn shape_infer_fn_{nullptr};
+  ShapeInferenceFn grad_shape_infer_fn_{nullptr};
+  bool validated_{false};
+};
+
+class NOPMaker : public OpProtoAndCheckerMaker {
+ public:
+  NOPMaker(framework::OpProto* proto, framework::OpAttrChecker* op_checker)
+      : OpProtoAndCheckerMaker(proto, op_checker) {}
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/op_registry.h

@@ -24,8 +24,10 @@ limitations under the License. */
 #include "paddle/framework/framework.pb.h"
 #include "paddle/framework/grad_op_builder.h"
 #include "paddle/framework/op_info.h"
+#include "paddle/framework/op_proto_maker.h"
 #include "paddle/framework/operator.h"
 #include "paddle/framework/scope.h"
+#include "paddle/framework/shape_inference_impl.h"
 
 namespace paddle {
 namespace framework {
@@ -34,10 +36,12 @@ class OpRegistry {
  public:
   template <typename OpType, typename ProtoMakerType, typename GradOpType>
   static void RegisterOp(const std::string& op_type,
-                         const std::string& grad_op_type) {
+                         const std::string& grad_op_type,
+                         const ShapeInferenceFn fn) {
     PADDLE_ENFORCE(!OpInfoMap::Instance().Has(op_type),
                    "'%s' is registered more than once.", op_type);
     OpInfo op_info;
+    ShapeInferenceFn grad_op_inferer = nullptr;
     op_info.creator_ = [](
         const std::string& type, const VariableNameMap& inputs,
         const VariableNameMap& outputs, const AttributeMap& attrs) {
@@ -51,18 +55,21 @@ class OpRegistry {
       auto maker = ProtoMakerType(op_info.proto_, op_info.checker_);
       maker.Validate();
       op_info.proto_->set_type(op_type);
+      op_info.shapeInferFn_ = maker.GetShapeInferenceFn();
+      grad_op_inferer = maker.GetGradShapeInferenceFn();
       PADDLE_ENFORCE(
           op_info.proto_->IsInitialized(),
           "Fail to initialize %s's OpProto, because %s is not initialized",
           op_type, op_info.proto_->InitializationErrorString());
     } else {
       op_info.proto_ = nullptr;
       op_info.checker_ = nullptr;
+      op_info.shapeInferFn_ = fn;
     }
     OpInfoMap::Instance().Insert(op_type, op_info);
     // register gradient op
     if (!grad_op_type.empty()) {
-      RegisterOp<GradOpType, NOPMaker, NOP>(grad_op_type, "");
+      RegisterOp<GradOpType, NOPMaker, NOP>(grad_op_type, "", grad_op_inferer);
     }
   }
 
@@ -74,6 +81,20 @@ class OpRegistry {
   static std::unique_ptr<OperatorBase> CreateOp(const OpDesc& op_desc);
 
   static std::unique_ptr<OperatorBase> CreateGradOp(const OperatorBase& op);
+
+  // compile time InferShape
+  static void InferShape(const OpDesc& op_desc,
+                         std::map<std::string, VarDesc*>& var_descs) {
+    auto& info = OpInfoMap::Instance().Get(op_desc.type());
+    auto op = OpRegistry::CreateOp(op_desc);
+    info.shapeInferFn_(CompileTimeInferShapeContext(op, var_descs));
+  }
+
+  // runtime InferShape
+  static void InferShape(const OperatorBase& op, const Scope& scope) {
+    auto& info = OpInfoMap::Instance().Get(op.Type());
+    info.shapeInferFn_(RunTimeInferShapeContext(op, scope));
+  }
 };
 
 class Registrar {
@@ -94,8 +115,8 @@ class OpRegistrar : public Registrar {
  public:
   explicit OpRegistrar(const char* op_type) { OpRegistrar(op_type, ""); }
   OpRegistrar(const char* op_type, const char* grad_op_type) {
-    OpRegistry::RegisterOp<OpType, ProtoMakerType, GradOpType>(op_type,
-                                                               grad_op_type);
+    OpRegistry::RegisterOp<OpType, ProtoMakerType, GradOpType>(
+        op_type, grad_op_type, nullptr);
   }
 };
 

paddle/framework/op_registry_test.cc

@@ -174,4 +174,4 @@ TEST(OpRegistry, CustomChecker) {
   op->Run(scope, dev_ctx);
   int test_attr = op->Attr<int>("test_attr");
   ASSERT_EQ(test_attr, 4);
-}
+}

paddle/framework/operator.cc

@@ -13,8 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/framework/operator.h"
-#include <algorithm>
-#include "paddle/framework/op_registry.h"
 
 namespace paddle {
 namespace framework {
@@ -228,43 +226,5 @@ std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
   return res;
 }
 
-void OpProtoAndCheckerMaker::Validate() {
-  validated_ = true;
-  CheckNoDuplicatedInOutAttrs();
-}
-
-OpProtoAndCheckerMaker::VariableBuilder OpProtoAndCheckerMaker::AddInput(
-    const std::string& name, const std::string& comment) {
-  auto* input = proto_->add_inputs();
-  input->set_name(name);
-  input->set_comment(comment);
-  return OpProtoAndCheckerMaker::VariableBuilder{input};
-}
-
-OpProtoAndCheckerMaker::VariableBuilder OpProtoAndCheckerMaker::AddOutput(
-    const std::string& name, const std::string& comment) {
-  auto* output = proto_->add_outputs();
-  output->set_name(name);
-  output->set_comment(comment);
-  return OpProtoAndCheckerMaker::VariableBuilder{output};
-}
-
-void OpProtoAndCheckerMaker::CheckNoDuplicatedInOutAttrs() {
-  std::unordered_set<std::string> names;
-  auto checker = [&](const std::string& name) {
-    PADDLE_ENFORCE(!names.count(name), "[%s] is duplicated", name);
-    names.insert(name);
-  };
-  for (auto& attr : proto_->attrs()) {
-    checker(attr.name());
-  }
-  for (auto& input : proto_->inputs()) {
-    checker(input.name());
-  }
-  for (auto& output : proto_->outputs()) {
-    checker(output.name());
-  }
-}
-
 }  // namespace framework
 }  // namespace paddle