Input Pipeline Analysis: Flat Host Metrics Proto Changes

plugin/xprof/protobuf/BUILD

@@ -159,7 +159,10 @@ xprof_proto_library(
 xprof_proto_library(
     name = "steps_db_proto",
     srcs = ["steps_db.proto"],
-    protodeps = [":op_metrics_proto"],
+    protodeps = [
+        ":op_metrics_proto",
+        ":flat_op_metrics_proto",
+    ],
 )
 
 xprof_proto_library(

plugin/xprof/protobuf/flat_op_metrics.proto

@@ -160,6 +160,9 @@ message FlatOpMetricsDb {
 
   // Weighted average total op time if running on default pstate.
   uint64 normalized_total_op_time_ps = 7;
+
+  int64 total_host_infeed_enq_duration_ps = 8;
+  int64 total_host_infeed_enq_start_timestamp_ps_diff = 9;
 }
 
 message FlatOpMetricsDbMap {

plugin/xprof/protobuf/op_stats.proto

@@ -13,7 +13,7 @@ import "plugin/xprof/protobuf/steps_db.proto";
 import "plugin/xprof/protobuf/tf_function.proto";
 import "plugin/xprof/protobuf/topology.proto";
 
-// Performance environment, e.g the peak performance capabilities of the device.
+// Performance environment, e.g. the peak performance capabilities of the device.
 message PerfEnv {
   // Peak performance of a TPU core or a GPU in TFLOP/s.
   double peak_tera_flops_per_second = 1;
@@ -189,8 +189,13 @@ message OpStats {
   SourceStats source_stats = 14;
   // The result for the disaggregated serving latency.
   DisaggregatedServingLatency disaggregated_serving_latency = 15;
-  // The new path for device op metrics. device_op_metrics_db will be deprecated
-  // soon.
+  // The new path for device op metrics. TODO(b/529857396): device_op_metrics_db
+  // will be deprecated.
   tensorflow.profiler.FlatOpMetricsDb flat_device_op_metrics_db = 16;
+  // The new path for hlo metrics. TODO(b/529857396):
+  // hlo_metrics_db_complete_steps_only will be deprecated.
+  tensorflow.profiler.FlatOpMetricsDb flat_hlo_metrics_db_complete_steps_only =
+      17;
+  optional tensorflow.profiler.FlatOpMetricsDb flat_host_op_metrics_db = 18;
   reserved 7;
 }

plugin/xprof/protobuf/steps_db.proto

@@ -4,6 +4,7 @@ package tensorflow.profiler;
 
 import "google/protobuf/any.proto";
 import "plugin/xprof/protobuf/op_metrics.proto";
+import "plugin/xprof/protobuf/flat_op_metrics.proto";
 
 // Breakdown of step-time on generic hardware. Note that these components are
 // mutually exclusive so that adding them together is equal to the step time. If
@@ -128,7 +129,7 @@ message StepInfoResult {
   uint64 begin_ps = 3;
   // Breakdown of the step-time. Can be unpacked into a GenericStepBreakdown.
   google.protobuf.Any step_breakdown = 4;
-  // Total time/bytes/occurences for collectives. (All-Reduce, All-to-All etc)
+  // Total time/bytes/occurrences for collectives. (All-Reduce, All-to-All etc)
   DeviceMemoryTransfer collectives = 6;
 }
 
@@ -168,13 +169,16 @@ message PerCoreStepInfo {
   map<uint32, uint32> core_id_to_replica_id_map = 5;
   // A map from core_id to all-reduce ops.
   map<uint32, AllReduceDbResult> all_reduce_db_per_core = 6;
-  // Information about deivce memory transfers, categoried by source and
+  // Information about device memory transfers, categoried by source and
   // destination. Ordered by following categories:
   // 1. HostToDevice
   // 2. DeviceToHost
   // 3. DeviceToDevice
   // Cores are normally sharing host interfaces (i.e. PCIe).
   repeated DeviceMemoryTransfer device_memory_transfers = 7;
+  // The new path for per-step HLO metrics. hlo_metrics_db will be deprecated
+  // soon.
+  tensorflow.profiler.FlatOpMetricsDb flat_hlo_metrics_db = 8;
 
   reserved 4;
 }

xprof/convert/BUILD

@@ -942,6 +942,35 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "flat_op_metrics_to_record",
+    srcs = ["flat_op_metrics_to_record.cc"],
+    hdrs = ["flat_op_metrics_to_record.h"],
+    deps = [
+        ":op_metrics_to_record",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc",
+        "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc",
+        "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc",
+        "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc",
+        "@xla//xla/tsl/profiler/utils:math_utils",
+    ],
+)
+
+cc_test(
+    name = "flat_op_metrics_to_record_test",
+    srcs = ["flat_op_metrics_to_record_test.cc"],
+    deps = [
+        ":flat_op_metrics_to_record",
+        "//net/proto2/contrib/parse_proto:parse_text_proto",
+        "@com_google_googletest//:gtest_main",
+        "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc",
+        "@org_xprof//plugin/xprof/protobuf:roofline_model_proto_cc",
+    ],
+)
+
 cc_library(
     name = "op_stack",
     hdrs = ["op_stack.h"],
@@ -995,11 +1024,15 @@ cc_library(
     hdrs = ["op_stats_to_roofline_model.h"],
     deps = [
         ":data_table_utils",
+        ":flat_op_metrics_db_combiner",
+        ":flat_op_metrics_to_record",
         ":op_metrics_db_combiner",
         ":op_metrics_to_record",
         ":source_info_utils",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings",
+        "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc",
@@ -1272,13 +1305,15 @@ cc_library(
     srcs = ["step_events_to_steps_db.cc"],
     hdrs = ["step_events_to_steps_db.h"],
     deps = [
+        ":flat_op_metrics_db_combiner",
         ":op_metrics_db_combiner",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc",
         "@org_xprof//xprof/utils:event_span",
+        "@org_xprof//xprof/utils:flat_op_metrics_db_utils",
         "@org_xprof//xprof/utils:op_metrics_db_utils",
         "@xla//xla/tsl/lib/gtl:map_util",
         "@xla//xla/tsl/platform:logging",
@@ -1331,6 +1366,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -1440,9 +1476,9 @@ cc_library(
         "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc",
         "@org_xprof//xprof/utils:event_span",
+        "@org_xprof//xprof/utils:flat_op_metrics_db_utils",
         "@org_xprof//xprof/utils:op_metrics_db_utils",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
-        "@xla//xla/tsl/platform:types",
         "@xla//xla/tsl/profiler/utils:tf_op_utils",
         "@xla//xla/tsl/profiler/utils:tf_xplane_visitor",
         "@xla//xla/tsl/profiler/utils:timespan",
@@ -1641,6 +1677,7 @@ cc_test(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
+        "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:power_metrics_proto_cc",
@@ -3035,6 +3072,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
         "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc",
         "@org_xprof//xprof/utils:flat_op_metrics_db_utils",
         "@org_xprof//xprof/utils:gpu_event_stats",

xprof/convert/flat_op_metrics_db_combiner.h

@@ -35,6 +35,11 @@ class FlatOpMetricsDbCombiner : public FlatOpMetricsDbBuilder {
   // occurs.
   void Combine(const FlatOpMetricsDb& src, bool update_num_cores = true);
 
+  // Combines the memory access breakdown for FlatOpMetrics.
+  static void CombineMemoryAccessedBreakdown(
+      const tsl::protobuf::RepeatedPtrField<FlatOpMetrics::MemoryAccessed>& src,
+      tsl::protobuf::RepeatedPtrField<FlatOpMetrics::MemoryAccessed>* dst);
+
  private:
   // Copies FlatOpMetrics metadata (e.g., category, provenance) from src to dst.
   static void CopyFlatOpMetricsMetadata(const FlatOpMetrics& src,
@@ -43,11 +48,6 @@ class FlatOpMetricsDbCombiner : public FlatOpMetricsDbBuilder {
   // Combines FlatOpMetrics data from src into dst.
   static void CombineFlatOpMetrics(const FlatOpMetrics& src, FlatOpMetrics* dst,
                                    bool update_num_cores);
-
-  // Combines the memory access breakdown for FlatOpMetrics.
-  static void CombineMemoryAccessedBreakdown(
-      const tsl::protobuf::RepeatedPtrField<FlatOpMetrics::MemoryAccessed>& src,
-      tsl::protobuf::RepeatedPtrField<FlatOpMetrics::MemoryAccessed>* dst);
 };
 
 }  // namespace profiler

xprof/convert/flat_op_metrics_to_record.cc

@@ -0,0 +1,58 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xprof/convert/flat_op_metrics_to_record.h"
+
+#include <tuple>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+
+namespace tensorflow {
+namespace profiler {
+
+
+std::vector<const FlatOpMetrics*> SortedOpMetricsDb(
+    const FlatOpMetricsDb& metrics_db, int max_records) {
+  std::vector<const FlatOpMetrics*> result;
+  // Exclude SparseCore ops to maintain parity with legacy OpMetricsDb behavior.
+  // SparseCore operations are typically aggregated as children of TensorCore
+  // operations and are not listed as independent top-level records in this
+  // view.
+  result.reserve(absl::c_count_if(
+      metrics_db.op_instances(), [](const FlatOpMetrics& metrics) {
+        return metrics.core_type() != FlatOpMetrics::SPARSE_CORE;
+      }));
+  for (const FlatOpMetrics& metrics : metrics_db.op_instances()) {
+    if (metrics.core_type() == FlatOpMetrics::SPARSE_CORE) continue;
+    result.push_back(&metrics);
+  }
+
+  auto comp = [](const FlatOpMetrics* a, const FlatOpMetrics* b) {
+    return std::make_tuple(a->self_time_ps(), a->hlo_name()) >
+           std::make_tuple(b->self_time_ps(), b->hlo_name());
+  };
+  int result_size = result.size();
+  if (max_records > 0 && result_size > max_records) {
+    absl::c_partial_sort(result, result.begin() + max_records, comp);
+    result.resize(max_records);
+  } else {
+    absl::c_sort(result, comp);
+  }
+  return result;
+}
+
+}  // namespace profiler
+}  // namespace tensorflow