Merge pull request #150 from jinsun-yoo/et_generator

In et_generator, change 'comm_size' attr from uint64 to int64
mlcommons · Oct 23, 2024 · 47c8154 · 47c8154
2 parents f256593 + 796903a
commit 47c8154
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/converter/text_converter.py b/src/converter/text_converter.py
@@ -104,7 +104,7 @@ def get_comm_type(self, comm_type: str) -> int:
     def get_comm_coll_node(self, layer_name: str, comm_type: str, comm_size: int) -> Any:
         node = self.get_node(f"COMM_COLL_NODE_{layer_name}_{comm_type}", COMM_COLL_NODE)
         node.attr.append(ChakraAttr(name="comm_type", int64_val=self.get_comm_type(comm_type)))
-        node.attr.append(ChakraAttr(name="comm_size", uint64_val=comm_size))
+        node.attr.append(ChakraAttr(name="comm_size", int64_val=comm_size))
         return node
 
     def add_parent(self, child_node: Any, parent_node: Any) -> None:

diff --git a/src/feeder/et_feeder_node.cpp b/src/feeder/et_feeder_node.cpp
@@ -25,7 +25,7 @@ ETFeederNode::ETFeederNode(std::shared_ptr<ChakraProtoMsg::Node> node) {
     } else if (attr_name == "comm_priority") {
       this->comm_priority_ = static_cast<uint32_t>(attr.int32_val());
     } else if (attr_name == "comm_size") {
-      this->comm_size_ = attr.int64_val();
+      this->comm_size_ = static_cast<uint64_t>(attr.int64_val());
     } else if (attr_name == "comm_src") {
       this->comm_src_ = static_cast<uint32_t>(attr.int32_val());
     } else if (attr_name == "comm_dst") {

diff --git a/src/generator/generator.py b/src/generator/generator.py
@@ -191,7 +191,7 @@ def generate_comm_coll_node(num_npus: int, comm_size: int, comm_type: int, node_
 
             node = get_node(node_name, COMM_COLL_NODE)
             node.attr.append(ChakraAttr(name="is_cpu_op", bool_val=False))
-            node.attr.extend([get_comm_type_attr(comm_type), ChakraAttr(name="comm_size", uint64_val=comm_size)])
+            node.attr.extend([get_comm_type_attr(comm_type), ChakraAttr(name="comm_size", int64_val=comm_size)])
             encode_message(et, node)