add SymbolTable for Block

paddle/framework/CMakeLists.txt

@@ -40,3 +40,6 @@ add_custom_command(TARGET framework_py_proto POST_BUILD
 
 cc_library(backward SRCS backward.cc DEPS net_op)
 cc_test(backward_test SRCS backward_test.cc DEPS backward recurrent_op device_context)
+
+cc_library(block SRCS block.cc DEPS framework_proto)
+cc_test(block_test SRCS block_test.cc DEPS block)

paddle/framework/block.cc

@@ -0,0 +1,64 @@
+/* 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/block.h"
+#include "paddle/platform/enforce.h"
+
+namespace paddle {
+namespace framework {
+
+const OpDesc* SymbolTable::NewOp() {
+  ops_.emplace_back();
+  return &ops_.back();
+}
+
+const VarDesc* SymbolTable::NewVar(const std::string& name) {
+  PADDLE_ENFORCE(vars_.insert(std::make_pair(name, VarDesc())).second,
+                 "Var called %s duplicates", name.c_str());
+  auto& var = vars_[name];
+  var.set_name(name);
+  return &var;
+}
+
+const VarDesc* SymbolTable::FindVar(const std::string& name,
+                                    bool recursive) const {
+  auto it = vars_.find(name);
+  if (it != vars_.end()) {
+    return &(it->second);
+  }
+  if (recursive && parent_) {
+    return parent_->FindVar(name, true);
+  }
+  // not found
+  return nullptr;
+}
+
+const OpDesc* SymbolTable::FindOp(size_t idx) const {
+  PADDLE_ENFORCE_LT(idx, ops_.size());
+  return &ops_[idx];
+}
+
+BlockDesc SymbolTable::Compile() const {
+  BlockDesc res;
+  for (const auto& op : ops_) {
+    *res.add_ops() = op;
+  }
+  for (const auto& it : vars_) {
+    *res.add_vars() = it.second;
+  }
+  return res;
+}
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/block.h

@@ -0,0 +1,71 @@
+/* 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 <map>
+#include <vector>
+#include "paddle/framework/framework.pb.h"
+
+namespace paddle {
+namespace framework {
+
+/*
+ * SymbolTable is the container of OpDesc and VarDesc.
+ */
+class SymbolTable {
+ public:
+  SymbolTable() {}
+  explicit SymbolTable(SymbolTable* parent) : parent_(parent) {}
+
+  /*
+   * Create a new OpDesc.
+   */
+  const OpDesc* NewOp();
+
+  /*
+   * Create a new VarDesc.
+   */
+  const VarDesc* NewVar(const std::string& name);
+
+  /*
+   * Find a VarDesc by name, if recursive true, find parent's SymbolTable
+   * recursively. This interface is proposed to support InferShape, find
+   * protobuf messages or variables and operators, pass pointers into
+   * InferShape.
+   *
+   * NOTE(superjom) maybe some C++ classes such as VarDescBuilder and
+   * OpDescBuilder should be proposed and embedded into pybind to enable python
+   * operate on C++ pointers.
+   */
+  const VarDesc* FindVar(const std::string& name, bool recursive = true) const;
+
+  /*
+   * Find a OpDesc by idx.
+   */
+  const OpDesc* FindOp(size_t idx) const;
+
+  /*
+   * Create a BlockDesc based on the registered VarDesc and OpDesc.
+   */
+  BlockDesc Compile() const;
+
+ private:
+  SymbolTable* parent_{nullptr};
+  std::vector<OpDesc> ops_;
+  std::map<std::string, VarDesc> vars_;
+};
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/block_test.cc

@@ -0,0 +1,51 @@
+/* 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/block.h"
+
+#include <gtest/gtest.h>
+
+namespace paddle {
+namespace framework {
+
+// there may be some complex init latter.
+class SymbolTableTester : public ::testing::Test {
+ protected:
+  SymbolTable table_;
+};
+
+TEST_F(SymbolTableTester, NewOp) { ASSERT_NE(table_.NewOp(), nullptr); }
+
+TEST_F(SymbolTableTester, NewVar) { ASSERT_NE(table_.NewVar("var1"), nullptr); }
+
+TEST_F(SymbolTableTester, FindOp) {
+  table_.NewOp();
+  ASSERT_NE(table_.FindOp(0), nullptr);
+}
+
+TEST_F(SymbolTableTester, Compile) {
+  table_.NewOp();
+  table_.NewOp();
+  table_.NewVar("var1");
+  table_.NewVar("var2");
+
+  auto block_desc = table_.Compile();
+  ASSERT_EQ(block_desc.ops_size(), 2);
+  ASSERT_EQ(block_desc.vars_size(), 2);
+  ASSERT_TRUE(block_desc.vars(0).name() == "var1" ||
+              block_desc.vars(1).name() == "var2");
+}
+
+}  // namespace framework
+}  // namespace paddle

paddle/framework/framework.proto

@@ -108,3 +108,8 @@ message VarDesc {
   required string name = 1;
   optional LoDTensorDesc lod_tensor = 2;
 }
+
+message BlockDesc {
+  repeated VarDesc vars = 1;
+  repeated OpDesc ops = 2;
+}