JSONNode and WrapperNode

Author: Vinay Ramakrishnaiah (AMD)

Introduction

This document details two classes that complement the ETFeeder module:

• JSONNode: A lightweight representation of a trace node parsed directly from a JSON-formatted Chakra trace.
• WrapperNode: A format-agnostic wrapper that abstracts over both protobuf-based ETFeederNode/ETFeeder and JSONNode-based feeding, providing a unified API to issue nodes, resolve dependencies, and manage queues.

---

JSONNode Class

Overview

The JSONNode class represents an operation/node within a JSON-formatted Chakra execution trace. It encapsulates node identifiers, attributes, and dependency tracking, and provides methods to manage child relationships and unresolved parent dependencies. It is designed to be stored in ordered containers and priority queues to support dependency resolution and node issuing.

Key Attributes

• Identification and metadata

  • uint64_t node_id: Unique identifier.
  • std::string node_name: The name of the operation.
  • int node_type: Node type; communication nodes use specific values.
  • bool is_cpu_op: Indicates whether the operation executes on CPU.
  • uint64_t runtime: Duration (if present in trace).
  • uint64_t num_ops: Operation count (may be absent in JSON input depending on schema).

• Communication metadata (only set for communication nodes)

  • Node types considered communication:
    • Enums:
      • COMM_SEND_NODE = 5
      • COMM_RECV_NODE = 6
      • COMM_COLL_NODE = 7
  • uint64_t tensor_size: Tensor size
  • int64_t comm_type: communication operator type
  • uint32_t comm_priority (defaults to 0 if not present)
  • uint64_t comm_size: communication size
  • uint32_t comm_src: source
  • uint32_t comm_dst: destination
  • uint32_t comm_tag: communication tag

• Dependency management

  • std::vector<uint64_t> data_deps: IDs of parent nodes this node depends on (parsed from JSON data_deps).
  • std::vector<uint64_t> dep_unresolved_parent_ids_json: Parents not yet loaded or found.
  • std::vector<JSONNode> children_vec_json: Child nodes that depend on this node.
  • std::set<JSONNode, CompareJSONNodesLT> children_set_json: Ordered set to prevent duplicate children. CompareJSONNodesLT orders by node_id.

• Ordering/comparison helpers

  • struct CompareJSONNodesLT: a.node_id < b.node_id.
  • struct CompareJSONNodesGT: Used for priority queue to order by greater id(); enables min-issue behavior by flipping priority.

• Hashing

  • namespace std::hash<JSONNode>: Custom hasher combines id, name, type, is_cpu_op, and runtime for use in unordered containers.

Key Methods

• Constructors

  • JSONNode(): Default.
  • JSONNode(const JSONNode& t): Copy constructor; copies metadata, deps, children; communication fields only if node_type is 5/6/7.
  • JSONNode(json data, uint64_t id): Constructs from parsed JSON document and node index (continuous index in workload_graph). Extracts fields with error tolerance.

• Accessors

  • 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

• Dependency helpers

  • void addDepUnresolvedParentID(uint64_t node_id): Track missing parent.
  • std::vector<uint64_t> getDepUnresolvedParentIDs(): Get the list of unresolved parent IDs
  • void setDepUnresolvedParentIDs(const std::vector<uint64_t>& dep_unresolved_parent_ids): Update the list of unresolved parent IDs

• Child management

  • void addChild(JSONNode node): Adds child if not already present (checked via children_set_json).
  • std::vector<JSONNode> getChildren(): Get children nodes

• Operators

  • bool operator==(const JSONNode& other) const: Deep equality across all relevant fields and containers.
  • JSONNode& operator=(const JSONNode& other): Assignment copies all fields and containers.

Important Considerations

• Conditional communication fields: Communication metadata is only valid/populated if node_type is send/recv/coll (types 5/6/7).
• Input tolerance: JSON parsing is guarded by try/catch; missing fields either default (e.g., comm_priority = 0) or emit diagnostic messages.
• Deduplication: addChild prevents duplicates via children_set_json; children_vec_json mirrors the set order for iteration.
• Ordering for issuance: CompareJSONNodesGT is used to prioritize nodes by ID; coupled with dependency-resolution logic, it ensures deterministic issuance order.

---

WrapperNode Class

Overview

WrapperNode unifies trace feeding across two formats:

• Protobuf (.et files) via Chakra::ETFeeder and Chakra::ETFeederNode.
• JSON (.json files) via JSONNode.

It abstracts the lifecycle (initialization, reading windows, dependency resolution, node issuance), so downstream simulators can consume nodes without format-specific code. It mirrors much of ETFeeder’s API while adding JSON-specific data structures and logic.

Key Attributes

• Format selection and core handles

  • enum format format_type_: Protobuf or JSON (based on filename extension *.et or *.json).
  • Chakra::ETFeeder* et_feeder_: Feeder for protobuf traces.
  • std::shared_ptr<Chakra::ETFeederNode> node_: Current protobuf node.
  • std::ifstream jsonfile_: Input stream for JSON traces.
  • json data_: Parsed JSON document (nlohmann::json).
  • JSONNode json_node_: Current JSON node.
  • int64_t node_idx_ = -1: Index in workload_graph array for the current JSON node.

• Queues and dependency tracking (protobuf vs JSON)

  • Protobuf queues
    • std::queue<std::shared_ptr<Chakra::ETFeederNode>> push_back_queue_proto
  • JSON queues/graphs
    • std::queue<JSONNode> push_back_queue_json
    • std::unordered_map<uint64_t, JSONNode> dep_graph_json: Loaded nodes by ID.
    • std::unordered_set<uint64_t> dep_free_node_id_set_json: IDs currently dependency-free.
    • std::priority_queue<JSONNode, std::vector<JSONNode>, CompareJSONNodesGT> dep_free_node_queue_json: JSON nodes ready to issue (ordered by id()).
    • std::unordered_set<JSONNode, std::hash<JSONNode>> dep_unresolved_node_set_json: Nodes with unresolved parents (hashed by custom std::hash<JSONNode>).
  • Window and completion flags
    • int window_size_json: Number of nodes in JSON workload_graph.
    • bool json_et_complete_: Whether all JSON nodes have been read.

Key Methods

• Construction and initialization

  • WrapperNode(): Default.
  • WrapperNode(const WrapperNode& t): Copies format selection, queues, graphs, and current node state.
  • WrapperNode(std::string filename): Calls createWrapper(filename).
  • void createWrapper(std::string filename): Detects format from extension:
    • .et → Protobuf: initializes ETFeeder.
    • .json → JSON: parses the entire document, sets window_size_json, and invokes readNextWindow() (legacy behavior: full-file read).
  • ~WrapperNode(): Destructor
  • void releaseMemory(): Releases underlying resources (delete et_feeder_ for protobuf, closes JSON file stream for JSON).

• Node accessors (current node)

  • std::shared_ptr<Chakra::ETFeederNode> getProtobufNode()
  • JSONNode getJSONNode()

• Graph management

  • Protobuf:
    • void addNode(std::shared_ptr<Chakra::ETFeederNode> node): Delegates to ETFeeder::addNode.
    • void removeNode(uint64_t node_id): Delegates; may trigger window reads.
  • JSON:
    • void addNode(JSONNode node): Inserts into dep_graph_json.
    • void removeNode(uint64_t node_id): Erases from dep_graph_json; if not complete and queues underflow, triggers readNextWindow().

• Reading/parsing (JSON)

  • int64_t findNodeIndexJSON(uint64_t node_id): Linear scan of workload_graph to find continuous index.
  • JSONNode readNode(uint64_t node_idx): Builds JSONNode from data_ at index:
    • For each data_deps parent ID:
      • If found in dep_graph_json, add current node as child to parent.
      • Else, mark as unresolved via addDepUnresolvedParentID.
    • Unresolved nodes are tracked in dep_unresolved_node_set_json.
  • void readNextWindow(): Reads nodes and resolves deps in a loop:
    • Reads all nodes up to window_size_json.
    • After each addition, calls resolveDep().
    • Uses an arbitrarily large loop bound (256 * window_size_json) to ensure repeated resolution attempts.
    • After reading, any node with zero data_deps is added to dep_free_node_id_set_json and dep_free_node_queue_json.

• Dependency resolution

  • void resolveDep(): For JSON:
    • Iterates unresolved nodes; for each unresolved parent ID, checks if now present in dep_graph_json.
    • If present, adds current node as child to the parent and removes that unresolved parent ID.
    • If a node’s unresolved list becomes empty, it is removed from dep_unresolved_node_set_json; otherwise, its unresolved list is updated.
  • Protobuf path delegates to ETFeeder::resolveDep().

