Skip to content

Commit

Permalink
[ETFeeder] Added API to get other attrs.
Browse files Browse the repository at this point in the history
	modified:   et_feeder/et_feeder_node.cpp
	modified:   et_feeder/et_feeder_node.h
  • Loading branch information
changhai0109 committed Mar 7, 2024
1 parent e053939 commit 9dc4d06
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 0 deletions.
16 changes: 16 additions & 0 deletions et_feeder/et_feeder_node.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ ETFeederNode::ETFeederNode(std::shared_ptr<ChakraProtoMsg::Node> node) {
this->comm_dst_ = static_cast<uint32_t>(attr.int32_val());
} else if (attr_name == "comm_tag") {
this->comm_tag_ = static_cast<uint32_t>(attr.int32_val());
} else {
this->other_attrs_.emplace(attr_name, attr);
}
}
}
Expand Down Expand Up @@ -73,6 +75,20 @@ void ETFeederNode::setDepUnresolvedParentIDs(
dep_unresolved_parent_ids_ = dep_unresolved_parent_ids;
}

ChakraProtoMsg::AttributeProto& ETFeederNode::get_other_attr(
const string& attr_name) const {
if (this->has_other_attr(attr_name))
return this->other_attrs_.at(attr_name);
throw std::runtime_error(
"Asked for attr \"" + attr_name + "\" from node " +
std::to_string(this->id_) + ", which do not exist");
}

bool ETFeederNode::has_other_attr(const string& attr_name) const {
const auto& item = this->other_attrs_.find(attr_name);
return item != this->other_attrs_.end();
}

uint64_t ETFeederNode::id() {
return id_;
}
Expand Down
6 changes: 6 additions & 0 deletions et_feeder/et_feeder_node.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ class ETFeederNode {
void setDepUnresolvedParentIDs(
std::vector<uint64_t> const& dep_unresolved_parent_ids);

ChakraProtoMsg::AttributeProto& get_other_attr(
const std::string& attr_name) const;
bool has_other_attr(const std::string& attr_name) const;

uint64_t id();
std::string name();
bool is_cpu_op();
Expand Down Expand Up @@ -46,6 +50,8 @@ class ETFeederNode {
std::unordered_set<std::shared_ptr<ETFeederNode>> children_set_{};
std::vector<std::shared_ptr<ETFeederNode>> children_vec_{};
std::vector<uint64_t> dep_unresolved_parent_ids_{};
std::unordered_map<std::string, ChakraProtoMsg::AttributeProto&>
other_attrs_{};

uint64_t id_;
std::string name_;
Expand Down

0 comments on commit 9dc4d06

Please sign in to comment.