JSON format for Chakra ET

src/feeder/et_feeder.cpp

@@ -173,4 +173,4 @@ void ETFeeder::readNextWindow() {
       dep_free_node_queue_.emplace(node);
     }
   }
-}
+}

src/feeder/et_feeder.h

@@ -33,13 +33,12 @@ class ETFeeder {
   void pushBackIssuableNode(uint64_t node_id);
   std::shared_ptr<ETFeederNode> lookupNode(uint64_t node_id);
   void freeChildrenNodes(uint64_t node_id);
-
- private:
   void readGlobalMetadata();
   std::shared_ptr<ETFeederNode> readNode();
   void readNextWindow();
   void resolveDep();
 
+ private:
   ProtoInputStream trace_;
   const uint32_t window_size_;
   bool et_complete_;
@@ -54,4 +53,4 @@ class ETFeeder {
   std::unordered_set<std::shared_ptr<ETFeederNode>> dep_unresolved_node_set_{};
 };
 
-} // namespace Chakra
+} // namespace Chakra

src/feeder/json_node.cpp

@@ -0,0 +1,194 @@
+#include "json_node.h"
+
+// JSONNode default constructor
+JSONNode::JSONNode() {}
+
+// JSONNode copy constructor
+JSONNode::JSONNode(const JSONNode& t) {
+  node_id = t.node_id;
+  node_name = t.node_name;
+  node_type = t.node_type;
+  is_cpu_op = t.is_cpu_op;
+  runtime = t.runtime;
+  data_deps = t.data_deps;
+  dep_unresolved_parent_ids_json = t.dep_unresolved_parent_ids_json;
+  children_vec_json = t.children_vec_json;
+  children_set_json = t.children_set_json;
+
+  if (node_type == NodeType::COMM_SEND_NODE ||
+      node_type == NodeType::COMM_RECV_NODE ||
+      node_type == NodeType::COMM_COLL_NODE) {
+    tensor_size = t.tensor_size;
+    comm_type = t.comm_type;
+    comm_priority = t.comm_priority;
+    comm_size = t.comm_size;
+    comm_src = t.comm_src;
+    comm_dst = t.comm_dst;
+    comm_tag = t.comm_tag;
+  }
+}
+
+// JSONNode constructor
+JSONNode::JSONNode(json data, uint64_t id) {
+  try {
+    node_id = data["workload_graph"][id]["Id"];
+  } catch (...) {
+    std::cerr << "node_id not specified in ET" << std::endl;
+  }
+  try {
+    node_name = data["workload_graph"][id]["Name"];
+  } catch (...) {
+    std::cerr << "node_name not specified in ET" << std::endl;
+  }
+  try {
+    node_type = data["workload_graph"][id]["NodeType"];
+  } catch (...) {
+    std::cerr << "node_type not specified in ET" << std::endl;
+  }
+  try {
+    is_cpu_op = data["workload_graph"][id]["is_cpu_op"];
+  } catch (...) {
+    std::cerr << "is_cpu_op not specified in ET" << std::endl;
+  }
+  try {
+    runtime = data["workload_graph"][id]["runtime"];
+  } catch (...) {
+  }
+  try {
+    data_deps =
+        data["workload_graph"][id]["data_deps"].get<std::vector<uint64_t>>();
+  } catch (...) {
+    std::cerr << "data deps not specified in ET" << std::endl;
+  }
+
+  if (node_type == NodeType::COMM_SEND_NODE ||
+      node_type == NodeType::COMM_RECV_NODE ||
+      node_type == NodeType::COMM_COLL_NODE) {
+    try {
+      tensor_size = data["workload_graph"][id]["tensor_size"];
+    } catch (...) {
+    }
+    try {
+      comm_type = data["workload_graph"][id]["comm_type"];
+    } catch (...) {
+    }
+    try {
+      comm_priority = data["workload_graph"][id]["comm_priority"];
+    } catch (...) {
+      comm_priority = 0; // Protobuf defaults to 0
+    }
+    try {
+      comm_size = data["workload_graph"][id]["comm_size"];
+    } catch (...) {
+    }
+    try {
+      comm_src = data["workload_graph"][id]["comm_src"];
+    } catch (...) {
+    }
+    try {
+      comm_dst = data["workload_graph"][id]["comm_dst"];
+    } catch (...) {
+    }
+    try {
+      comm_tag = data["workload_graph"][id]["comm_tag"];
+    } catch (...) {
+    }
+  }
+}
+
+// Node id
+uint64_t JSONNode::id() const {
+  return node_id;
+}
+
+// Node name
+std::string JSONNode::name() const {
+  return node_name;
+}
+
+// Node type
+int JSONNode::type() const {
+  return node_type;
+}
+
+// Check if CPU OP
+bool JSONNode::isCPUOp() const {
+  return is_cpu_op;
+}
+
+// Runtime
+uint64_t JSONNode::getRuntime() const {
+  return runtime;
+}
+
+// Num ops
+uint64_t JSONNode::getNumOps() const {
+  return num_ops;
+}
+
+// Tensor size
+uint64_t JSONNode::getTensorSize() const {
+  return tensor_size;
+}
+
+// Comm type
+int64_t JSONNode::getCommType() const {
+  return comm_type;
+}
+
+// Comm priority
+uint32_t JSONNode::getCommPriority() const {
+  return comm_priority;
+}
+
+// Comm size
+uint64_t JSONNode::getCommSize() const {
+  return comm_size;
+}
+
+// Comm src
+uint32_t JSONNode::getCommSrc() const {
+  return comm_src;
+}
+
+// Comm dst
+uint32_t JSONNode::getCommDst() const {
+  return comm_dst;
+}
+
+// Comm tag
+uint32_t JSONNode::getCommTag() const {
+  return comm_tag;
+}
+
+// Dependency unresolved parent IDs
+void JSONNode::addDepUnresolvedParentID(uint64_t node_id) {
+  dep_unresolved_parent_ids_json.emplace_back(node_id);
+}
+
+// Get dependency unresolved parent IDs
+std::vector<uint64_t> JSONNode::getDepUnresolvedParentIDs() {
+  return dep_unresolved_parent_ids_json;
+}
+
+// Set dependency unresolved parent IDs
+void JSONNode::setDepUnresolvedParentIDs(
+    std::vector<uint64_t> const& dep_unresolved_parent_ids) {
+  dep_unresolved_parent_ids_json = dep_unresolved_parent_ids;
+}
+
+// Add child
+void JSONNode::addChild(JSONNode node) {
+  // Avoid adding the same child node multiple times
+  // addChild is called multiple times to resolve dependencies
+  if (children_set_json.find(node) != children_set_json.end()) {
+    return;
+  }
+  children_vec_json.emplace_back(node);
+  children_set_json.emplace(node);
+}
+
+// Get children vector
+std::vector<JSONNode> JSONNode::getChildren() {
+  return children_vec_json;
+}

src/feeder/json_node.h

@@ -0,0 +1,150 @@
+#pragma once
+
+#include <json/json.hpp>
+#include <fstream>
+#include <functional>
+#include <iostream>
+#include <queue>
+#include <set>
+#include <string>
+
+using json = nlohmann::json;
+
+enum NodeType : int {
+  INVALID_NODE = 0,
+  METADATA_NODE = 1,
+  MEM_LOAD_NODE = 2,
+  MEM_STORE_NODE = 3,
+  COMP_NODE = 4,
+  COMM_SEND_NODE = 5,
+  COMM_RECV_NODE = 6,
+  COMM_COLL_NODE = 7
+};
+
+class JSONNode {
+ private:
+  uint64_t node_id;
+  std::string node_name;
+  int node_type;
+  bool is_cpu_op;
+  uint64_t runtime;
+  uint64_t num_ops;
+  uint64_t tensor_size;
+  int64_t comm_type;
+  uint32_t comm_priority;
+  uint64_t comm_size;
+  uint32_t comm_src;
+  uint32_t comm_dst;
+  uint32_t comm_tag;
+
+ public:
+  std::vector<uint64_t> data_deps{};
+  std::vector<uint64_t> dep_unresolved_parent_ids_json{};
+  std::vector<JSONNode> children_vec_json{};
+
+  // Compare function for set
+  struct CompareJSONNodesLT {
+    bool operator()(const JSONNode& a, const JSONNode& b) const {
+      return a.node_id < b.node_id;
+    }
+  };
+  std::set<JSONNode, CompareJSONNodesLT> children_set_json{};
+
+  JSONNode();
+  JSONNode(const JSONNode& t);
+  JSONNode(json data, uint64_t id);
+  uint64_t id() const;
+  std::string name() const;
+  int type() const;
+  bool isCPUOp() const;
+  uint64_t getRuntime() const;
+  uint64_t getNumOps() const;
+  uint64_t getTensorSize() const;
+  int64_t getCommType() const;
+  uint32_t getCommPriority() const;
+  uint64_t getCommSize() const;
+  uint32_t getCommSrc() const;
+  uint32_t getCommDst() const;
+  uint32_t getCommTag() const;
+  void addDepUnresolvedParentID(uint64_t node_id);
+  std::vector<uint64_t> getDepUnresolvedParentIDs();
+  void setDepUnresolvedParentIDs(
+      std::vector<uint64_t> const& dep_unresolved_parent_ids);
+  void addChild(JSONNode node);
+  std::vector<JSONNode> getChildren();
+
+  // Define the == operator for comparison
+  bool operator==(const JSONNode& other) const {
+    return node_id == other.node_id && node_name == other.node_name &&
+        node_type == other.node_type && is_cpu_op == other.is_cpu_op &&
+        runtime == other.runtime && num_ops == other.num_ops &&
+        tensor_size == other.tensor_size && comm_type == other.comm_type &&
+        comm_priority == other.comm_priority && comm_size == other.comm_size &&
+        comm_src == other.comm_src && comm_dst == other.comm_dst &&
+        comm_tag == other.comm_tag && data_deps == other.data_deps &&
+        dep_unresolved_parent_ids_json ==
+        other.dep_unresolved_parent_ids_json &&
+        children_vec_json == other.children_vec_json &&
+        children_set_json == other.children_set_json;
+  }
+
+  // Overload the assignment operator
+  JSONNode& operator=(const JSONNode& other) {
+    if (this != &other) {
+      // Copy all member variables
+      node_id = other.node_id;
+      node_name = other.node_name;
+      node_type = other.node_type;
+      is_cpu_op = other.is_cpu_op;
+      runtime = other.runtime;
+      num_ops = other.num_ops;
+      tensor_size = other.tensor_size;
+      comm_type = other.comm_type;
+      comm_priority = other.comm_priority;
+      comm_size = other.comm_size;
+      comm_src = other.comm_src;
+      comm_dst = other.comm_dst;
+      comm_tag = other.comm_tag;
+      data_deps = other.data_deps;
+      dep_unresolved_parent_ids_json = other.dep_unresolved_parent_ids_json;
+      children_vec_json = other.children_vec_json;
+      children_set_json = other.children_set_json;
+    }
+    return *this;
+  }
+};
+
+// Define a custom hash function for unordered set
+namespace std {
+template <>
+struct hash<JSONNode> {
+  std::size_t operator()(const JSONNode& node) const {
+    std::size_t h1 = std::hash<int64_t>()(node.id());
+    std::size_t h2 = std::hash<std::string>()(node.name());
+    std::size_t h3 = std::hash<int>()(node.type());
+    std::size_t h4 = std::hash<bool>()(node.isCPUOp());
+    std::size_t h5 = std::hash<int64_t>()(node.getRuntime());
+
+    // A prime number for bit manipulation
+    const std::size_t prime = 31;
+
+    // Combine the hash of the current member with the hashes of the previous
+    // members
+    std::size_t hash = h1;
+    hash = hash * prime + h2;
+    hash = hash * prime + h3;
+    hash = hash * prime + h4;
+    hash = hash * prime + h5;
+
+    return hash;
+  }
+};
+} // namespace std
+
+// Compare function for JSON node for priority queue
+struct CompareJSONNodesGT
+    : public std::binary_function<JSONNode, JSONNode, bool> {
+  bool operator()(const JSONNode lhs, const JSONNode rhs) const {
+    return lhs.id() > rhs.id();
+  }
+};