• Issuable nodes handling

  • void pushBackIssuableNode(uint64_t node_id): Marks a node as dependency-free and enqueues to dep_free_node_queue_json (JSON) or delegates (protobuf).
  • void freeChildrenNodes(uint64_t node_id): Upon completing a node:
    • For each child: remove the completed node ID from child.data_deps.
    • If a child becomes dependency-free, add to dep_free_node_id_set_json and dep_free_node_queue_json.
    • Protobuf path delegates to ETFeeder::freeChildrenNodes.
  • void getNextIssuableNode(): Retrieves the next dependency-free node:
    • Protobuf: node_ = et_feeder_->getNextIssuableNode().
    • JSON: Pops top of dep_free_node_queue_json, sets json_node_, updates node_idx_ via findNodeIndexJSON, and erases from ID set. If queue empty, sets node_idx_ = -1.
  • bool hasNodesToIssue(): Checks whether there are nodes ready to issue:
    • Protobuf: delegates to ETFeeder::hasNodesToIssue().
    • JSON: Returns !(dep_graph_json.empty() && dep_free_node_queue_json.empty()).

• Queue utilities (format-generic wrappers)

  • void push_to_queue(): Pushes current node to push_back_queue_*.
  • bool is_queue_empty(): Checks the appropriate push-back queue.
  • void queue_front(): Sets current node from the front of the queue.
  • void pop_from_queue(): Pops the front element.

• Common metadata accessors (format-agnostic)

  • uint64_t getNodeID()
  • std::string getNodeName()
  • int getNodeType()
  • bool isCPUOp()
  • uint64_t getRuntime()
  • uint64_t getNumOps()
  • uint64_t getTensorSize()
  • int64_t getCommType()
  • uint32_t getCommPriority()
  • uint64_t getCommSize()
  • uint32_t getCommSrc()
  • uint32_t getCommDst()
  • uint32_t getCommTag()
  • Protobuf-only:
    • uint32_t getCommSrcOff()
    • uint32_t getCommDstOff()

• Children access

  • Protobuf: void getChildren(std::vector<std::shared_ptr<Chakra::ETFeederNode>>& childrenNodes)
  • JSON: void getChildren(std::vector<JSONNode>& childrenNodes)

• Lookup (retrieve by ID)

  • void lookupNode(uint64_t node_id): Protobuf delegates to ETFeeder::lookupNode.
  • JSON uses dep_graph_json.at(node_id); throws and logs if not loaded yet.

Usage Workflow

1. Initialization

   • Create a WrapperNode with the trace filename:
     • WrapperNode wrapper("path/to/trace.et") or WrapperNode wrapper("path/to/trace.json").
   • createWrapper detects format (.et → Protobuf, .json → JSON).
   • For JSON, the full document is parsed and readNextWindow() builds the dependency graph.

2. Simulation loop

   • Check readiness: hasNodesToIssue().
   • Retrieve next issuable node: getNextIssuableNode().
     • Protobuf: returns ETFeederNode via internal state.
     • JSON: returns JSONNode via getJSONNode().
   • Process the node in your simulator.
   • After completion:
     • Free children dependencies: freeChildrenNodes(node_id).
     • Optionally remove from graph: removeNode(node_id) (JSON will read more if needed).
   • Continue until hasNodesToIssue() returns false.

3. Node inspection

   • Access common metadata (getNodeID, getNodeName, etc.) regardless of format.
   • Retrieve children via the appropriate getChildren overload.

Important Considerations

• Format detection: WrapperNode determines trace format from the filename extension. Ensure correct extensions (.et for protobuf, .json for JSON).
• JSON reading strategy: Legacy design reads the entire JSON file at once and performs iterative dependency resolution across the whole graph.
• Issuance order: JSON dependency-free nodes are managed in a priority queue ordered by id() (CompareJSONNodesGT), ensuring deterministic issuance.
• Duplicate prevention: The JSON path uses sets (dep_free_node_id_set_json, children sets) to avoid re-enqueuing or adding duplicate references.
• Error handling: Many operations log to std::cerr and exit(-1) on fatal path errors (e.g., unsupported format), so ensure inputs are valid.
• Memory management: Call releaseMemory() when done to close streams (JSON) or delete feeder (protobuf).
• Protobuf-only fields: getCommSrcOff() and getCommDstOff() are only valid in protobuf mode; calling them in JSON mode triggers an error.

---

Relationship to ETFeeder

WrapperNode mirrors ETFeeder concepts (windows, dependency graph, dependency-free priority queue, unresolved set) while providing a uniform API across formats. For protobuf, most operations delegate to ETFeeder/ETFeederNode. For JSON, analogous data structures in WrapperNode and JSONNode implement equivalent functionality.

For broader context on ETFeeder's design, see the Chakra ETFeeder documentation:

• Chakra ETFeeder wiki (overview, methods, and workflow).

---

Conclusion

JSONNode and WrapperNode extend the Chakra Feeder's capabilities to handle JSON traces and provide a format-agnostic interface for simulators. Understanding their attributes and methods will help integrate Chakra traces into downstream tools while preserving correctness regarding dependency resolution and node issuance.