Refactoring InferShape

paddle/framework/CMakeLists.txt

@@ -23,11 +23,11 @@ cc_library(op_proto_maker SRCS op_proto_maker.cc DEPS framework_proto attribute)
 cc_test(op_proto_maker_test SRCS op_proto_maker_test.cc DEPS op_proto_maker)
 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 op_proto_maker)
-cc_test(op_registry_test SRCS op_registry_test.cc DEPS op_registry)
+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/ddim.cc

@@ -195,6 +195,21 @@ std::vector<int64_t> vectorize(const DDim& ddim) {
   return result;
 }
 
+std::string debug_str(const DDim& ddim) {
+  auto ret = vectorize(ddim);
+  std::stringstream ss;
+  ss << "[";
+  auto size = ddim.size();
+  for (size_t i = 0; i < size; ++i) {
+    ss << ret[i];
+    if (i < size - 1) {
+      ss << ", ";
+    }
+  }
+  ss << "]";
+  return ss.str();
+}
+
 struct ProductVisitor : public boost::static_visitor<int64_t> {
   template <int D>
   int64_t operator()(const Dim<D>& dim) {

paddle/framework/ddim.h

@@ -94,6 +94,8 @@ void set(DDim& dim, int idx, int val);
 
 std::vector<int64_t> vectorize(const DDim& ddim);
 
+std::string debug_str(const DDim& ddim);
+
 int64_t product(const DDim& ddim);
 
 /**

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 {

paddle/framework/op_proto_maker.cc

@@ -1,4 +1,5 @@
 /* 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

paddle/framework/op_proto_maker.h

@@ -1,4 +1,5 @@
 /* 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
@@ -15,6 +16,7 @@ limitations under the License. */
 
 #include "paddle/framework/attribute.h"
 #include "paddle/framework/framework.pb.h"
+#include "paddle/framework/shape_inference.h"
 
 namespace paddle {
 namespace framework {

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 {
@@ -33,6 +31,15 @@ ExecutionContext::GetEigenDevice<platform::GPUPlace, Eigen::GpuDevice>() const {
 }
 #endif
 
+const Tensor* GetTensorFromVar(const Variable* var) {
+  if (var->IsType<LoDTensor>()) {
+    return &var->Get<LoDTensor>();
+  }
+  PADDLE_ENFORCE(var->IsType<Tensor>(),
+                 "The Input must be LoDTensor or Tensor.");
+  return &var->Get<Tensor>();
+}
+
 std::string OperatorBase::Input(const std::string& name) const {
   auto& ins = Inputs(name);
   PADDLE_ENFORCE_LE(ins.size(), 1UL,

paddle/framework/operator.h

@@ -24,6 +24,7 @@ limitations under the License. */
 #include "paddle/framework/framework.pb.h"
 #include "paddle/framework/lod_tensor.h"
 #include "paddle/framework/scope.h"
+#include "paddle/framework/shape_inference.h"
 #include "paddle/framework/tensor.h"
 #include "paddle/platform/device_context.h"
 #include "paddle/platform/place.h"
@@ -56,6 +57,8 @@ class OperatorBase;
 class InferShapeContext;
 class ExecutionContext;
 
+extern const Tensor* GetTensorFromVar(const Variable* var);
+
 /**
  * OperatorBase has the basic element that Net will call to do computation.
  * Only CreateOperator from OpRegistry will new Operator directly. User
@@ -262,15 +265,6 @@ class InferShapeContext {
     return res;
   }
 
-  const Tensor* GetTensorFromVar(const Variable* var) const {
-    if (var->IsType<LoDTensor>()) {
-      return &var->Get<LoDTensor>();
-    }
-    PADDLE_ENFORCE(var->IsType<Tensor>(),
-                   "The Input(%s) must be LoDTensor or Tensor.");
-    return &var->Get<Tensor>();
-  }
-
  private:
   const OperatorBase& op_;
   const Scope& scope_;
@@ -347,6 +341,154 @@ template <>
 std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
     const std::string& name) const;
 
+class BlockDesc {
+ public:
+  explicit BlockDesc(const std::map<std::string, VarDesc*>& var_descs)
+      : var_descs_(var_descs) {}
+  ~BlockDesc() {}
+
+  VarDesc* GetVar(const std::string& name) const {
+    PADDLE_ENFORCE(var_descs_.count(name) == 1, "%s must be in Block", name);
+    return var_descs_.at(name);
+  }
+
+  bool HasVar(const std::string& name) const {
+    return var_descs_.count(name) > 0;
+  }
+
+ private:
+  const std::map<std::string, VarDesc*>& var_descs_;
+};
+
+class CompileTimeInferShapeContext : public InferShapeContextBase {
+ public:
+  CompileTimeInferShapeContext(const OperatorBase& op,
+                               const BlockDesc& block_desc)
+      : op_(op), block_desc_(block_desc) {}
+
+  bool HasInput(const std::string& name) const {
+    return block_desc_.HasVar(name);
+  }
+
+  bool HasOutput(const std::string& name) const {
+    return block_desc_.HasVar(name);
+  }
+
+  DDim GetInputDim(const std::string& name) const {
+    return GetDim(op_.Input(name));
+  }
+
+  void SetInputDim(const std::string& name, const DDim& dim) const {
+    SetDim(op_.Input(name), dim);
+  }
+
+  DDim GetOutputDim(const std::string& name) const {
+    return GetDim(op_.Output(name));
+  }
+
+  void SetOutputDim(const std::string& name, const DDim& dim) const {
+    SetDim(op_.Output(name), dim);
+  }
+
+  AttrReader Attrs() const { return AttrReader(op_.Attrs()); }
+
+  const std::vector<std::string>& Inputs(const std::string& name) const {
+    return op_.Inputs(name);
+  }
+
+  const std::vector<std::string>& Outputs(const std::string& name) const {
+    return op_.Outputs(name);
+  }
+
+ private:
+  DDim GetDim(const std::string& name) const {
+    VarDesc* desc = block_desc_.GetVar(name);
+    std::vector<int64_t> dim;
+    int length = desc->lod_tensor().dims().size();
+    dim.reserve(length);
+    std::copy(desc->lod_tensor().dims().begin(),
+              desc->lod_tensor().dims().end(), std::back_inserter(dim));
+    return make_ddim(dim);
+  }
+
+  void SetDim(const std::string& name, const DDim& dim) const {
+    VarDesc* desc = block_desc_.GetVar(name);
+    auto tensor = desc->mutable_lod_tensor();
+    tensor->clear_dims();
+    for (int i = 0; i < dim.size(); ++i) {
+      tensor->add_dims(static_cast<int>(dim[i]));
+    }
+  }
+
+  const OperatorBase& op_;
+  const BlockDesc& block_desc_;
+};
+
+class RunTimeInferShapeContext : public InferShapeContextBase {
+ public:
+  RunTimeInferShapeContext(const OperatorBase& op, const Scope& scope)
+      : op_(op), scope_(scope) {}
+
+  bool HasInput(const std::string& name) const {
+    auto ipt = op_.Input(name);
+    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
+    return var != nullptr;
+  }
+
+  bool HasOutput(const std::string& name) const {
+    auto ipt = op_.Output(name);
+    auto* var = ipt == kEmptyVarName ? nullptr : scope_.FindVar(ipt);
+    return var != nullptr;
+  }
+
+  DDim GetInputDim(const std::string& name) const {
+    return GetDim(op_.Input(name));
+  }
+
+  void SetInputDim(const std::string& name, const DDim& dim) const {
+    SetDim(op_.Input(name), dim);
+  }
+
+  DDim GetOutputDim(const std::string& name) const {
+    return GetDim(op_.Output(name));
+  }
+
+  void SetOutputDim(const std::string& name, const DDim& dim) const {
+    SetDim(op_.Output(name), dim);
+  }
+
+  AttrReader Attrs() const { return AttrReader(op_.Attrs()); }
+
+  const std::vector<std::string>& Inputs(const std::string& name) const {
+    return op_.Inputs(name);
+  }
+
+  const std::vector<std::string>& Outputs(const std::string& name) const {
+    return op_.Outputs(name);
+  }
+
+ private:
+  Tensor* GetTensor(const std::string& name) const {
+    Tensor* t;
+    auto* var = scope_.FindVar(name);
+    if (!var->IsType<LoDTensor>() && !var->IsType<Tensor>()) {
+      t = var->GetMutable<LoDTensor>();
+    } else {
+      t = const_cast<Tensor*>(GetTensorFromVar(scope_.FindVar(name)));
+    }
+    return t;
+  }
+
+  DDim GetDim(const std::string& name) const { return GetTensor(name)->dims(); }
+
+  void SetDim(const std::string& name, const DDim& dim) const {
+    GetTensor(name)->Resize(dim);
+  }
+
+  const OperatorBase& op_;
+  const Scope& scope_;
+};
+
 class OpKernel {
  public:
   /**
@@ -390,8 +532,14 @@ class OperatorWithKernel : public OperatorBase {
                      const VariableNameMap& outputs, const AttributeMap& attrs)
       : OperatorBase(type, inputs, outputs, attrs) {}
 
+  // runtime infershape
   void InferShape(const Scope& scope) const override {
-    InferShape(InferShapeContext(*this, scope));
+    InferShape(RunTimeInferShapeContext(*this, scope));
+  }
+
+  // compile time infershape
+  void InferShape(const BlockDesc& block_desc) const {
+    InferShape(CompileTimeInferShapeContext(*this, block_desc));
   }
 
   void Run(const Scope& scope,
@@ -413,7 +561,8 @@ class OperatorWithKernel : public OperatorBase {
   }
 
  protected:
-  virtual void InferShape(const InferShapeContext& ctx) const = 0;
+  virtual void InferShape(const InferShapeContext& ctx) const {}
+  virtual void InferShape(const InferShapeContextBase& ctx) const {}
 };
 
 }  // namespace framework

paddle/framework/operator_test.cc

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/framework/operator.h"
 #include "gtest/gtest.h"
+#include "paddle/framework/op_info.h"
 #include "paddle/framework/op_registry.h"
 
 namespace paddle {