Graph Format

graph_format.md

struct Edge {
  int src;
  int dst;
  float feature[100];
}

struct Node {
  int node_type;
  int id;
  float feature[100];
}

std::vector<Edge> edge_list;
std::vector<Node> node_list;

// See also feature.cc::SerializeFeatures 
TVMByteArray SerializeGraph(const std::vector<Edge>& edge_list, const std::vector<Node>& node_list, std::vector<char>* out) {
   size_t pt = 0;

   /* Serialization format
    * {
    *   int n_edges;
    *   int n_nodes;
    *   Edge edges[n_edges];
    *   Node nodes[n_nodes];
    * }
    */
   
   size_t total_bytes = sizeof(int) + sizeof(int) + sizeof(Edge) * edge_list.size() + sizeof(Node) * node_list.size();
   
   out_data->reserve(total_bytes);
   char* ptr = out_data->data();
   
   int len = static_cast<int>(edge_list.size());
   memcpy(ptr, &len, sizeof(int));
   ptr += sizeof(int);
   
   len = static_cast<int>(node_list.size());
   memcpy(ptr, &len, sizeof(int));
   ptr += sizeof(int);

   // Serialize edge_list
   for (const Edge& edge : edge_list) {
      memcpy(ptr, &edge.src, sizeof(int));
      ptr += sizeof(int);
      memcpy(ptr, &edge.dst, sizeof(int));
      ptr += sizeof(int);
      memcpy(ptr, edge.feature, sizeof(float) * 100);
      ptr += sizeof(float) * 100;
   }

   // Serialize node_list
   ...

   CHECK_EQ(ptr - out_data->data(), total_bytes);
   
   return TVMByteArray{out_data->data(), total_bytes};
}

Edge = namedtuple("Edge", ["src", "dst", "feature"])
Node = namedtuple("Node", ["node_type", "id", "feature"])


import struct

# See also feature.py::unpack_feature
def DeserializeGraph(byte_array) -> Tuple[List[Edge], List[Node]]:

  edge_list = []
  node_list = []

  size_of_int = 4
  size_of_float = 4
  vec_len = 100
    
  # unpack sizes
  offset = 0
  n_edges = struct.unpack_from("1i", byte_arr, offset=offset)[0]
  offset += size_of_int

  n_nodes = struct.unpack_from("1i", byte_arr, offset=offset)[0]
  offset += size_of_int    
    
  # unpack edge list
  for i in range(n_edges):
     src = struct.unpack_from("1i", byte_arr, offset=offset)[0]
     offset += size_of_int
     dst = struct.unpack_from("1i", byte_arr, offset=offset)[0]
     offset += size_of_int
     feature = struct.unpack_from("%df" % vec_len, byte_arr, offset=offset)[0]
     offset += size_of_float * vec_len
     edge_list.append(Edge(src, dst, feature))

  # unpack node list
  ...
  
  return edge_list, node_list

Translating an AST to a graph with visitor pattern

class NodeGather: public StmtExprVisitor {

  void VisitStmt(const Stmt& n) {
    node_list.push_back(Node(..));
    node_to_index[n->get()] = node_to_index.size();
    StmtExprVisitor::VisitExpr(n);
  }

  void VisitExpr(const PrimExpr& n) {
    node_list.push_back(Node(..));
    node_to_index[n->get()] = node_to_index.size();
    StmtExprVisitor::VisitExpr(n);
  }

  std::unordered_map<const Object*, int> node_to_index;
  std::vector<Node> node_list;
}

class EdgeGather: public StmtExprVisitor {
  void VisitExpr_(const AddNode* op) {
      int src_id = node_to_index[static_cast<const Object*>(op)];
      int dst_id = node_to_index[op->a->get()];
      edge_list.push_back(Edge(src_id, dst_id));

      dst_id = node_to_index[op->b->get()];
      edge_list.push_back(Edge(src_id, dst_id));

      StmtExprVisitor::VisitExpr(n);
  }

  std::vector<Edge> edge_list;
}