Internal code cleanup and refactoring.

xprof/convert/BUILD

@@ -757,6 +757,7 @@ cc_library(
     hdrs = ["op_stats_to_overview_page.h"],
     deps = [
         ":data_table_utils",
+        ":op_metrics_db_combiner",
         ":op_metrics_to_record",
         ":op_stats_to_input_pipeline_analysis",
         "@com_google_absl//absl/algorithm:container",
@@ -780,6 +781,7 @@ cc_library(
         "@org_xprof//xprof/utils:op_metrics_db_utils",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
         "@xla//xla/tsl/platform:types",
+        "@xla//xla/tsl/profiler/convert:xla_op_utils",
         "@xla//xla/tsl/profiler/utils:format_utils",
         "@xla//xla/tsl/profiler/utils:math_utils",
         "@xla//xla/tsl/profiler/utils:tf_op_utils",

xprof/convert/op_stats_to_overview_page.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -33,11 +34,13 @@ limitations under the License.
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "xla/tsl/platform/types.h"
+#include "xla/tsl/profiler/convert/xla_op_utils.h"
 #include "xla/tsl/profiler/utils/format_utils.h"
 #include "xla/tsl/profiler/utils/math_utils.h"
 #include "xla/tsl/profiler/utils/tf_op_utils.h"
 #include "xprof/convert/data_table_utils.h"
 #include "xprof/convert/op_metrics_to_record.h"
+#include "xprof/convert/op_metrics_db_combiner.h"
 #include "xprof/convert/op_stats_to_input_pipeline_analysis.h"
 #include "plugin/xprof/protobuf/hardware_types.pb.h"
 #include "plugin/xprof/protobuf/input_pipeline.pb.h"
@@ -185,7 +188,8 @@ OverviewPageRecommendation ComputeGenericRecommendation(
 }
 
 bool ComputeTpuAnalysisResult(const OpStats& op_stats,
-                              OverviewPageAnalysis* analysis) {
+                              OverviewPageAnalysis* analysis,
+                              std::optional<TpuPerformanceLimits> limits) {
   analysis->set_device_duty_cycle_percent(
       tsl::profiler::SafeDivide(
           op_stats.device_op_metrics_db().busy_time_ps(),
@@ -198,6 +202,42 @@ bool ComputeTpuAnalysisResult(const OpStats& op_stats,
   analysis->set_host_idle_time_percent(
       IdleTimeRatio(op_stats.host_op_metrics_db()) * 100.0);
 
+  if (limits.has_value() && limits->max_gigaflops_per_second > 0.0 &&
+      limits->max_gigabytes_per_second > 0.0) {
+    const tensorflow::profiler::OpMetricsDb& hlo_db_complete_steps_only =
+        op_stats.hlo_metrics_db_complete_steps_only();
+    uint64_t total_time_ps = hlo_db_complete_steps_only.total_time_ps();
+    OpMetrics aggregated;
+    for (const auto& metrics : hlo_db_complete_steps_only.metrics_db()) {
+      if (metrics.occurrences() == 0) continue;
+      if (tsl::profiler::MayHaveInnerOps(metrics.category())) continue;
+      CombineOpMetrics(metrics, &aggregated, /*update_num_cores=*/false);
+    }
+    double total_flops = aggregated.flops_v2();
+    uint64_t total_bytes_accessed =
+        BytesAccessedPerCore(aggregated, MemorySpace::MEMORY_SPACE_HBM,
+                             OpMetrics::MemoryAccessed::READ) +
+        BytesAccessedPerCore(aggregated, MemorySpace::MEMORY_SPACE_HBM,
+                             OpMetrics::MemoryAccessed::WRITE);
+    double operational_intensity_flops_per_byte =
+        tsl::profiler::SafeDivide(total_flops, total_bytes_accessed);
+    double measured_gigaflops_per_second = tsl::profiler::SafeDivide(
+        total_flops, tsl::profiler::PicoToNano(total_time_ps));
+    double measured_gigabytes_per_second = tsl::profiler::SafeDivide(
+        total_bytes_accessed, tsl::profiler::PicoToNano(total_time_ps));
+
+    double optimal_gigaflops_per_second =
+        std::min(operational_intensity_flops_per_byte *
+                     limits->max_gigabytes_per_second,
+                 limits->max_gigaflops_per_second);
+    analysis->set_flop_rate_utilization_relative_to_roofline_percent(
+        100.0 * tsl::profiler::SafeDivide(measured_gigaflops_per_second,
+                                          optimal_gigaflops_per_second));
+    analysis->set_memory_bw_utilization_relative_to_hw_limit_percent(
+        100.0 * tsl::profiler::SafeDivide(measured_gigabytes_per_second,
+                                          limits->max_gigabytes_per_second));
+  }
+
   return true;
 }
 

xprof/convert/op_stats_to_overview_page.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define XPROF_CONVERT_OP_STATS_TO_OVERVIEW_PAGE_H_
 
 #include <memory>
+#include <optional>
 #include <string>
 
 #include "absl/strings/string_view.h"
@@ -63,8 +64,14 @@ OverviewPageRecommendation ComputeGenericRecommendation(
 
 OverviewPageAnalysis ComputeAnalysisResult(const OpStats& op_stats);
 
-bool ComputeTpuAnalysisResult(const OpStats& op_stats,
-                              OverviewPageAnalysis* analysis);
+struct TpuPerformanceLimits {
+  double max_gigaflops_per_second = 0.0;
+  double max_gigabytes_per_second = 0.0;
+};
+
+bool ComputeTpuAnalysisResult(
+    const OpStats& op_stats, OverviewPageAnalysis* analysis,
+    std::optional<TpuPerformanceLimits> limits = std::nullopt);
 
 OverviewPageRunEnvironment ComputeRunEnvironment(
     const RunEnvironment& run_environment);

xprof/convert/op_stats_to_overview_page_test.cc

@@ -33,6 +33,122 @@ TEST(OpStatsToOverviewPageTest, TpuDutyCycle) {
 
   EXPECT_DOUBLE_EQ(overview_page.analysis().device_duty_cycle_percent(), 70.0);
 }
+TEST(OpStatsToOverviewPageTest, RooflineMetrics) {
+  OpStats op_stats;
+  OpMetricsDb* hlo_db = op_stats.mutable_hlo_metrics_db_complete_steps_only();
+  hlo_db->set_total_time_ps(1000000);  // 1 us
+
+  OpMetrics* metrics = hlo_db->add_metrics_db();
+  metrics->set_occurrences(1);
+  metrics->set_category("convolution");
+  metrics->set_flops_v2(1000);
+  metrics->set_time_ps(1000000);
+
+  auto* breakdown = metrics->add_memory_accessed_breakdown();
+  breakdown->set_memory_space(1);  // HBM
+  breakdown->set_operation_type(OpMetrics::MemoryAccessed::READ);
+  breakdown->set_bytes_accessed(500);
+
+  breakdown = metrics->add_memory_accessed_breakdown();
+  breakdown->set_memory_space(1);  // HBM
+  breakdown->set_operation_type(OpMetrics::MemoryAccessed::WRITE);
+  breakdown->set_bytes_accessed(500);
+
+  OverviewPageAnalysis analysis;
+  ComputeTpuAnalysisResult(op_stats, &analysis, TpuPerformanceLimits{2.0, 2.0});
+
+  EXPECT_DOUBLE_EQ(
+      analysis.flop_rate_utilization_relative_to_roofline_percent(), 50.0);
+  EXPECT_DOUBLE_EQ(
+      analysis.memory_bw_utilization_relative_to_hw_limit_percent(), 50.0);
+}
+
+TEST(OpStatsToOverviewPageTest, RooflineMetrics_FilteringAndAggregation) {
+  OpStats op_stats;
+  OpMetricsDb* hlo_db = op_stats.mutable_hlo_metrics_db_complete_steps_only();
+  hlo_db->set_total_time_ps(2000000);  // 2 us
+
+  // Valid Op 1
+  OpMetrics* metrics1 = hlo_db->add_metrics_db();
+  metrics1->set_occurrences(1);
+  metrics1->set_category("convolution");
+  metrics1->set_flops_v2(1000);
+  metrics1->set_time_ps(1000000);
+  auto* breakdown1 = metrics1->add_memory_accessed_breakdown();
+  breakdown1->set_memory_space(1);  // HBM
+  breakdown1->set_operation_type(OpMetrics::MemoryAccessed::READ);
+  breakdown1->set_bytes_accessed(500);
+  breakdown1 = metrics1->add_memory_accessed_breakdown();
+  breakdown1->set_memory_space(1);  // HBM
+  breakdown1->set_operation_type(OpMetrics::MemoryAccessed::WRITE);
+  breakdown1->set_bytes_accessed(500);
+
+  // Filtered Op (while)
+  OpMetrics* metrics2 = hlo_db->add_metrics_db();
+  metrics2->set_occurrences(1);
+  metrics2->set_category("while");
+  metrics2->set_flops_v2(5000);
+  metrics2->set_time_ps(5000000);
+
+  // Ignored Op (0 occurrences)
+  OpMetrics* metrics3 = hlo_db->add_metrics_db();
+  metrics3->set_occurrences(0);
+  metrics3->set_category("convolution");
+  metrics3->set_flops_v2(2000);
+
+  // Valid Op 2
+  OpMetrics* metrics4 = hlo_db->add_metrics_db();
+  metrics4->set_occurrences(1);
+  metrics4->set_category("convolution");
+  metrics4->set_flops_v2(1000);
+  metrics4->set_time_ps(1000000);
+  auto* breakdown4 = metrics4->add_memory_accessed_breakdown();
+  breakdown4->set_memory_space(1);  // HBM
+  breakdown4->set_operation_type(OpMetrics::MemoryAccessed::READ);
+  breakdown4->set_bytes_accessed(500);
+  breakdown4 = metrics4->add_memory_accessed_breakdown();
+  breakdown4->set_memory_space(1);  // HBM
+  breakdown4->set_operation_type(OpMetrics::MemoryAccessed::WRITE);
+  breakdown4->set_bytes_accessed(500);
+
+  OverviewPageAnalysis analysis;
+  ComputeTpuAnalysisResult(op_stats, &analysis, TpuPerformanceLimits{2.0, 2.0});
+
+  EXPECT_DOUBLE_EQ(
+      analysis.flop_rate_utilization_relative_to_roofline_percent(), 50.0);
+  EXPECT_DOUBLE_EQ(
+      analysis.memory_bw_utilization_relative_to_hw_limit_percent(), 50.0);
+}
+
+TEST(OpStatsToOverviewPageTest, RooflineMetrics_MemoryBound) {
+  OpStats op_stats;
+  OpMetricsDb* hlo_db = op_stats.mutable_hlo_metrics_db_complete_steps_only();
+  hlo_db->set_total_time_ps(1000000);  // 1 us
+
+  OpMetrics* metrics = hlo_db->add_metrics_db();
+  metrics->set_occurrences(1);
+  metrics->set_category("convolution");
+  metrics->set_flops_v2(500);  // Low FLOPs
+  metrics->set_time_ps(1000000);
+
+  auto* breakdown = metrics->add_memory_accessed_breakdown();
+  breakdown->set_memory_space(1);  // HBM
+  breakdown->set_operation_type(OpMetrics::MemoryAccessed::READ);
+  breakdown->set_bytes_accessed(1000);  // High bytes
+
+  breakdown = metrics->add_memory_accessed_breakdown();
+  breakdown->set_memory_space(1);  // HBM
+  breakdown->set_operation_type(OpMetrics::MemoryAccessed::WRITE);
+  breakdown->set_bytes_accessed(1000);
+
+  OverviewPageAnalysis analysis;
+  ComputeTpuAnalysisResult(op_stats, &analysis, TpuPerformanceLimits{2.0, 4.0});
+
+  EXPECT_DOUBLE_EQ(
+      analysis.flop_rate_utilization_relative_to_roofline_percent(), 50.0);
+  EXPECT_DOUBLE_EQ(
+      analysis.memory_bw_utilization_relative_to_hw_limit_percent(), 50.0);
+}
 
 }  // namespace
 }  // namespace profiler