Add CoreId-to-CoreDetails map to OpStats.

PiperOrigin-RevId: 711909143

Author: bmass02

tensorflow/core/profiler/convert/xplane_to_op_stats.cc

@@ -357,11 +357,25 @@ OpStats ConvertXSpaceToOpStats(const XSpace& space,
     }
   }
 
-  // TODO(bvandermoon): Add the TPU equivalent for setting core details hostname
   if (!is_tpu) {
     CoreDetails& details =
         (*op_stats.mutable_core_id_to_details())[kDefaultGpuLocalCoreId];
     details.set_hostname(Hostname(space));
+  } else {
+    std::string hostname = Hostname(space);
+    auto& core_id_to_details = *op_stats.mutable_core_id_to_details();
+    for (const XPlane* device_plane : device_planes) {
+      XPlaneVisitor visitor =
+          tsl::profiler::CreateTfXPlaneVisitor(device_plane);
+      auto stat = visitor.GetStat(StatType::kCoreDetails);
+      if (stat.has_value()) {
+        CoreDetails core_details;
+        // TODO: Switch to StrOrRefValue once protobuf version is updated.
+        core_details.ParseFromString(stat->ToString());
+        core_details.set_hostname(hostname);
+        core_id_to_details[device_plane->id()] = core_details;
+      }
+    }
   }
 
   // Set program_id_to_name map in OpStats from Xspace

tensorflow/core/profiler/convert/xplane_to_step_events.cc

@@ -301,11 +301,11 @@ StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace) {
         // one more step than the "Step" line. We need to intersect them to get
         // the common step numbers.
         stream_step_events =
-            ConvertTpuDeviceTraceXLineToStepEvents(*tpu_core_id, line);
+            ConvertTpuDeviceTraceXLineToStepEvents(plane.Id(), line);
         IntersectCombineStepEvents(stream_step_events, &device_step_events);
       } else if (sc_core_id.has_value()) {
-        stream_step_events = ConvertTpuDeviceTraceXLineToStepEvents(
-            kSparseCoreIndexStart + *sc_core_id, line);
+        stream_step_events =
+            ConvertTpuDeviceTraceXLineToStepEvents(plane.Id(), line);
         IntersectCombineStepEvents(stream_step_events, &device_step_events);
       } else {
         stream_step_events =

third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc

@@ -360,7 +360,8 @@ const StatTypeMap& GetStatTypeMap() {
        {"source_stack", kSourceStack},
        {"device_offset_ps", kDeviceOffsetPs},
        {"device_duration_ps", kDeviceDurationPs},
-       {"scope_range_id", kScopeRangeId}});
+       {"scope_range_id", kScopeRangeId},
+       {"core_details", kCoreDetails}});
   DCHECK_EQ(stat_type_map->size(), kNumStatTypes);
   return *stat_type_map;
 }

third_party/xla/xla/tsl/profiler/utils/xplane_schema.h

@@ -346,7 +346,8 @@ enum StatType {
   kDeviceOffsetPs,
   kDeviceDurationPs,
   kScopeRangeId,
-  kLastStatType = kScopeRangeId,
+  kCoreDetails,
+  kLastStatType = kCoreDetails,
 };
 
 enum MegaScaleStatType : uint8_t {

third_party/xla/xla/tsl/profiler/utils/xplane_utils.cc

@@ -556,6 +556,7 @@ void AggregateXPlane(const XPlane& full_trace, XPlane& aggregated_trace) {
   const XPlaneVisitor& plane = CreateTfXPlaneVisitor(&full_trace);
   XPlaneBuilder aggregated_plane(&aggregated_trace);
   aggregated_plane.SetName(plane.Name());
+  aggregated_plane.SetId(plane.Id());
 
   uint64_t first_op_start_ps = kint64max;
   uint64_t last_op_end_ps = 0;

third_party/xla/xla/tsl/profiler/utils/xplane_utils_test.cc

@@ -397,6 +397,7 @@ TEST(XplaneUtilsTest, FindMutablePlanesWithPredicate) {
 TEST(XplaneUtilsTest, TestAggregateXPlanes) {
   XPlane xplane;
   XPlaneBuilder builder(&xplane);
+  builder.SetId(123);
   auto& event_metadata1 = *builder.GetOrCreateEventMetadata("EventMetadata1");
   auto& event_metadata2 = *builder.GetOrCreateEventMetadata("EventMetadata2");
   auto& event_metadata3 = *builder.GetOrCreateEventMetadata("EventMetadata3");
@@ -442,14 +443,16 @@ TEST(XplaneUtilsTest, TestAggregateXPlanes) {
   XPlane aggregated_xplane;
   AggregateXPlane(xplane, aggregated_xplane);
 
+  EXPECT_EQ(aggregated_xplane.id(), 123);
 // Protobuf matchers are unavailable in OSS (b/169705709)
 #if defined(PLATFORM_GOOGLE)
   // TODO(b/238349654): Proto matcher are ineffective for XPlanes.
   ASSERT_THAT(
       aggregated_xplane,
       IgnoringFields(
           {"tensorflow.profiler.XEvent.metadata_id",
-           "tensorflow.profiler.XPlane.event_metadata"},
+           "tensorflow.profiler.XPlane.event_metadata",
+           "tensorflow.profiler.XPlane.id"},
           IgnoringRepeatedFieldOrdering(EqualsProto(
               R"pb(lines {
                      id: 1