diff --git a/et_feeder/et_feeder_node.cpp b/et_feeder/et_feeder_node.cpp index aa7b6d9b..5afda544 100644 --- a/et_feeder/et_feeder_node.cpp +++ b/et_feeder/et_feeder_node.cpp @@ -38,6 +38,8 @@ ETFeederNode::ETFeederNode(std::shared_ptr node) { this->comm_dst_ = static_cast(attr.int32_val()); } else if (attr_name == "comm_tag") { this->comm_tag_ = static_cast(attr.int32_val()); + } else { + this->other_attrs_.emplace(attr_name, attr); } } } @@ -73,6 +75,20 @@ void ETFeederNode::setDepUnresolvedParentIDs( dep_unresolved_parent_ids_ = dep_unresolved_parent_ids; } +const 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_; } diff --git a/et_feeder/et_feeder_node.h b/et_feeder/et_feeder_node.h index 3e738eff..470ba3fd 100644 --- a/et_feeder/et_feeder_node.h +++ b/et_feeder/et_feeder_node.h @@ -19,6 +19,10 @@ class ETFeederNode { void setDepUnresolvedParentIDs( std::vector const& dep_unresolved_parent_ids); + const 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(); @@ -46,6 +50,8 @@ class ETFeederNode { std::unordered_set> children_set_{}; std::vector> children_vec_{}; std::vector dep_unresolved_parent_ids_{}; + std::unordered_map + other_attrs_{}; uint64_t id_; std::string name_;