From cb4c26f273254cbf12ef0214abf361bf7d3cb54c Mon Sep 17 00:00:00 2001
From: Clive Verghese <cliveverghese@google.com>
Date: Sat, 14 Feb 2026 11:19:02 -0800
Subject: [PATCH] Refractor AddHloProto to get module_id from module

PiperOrigin-RevId: 870213930
---
 xprof/utils/hlo_proto_map.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/xprof/utils/hlo_proto_map.cc b/xprof/utils/hlo_proto_map.cc
index 3bfa473f..8f7320ca 100644
--- a/xprof/utils/hlo_proto_map.cc
+++ b/xprof/utils/hlo_proto_map.cc
@@ -70,7 +70,8 @@ ParseHloProtosFromXSpace(const XSpace& space) {
               absl::string_view byte_value = hlo_proto_stat->BytesValue();
               if (hlo_proto->ParseFromString(byte_value)) {
                 if (!hlo_protos
-                         .try_emplace(event_metadata.Id(), std::move(hlo_proto))
+                         .try_emplace(hlo_proto->hlo_module().id(),
+                                      std::move(hlo_proto))
                          .second) {
                   LOG(WARNING) << "Insert failed for hlo_proto with program_id"
                                << event_metadata.Id();