Project import generated by Copybara

PiperOrigin-RevId: 929430465

Author: bmass02

xprof/convert/BUILD

@@ -1518,6 +1518,7 @@ cc_test(
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
         "@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",
         "@org_xprof//plugin/xprof/protobuf:power_metrics_proto_cc",
         "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc",

xprof/convert/op_metrics_db_combiner.cc

@@ -61,9 +61,6 @@ void CopyOpMetricsMetadata(const OpMetrics& src, OpMetrics* dst) {
   if (!dst->has_layout() && src.has_layout()) {
     *dst->mutable_layout() = src.layout();
   }
-  if (!dst->has_children() && src.has_children()) {
-    *dst->mutable_children() = src.children();
-  }
   if (!dst->has_source_info() && src.has_source_info()) {
     *dst->mutable_source_info() = src.source_info();
   }
@@ -101,6 +98,10 @@ void CombineOpMetrics(const OpMetrics& src, OpMetrics* dst,
   if (src.has_vdd_energy_j() || dst->has_vdd_energy_j()) {
     dst->set_vdd_energy_j(src.vdd_energy_j() + dst->vdd_energy_j());
   }
+  if (src.has_children()) {
+    OpMetricsDbCombiner combiner(dst->mutable_children());
+    combiner.Combine(src.children(), update_num_cores);
+  }
 }
 
 void CombineMemoryAccessedBreakdown(

xprof/convert/op_stats_combiner_test.cc

@@ -23,6 +23,7 @@ limitations under the License.
 #include "google/protobuf/util/message_differencer.h"
 #include "xla/tsl/platform/types.h"
 #include "plugin/xprof/protobuf/hardware_types.pb.h"
+#include "plugin/xprof/protobuf/op_metrics.pb.h"
 #include "plugin/xprof/protobuf/op_stats.pb.h"
 #include "plugin/xprof/protobuf/power_metrics.pb.h"
 #include "plugin/xprof/protobuf/steps_db.pb.h"
@@ -234,6 +235,74 @@ TEST(CombineAllOpStatsTest, CombineDisaggregatedServingStats) {
               1e-6);
 }
 
+TEST(CombineAllOpStatsTest, CombineDeviceOpMetricsDbWithChildren) {
+  OpStats dst_op_stats, op_stats_1, op_stats_2;
+
+  // Set up op_stats_1 with one device op that has children.
+  {
+    auto* db = op_stats_1.mutable_device_op_metrics_db();
+    auto* op = db->add_metrics_db();
+    op->set_name("parent_op");
+    op->set_hlo_module_id(123);
+    op->set_occurrences(1);
+    op->set_time_ps(100);
+
+    auto* child_db = op->mutable_children();
+    auto* child_op = child_db->add_metrics_db();
+    child_op->set_name("Add_1");
+    child_op->set_hlo_module_id(123);
+    child_op->set_occurrences(5);
+    child_op->set_time_ps(400);
+  }
+
+  // Set up op_stats_2 with the same device op that has other children.
+  {
+    auto* db = op_stats_2.mutable_device_op_metrics_db();
+    auto* op = db->add_metrics_db();
+    op->set_name("parent_op");
+    op->set_hlo_module_id(123);
+    op->set_occurrences(2);
+    op->set_time_ps(200);
+
+    auto* child_db = op->mutable_children();
+    auto* child_op = child_db->add_metrics_db();
+    child_op->set_name("Add_1");
+    child_op->set_hlo_module_id(123);
+    child_op->set_occurrences(10);
+    child_op->set_time_ps(800);
+  }
+
+  OpStatsInfo op_stats_info_1(&op_stats_1, TPU, 0);
+  OpStatsInfo op_stats_info_2(&op_stats_2, TPU, 1);
+
+  std::vector<OpStatsInfo> all_op_stats_info = {op_stats_info_1,
+                                                op_stats_info_2};
+
+  StepDatabaseResult dummy_step_db_result;
+  absl::flat_hash_map<uint32_t /*host_id*/, const StepDatabaseResult*> result;
+  result.insert({0, &dummy_step_db_result});
+  StepIntersection dummy_step_intersection = StepIntersection(1, result);
+
+  CombineAllOpStats(all_op_stats_info, dummy_step_intersection, &dst_op_stats);
+
+  // Verify that the combined op metrics has the parent_op.
+  const auto& combined_db = dst_op_stats.device_op_metrics_db();
+  ASSERT_EQ(combined_db.metrics_db_size(), 1);
+  const auto& combined_op = combined_db.metrics_db(0);
+  EXPECT_EQ(combined_op.name(), "parent_op");
+  EXPECT_EQ(combined_op.occurrences(), 3);
+  EXPECT_EQ(combined_op.time_ps(), 300);
+
+  // Verify that the children are merged without duplication.
+  const auto& combined_children_db = combined_op.children();
+  ASSERT_EQ(combined_children_db.metrics_db_size(), 1);
+
+  const auto& child_op = combined_children_db.metrics_db(0);
+  EXPECT_EQ(child_op.name(), "Add_1");
+  EXPECT_EQ(child_op.occurrences(), 15);
+  EXPECT_EQ(child_op.time_ps(), 1200);
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tensorflow

xprof/convert/smart_suggestion/BUILD

@@ -19,7 +19,6 @@ cc_library(
         ":tool_data_provider",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -205,7 +204,6 @@ cc_test(
         ":mock_tool_data_provider",
         ":signal_provider",
         ":sparse_core_offload_rule",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_googletest//:gtest_main",
         "@org_xprof//plugin/xprof/protobuf:memory_profile_proto_cc",

xprof/utils/op_metrics_db_utils.cc

@@ -336,7 +336,9 @@ void AdjustFlopsAndBytesAccessed(OpMetrics& op_metrics) {
 
 OpMetricsDbBuilder::OpMetricsDbBuilder(OpMetricsDb* db) : db_(db) {
   DCHECK_NE(db_, nullptr);
-  DCHECK_EQ(db_->metrics_db_size(), db->metrics_db_size());
+  for (OpMetrics& metrics : *db_->mutable_metrics_db()) {
+    op_metrics_map_[metrics.hlo_module_id()][metrics.name()] = &metrics;
+  }
 }
 
 OpMetrics* OpMetricsDbBuilder::LookupOrInsertNewOpMetrics(

xprof/utils/op_metrics_db_utils.h

@@ -45,7 +45,6 @@ TF_CONST_INIT extern const uint32_t kSparseCoreIndexStart;
 class OpMetricsDbBuilder {
  public:
   // Create with a borrowed op database.
-  // REQUIRED: The op database must be empty.
   explicit OpMetricsDbBuilder(OpMetricsDb* db);
 
  protected: