Merge pull request #4276 from Canpio/add_program_proto

Add `BlockDesc` and `ProgramDesc` to framework.proto

paddle/framework/attribute.cc

@@ -19,6 +19,15 @@ limitations under the License. */
 namespace paddle {
 namespace framework {
 
+static ProgramDesc* g_program_desc = nullptr;
+
+ProgramDesc& GetProgramDesc() {
+  if (g_program_desc == nullptr) {
+    g_program_desc = new ProgramDesc();
+  }
+  return *g_program_desc;
+}
+
 template <>
 AttrType AttrTypeID<int>() {
   return INT;
@@ -47,40 +56,44 @@ template <>
 AttrType AttrTypeID<std::vector<std::pair<int, int>>>() {
   return INT_PAIRS;
 }
+template <>
+AttrType AttrTypeID<BlockDesc>() {
+  return BLOCK;
+}
 
 Attribute GetAttrValue(const OpDesc::Attr& attr_desc) {
   switch (attr_desc.type()) {
-    case paddle::framework::AttrType::INT: {
+    case framework::AttrType::INT: {
       return attr_desc.i();
     }
-    case paddle::framework::AttrType::FLOAT: {
+    case framework::AttrType::FLOAT: {
       return attr_desc.f();
     }
-    case paddle::framework::AttrType::STRING: {
+    case framework::AttrType::STRING: {
       return attr_desc.s();
     }
-    case paddle::framework::AttrType::INTS: {
+    case framework::AttrType::INTS: {
       std::vector<int> val(attr_desc.ints_size());
       for (int i = 0; i < attr_desc.ints_size(); ++i) {
         val[i] = attr_desc.ints(i);
       }
       return val;
     }
-    case paddle::framework::AttrType::FLOATS: {
+    case framework::AttrType::FLOATS: {
       std::vector<float> val(attr_desc.floats_size());
       for (int i = 0; i < attr_desc.floats_size(); ++i) {
         val[i] = attr_desc.floats(i);
       }
       return val;
     }
-    case paddle::framework::AttrType::STRINGS: {
+    case framework::AttrType::STRINGS: {
       std::vector<std::string> val(attr_desc.strings_size());
       for (int i = 0; i < attr_desc.strings_size(); ++i) {
         val[i] = attr_desc.strings(i);
       }
       return val;
     }
-    case paddle::framework::AttrType::INT_PAIRS: {
+    case framework::AttrType::INT_PAIRS: {
       std::vector<std::pair<int, int>> val(attr_desc.int_pairs_size());
       for (int i = 0; i < attr_desc.int_pairs_size(); ++i) {
         val[i].first = attr_desc.int_pairs(i).first();
@@ -88,6 +101,9 @@ Attribute GetAttrValue(const OpDesc::Attr& attr_desc) {
       }
       return val;
     }
+    case framework::AttrType::BLOCK: {
+      return GetProgramDesc().mutable_blocks(attr_desc.block_idx());
+    }
   }
   PADDLE_ENFORCE(false, "Unknown OpDesc::AttrDesc::type !");
   return boost::blank();

paddle/framework/attribute.h

@@ -29,11 +29,13 @@ namespace framework {
 
 typedef boost::variant<boost::blank, int, float, std::string, std::vector<int>,
                        std::vector<float>, std::vector<std::string>,
-                       std::vector<std::pair<int, int>>>
+                       std::vector<std::pair<int, int>>, BlockDesc*>
     Attribute;
 
 typedef std::unordered_map<std::string, Attribute> AttributeMap;
 
+ProgramDesc& GetProgramDesc();
+
 template <typename T>
 AttrType AttrTypeID();
 

paddle/framework/framework.proto

@@ -23,6 +23,7 @@ enum AttrType {
   FLOATS = 4;
   STRINGS = 5;
   INT_PAIRS = 6;
+  BLOCK = 7;
 }
 
 message IntPair {
@@ -44,6 +45,7 @@ message OpDesc {
     repeated float floats = 7;
     repeated string strings = 8;
     repeated IntPair int_pairs = 9;
+    optional int32 block_idx = 10;
   };
 
   message Var {
@@ -108,3 +110,12 @@ message VarDesc {
   required string name = 1;
   optional LoDTensorDesc lod_tensor = 2;
 }
+
+message BlockDesc {
+  required int32 idx = 1;
+  required int32 parent_idx = 2;
+  repeated VarDesc vars = 3;
+  repeated OpDesc ops = 4;
+}
+
+message ProgramDesc { repeated BlockDesc blocks = 1; }