Refactored HloProtoMap::contains into ContainsOptimizedModule and ContainsOriginalModule to clearly distinguish between optimized and original HLO module lookups.

PiperOrigin-RevId: 825412534

Author: Profiler Team

xprof/utils/hlo_proto_map.cc

@@ -112,19 +112,13 @@ void HloProtoMap::AddHloProtosFromXSpace(const XSpace& space) {
   }
 }
 
-std::vector<absl::string_view> HloProtoMap::GetModuleList() const {
+std::vector<absl::string_view> HloProtoMap::GetModuleList(
+    bool is_original /*= false*/) const {
+  const auto& hlo_protos_map =
+      is_original ? original_hlo_protos_by_name_ : hlo_protos_by_name_;
   std::vector<absl::string_view> module_list;
-  module_list.reserve(hlo_protos_by_name_.size());
-  for (const auto& [name, hlo_proto] : hlo_protos_by_name_) {
-    module_list.push_back(name);
-  }
-  return module_list;
-}
-
-std::vector<absl::string_view> HloProtoMap::GetOriginalModuleList() const {
-  std::vector<absl::string_view> module_list;
-  module_list.reserve(original_hlo_protos_by_name_.size());
-  for (const auto& [name, hlo_proto] : original_hlo_protos_by_name_) {
+  module_list.reserve(hlo_protos_map.size());
+  for (const auto& [name, hlo_proto] : hlo_protos_map) {
     module_list.push_back(name);
   }
   return module_list;

xprof/utils/hlo_proto_map.h

@@ -49,19 +49,22 @@ class HloProtoMap {
   auto begin() const { return hlo_protos_by_program_id_.begin(); }
   auto end() const { return hlo_protos_by_program_id_.end(); }
 
-  bool contains(absl::string_view name) const {
+  bool ContainsOptimizedModule(absl::string_view name) const {
     return hlo_protos_by_name_.contains(name);
   }
 
+  bool ContainsOriginalModule(absl::string_view name) const {
+    return original_hlo_protos_by_name_.contains(name);
+  }
+
   bool contains(uint64_t program_id) const {
     return hlo_protos_by_program_id_.contains(program_id);
   }
 
-  // Returns a list of module names (not sorted).
-  std::vector<absl::string_view> GetModuleList() const;
-
-  // Returns a list of unoptimized/original module names (not sorted).
-  std::vector<absl::string_view> GetOriginalModuleList() const;
+  // Returns a list of HLO module names. If `is_original` is true, returns the
+  // names of the original/unoptimized modules; otherwise, returns the names of
+  // the optimized modules. The order of the modules is not guaranteed.
+  std::vector<absl::string_view> GetModuleList(bool is_original = false) const;
 
   // Returns a list of module names sorted alphabetically.
   std::vector<absl::string_view> GetSortedModuleList() const;

xprof/utils/hlo_proto_map_test.cc

@@ -19,7 +19,7 @@ using ::testing::status::StatusIs;
 
 TEST(HloProtoMapTest, GetOriginalModuleList) {
   HloProtoMap hlo_proto_map;
-  EXPECT_THAT(hlo_proto_map.GetOriginalModuleList(), IsEmpty());
+  EXPECT_THAT(hlo_proto_map.GetModuleList(true), IsEmpty());
 
   auto hlo_proto_1 = std::make_unique<xla::HloProto>();
   hlo_proto_1->mutable_hlo_module()->set_name("module1");
@@ -29,7 +29,7 @@ TEST(HloProtoMapTest, GetOriginalModuleList) {
   hlo_proto_2->mutable_hlo_module()->set_name("module2");
   hlo_proto_map.AddOriginalHloProto(2, std::move(hlo_proto_2));
 
-  EXPECT_THAT(hlo_proto_map.GetOriginalModuleList(),
+  EXPECT_THAT(hlo_proto_map.GetModuleList(true),
               UnorderedElementsAre("module1(1)", "module2(2)"));
 }