Enable GPU op profiling when use_profiler is used.

LiteRT-PiperOrigin-RevId: 818772840

Author: terryheo

litert/runtime/profiler.cc

@@ -76,8 +76,8 @@ void LiteRtProfilerT::EndEvent(uint32_t event_handle) {
 }
 
 void LiteRtProfilerT::AddEvent(const char* tag, EventType event_type,
-                              uint64_t elapsed_time, int64_t event_metadata1,
-                              int64_t event_metadata2) {
+                               uint64_t metric, int64_t event_metadata1,
+                               int64_t event_metadata2) {
   if (!profiling_enabled_ || !profile_buffer_) {
     return;
   }
@@ -86,8 +86,8 @@ void LiteRtProfilerT::AddEvent(const char* tag, EventType event_type,
   std::string s_tag(tag);
   auto [it, inserted] = owned_tags_set_.insert(std::move(s_tag));
   const char* owned_tag_ptr = it->c_str();
-  profile_buffer_->AddEvent(owned_tag_ptr, event_type, elapsed_time,
-                            event_metadata1, event_metadata2);
+  profile_buffer_->AddEvent(owned_tag_ptr, event_type, metric, event_metadata1,
+                            event_metadata2);
 }
 
 void LiteRtProfilerT::StartProfiling() {

litert/runtime/profiler.h

@@ -53,8 +53,12 @@ class LiteRtProfilerT : public tflite::Profiler {
 
   // tag is copied and owned by the profiler, caller does not need to keep
   // the string alive.
-  void AddEvent(const char* tag, EventType event_type, uint64_t
-    elapsed_time, int64_t event_metadata1, int64_t event_metadata2) override;
+  // `metric` field has different intreptation based on `event_type`.
+  // e.g. it means elapsed time for [DELEGATE_]OPERATOR_INVOKE_EVENT types,
+  // and interprets as source and status code for TELEMETRY_[DELEGATE_]EVENT
+  // event types.
+  void AddEvent(const char* tag, EventType event_type, uint64_t metric,
+                int64_t event_metadata1, int64_t event_metadata2) override;
 
   // Enables profiling. Events will start being recorded.
   void StartProfiling();

litert/tools/BUILD

@@ -370,6 +370,8 @@ cc_library(
         "//tflite/c:c_api_types",
         "//tflite/c:common",
         "//tflite/profiling:model_runtime_info",
+        "//tflite/profiling:profile_buffer",
+        "//tflite/profiling:profile_summarizer",
         "//tflite/tools:command_line_flags",
         "//tflite/tools:utils",
         "//tflite/tools/benchmark:benchmark_model_lib",

litert/tools/benchmark_litert_model.cc

@@ -39,6 +39,7 @@ limitations under the License.
 #include "tflite/c/c_api_types.h"
 #include "tflite/c/common.h"
 #include "tflite/interpreter.h"
+#include "tflite/profiling/profile_summarizer.h"
 
 namespace litert::benchmark {
 namespace {
@@ -89,6 +90,12 @@ Options CreateCompiledModelOptions(const BenchmarkParams& params) {
     if (gpu_low_priority) {
       gpu_options.SetGpuPriority(kLiteRtGpuPriorityLow);
     }
+
+    auto use_profiler = params.Get<bool>("use_profiler");
+    if (use_profiler) {
+      gpu_options.SetGpuPriority(kLiteRtGpuPriorityLow);
+    }
+
     compilation_options.AddOpaqueOptions(std::move(gpu_options));
   }
 
@@ -167,25 +174,29 @@ TfLiteStatus BenchmarkLiteRtModel::Init() {
   compiled_model_ =
       std::make_unique<litert::CompiledModel>(std::move(compiled_model_result));
 
+  LiteRtCompiledModelT* compiled_model_ptr = compiled_model_->Get();
+  if (compiled_model_ptr == nullptr) {
+    LITERT_LOG(LITERT_ERROR, "Compiled model is null");
+    return kTfLiteError;
+  }
+  LITERT_ASSIGN_OR_RETURN(interpreter_, GetInterpreter(compiled_model_ptr),
+                          AsTfLiteStatus(_ << "Failed to get interpreter."));
+
   if (!params_.Get<std::string>("model_runtime_info_output_file").empty()) {
-    ::tflite::Interpreter* interpreter_ptr = nullptr;
-    LiteRtCompiledModelT* compiled_model_ptr = compiled_model_->Get();
-    if (compiled_model_ptr == nullptr) {
-      LITERT_LOG(LITERT_ERROR, "Compiled model is null");
-      return kTfLiteError;
-    }
-    LITERT_ASSIGN_OR_RETURN(interpreter_ptr, GetInterpreter(compiled_model_ptr),
-                            AsTfLiteStatus(_ << "Failed to get interpreter."));
     model_runtime_info_listener_ =
-        std::make_unique<ModelRuntimeInfoListener>(interpreter_ptr);
+        std::make_unique<ModelRuntimeInfoListener>(interpreter_);
     AddListener(model_runtime_info_listener_.get());
   }
 
   auto use_profiler = params_.Get<bool>("use_profiler");
   if (use_profiler) {
+    run_summarizer_ = std::make_unique<tflite::profiling::ProfileSummarizer>();
     LITERT_ASSIGN_OR_ABORT(profiler_, compiled_model_->GetProfiler());
     profiler_.StartProfiling();
   }
+  log_output_ =
+      std::make_unique<BenchmarkLoggingListener>(run_summarizer_.get());
+  AddListener(log_output_.get());
 
   auto signature = params_.Get<std::string>("signature_to_run_for");
   if (signature.empty()) {

litert/tools/benchmark_litert_model.h

@@ -40,6 +40,8 @@ limitations under the License.
 #include "tflite/c/common.h"
 #include "tflite/interpreter.h"
 #include "tflite/profiling/model_runtime_info.h"
+#include "tflite/profiling/profile_buffer.h"
+#include "tflite/profiling/profile_summarizer.h"
 #include "tflite/tools/benchmark/benchmark_model.h"
 #include "tflite/tools/benchmark/benchmark_params.h"
 #include "tflite/tools/benchmark/proto/benchmark_result.pb.h"
@@ -54,8 +56,13 @@ using ::tflite::tools::benchmark::BenchmarkResult;
 class BenchmarkLoggingListener : public ::tflite::benchmark::BenchmarkListener {
  private:
   std::string result_file_path_ = "";
+  tflite::profiling::ProfileSummarizer* run_summarizer_;
 
  public:
+  explicit BenchmarkLoggingListener(
+      tflite::profiling::ProfileSummarizer* run_summarizer)
+      : run_summarizer_(run_summarizer) {}
+
   void OnBenchmarkStart(
       const ::tflite::benchmark::BenchmarkParams& params) override {
     if (!params.Get<std::string>("result_file_path").empty()) {
@@ -165,6 +172,11 @@ class BenchmarkLoggingListener : public ::tflite::benchmark::BenchmarkListener {
                    result_file_path_.c_str());
       }
     }
+
+    if (run_summarizer_) {
+      LITERT_LOG(LITERT_INFO, "\n%s",
+                 run_summarizer_->GetOutputString().c_str());
+    }
   }
 };
 
@@ -206,8 +218,7 @@ using ::tflite::utils::InputTensorData;
 class BenchmarkLiteRtModel : public BenchmarkModel {
  public:
   explicit BenchmarkLiteRtModel(BenchmarkParams params = DefaultParams())
-      : BenchmarkModel(std::move(params)), log_output_() {
-    AddListener(&log_output_);
+      : BenchmarkModel(std::move(params)) {
     model_runtime_info_listener_ = nullptr;
   }
   ~BenchmarkLiteRtModel() override = default;
@@ -336,10 +347,28 @@ class BenchmarkLiteRtModel : public BenchmarkModel {
   TfLiteStatus ResetInputsAndOutputs() override {
     if (profiler_) {
       profiler_.StopProfiling();
+
       auto events = profiler_.GetEvents();
+      std::vector<std::unique_ptr<tflite::profiling::ProfileEvent>>
+          tflite_events;
+      std::vector<const tflite::profiling::ProfileEvent*> tflite_ptr_events;
+      tflite_events.reserve(events->size());
+      tflite_ptr_events.reserve(events->size());
       for (const auto& event : *events) {
-        LITERT_LOG(LITERT_INFO, "Event: %s", event.tag);
+        auto tflite_event = std::make_unique<tflite::profiling::ProfileEvent>();
+        // Refer litert/litert/runtime/profiler.cc
+        tflite_event->tag = event.tag;
+        tflite_event->begin_timestamp_us = event.start_timestamp_us;
+        tflite_event->elapsed_time = event.elapsed_time_us;
+        tflite_event->event_type = event.event_type;
+        tflite_event->event_metadata = event.event_metadata1;
+        tflite_event->extra_event_metadata = event.event_metadata2;
+        tflite_event->begin_mem_usage = event.begin_mem_usage;
+        tflite_event->end_mem_usage = event.end_mem_usage;
+        tflite_ptr_events.push_back(tflite_event.get());
+        tflite_events.push_back(std::move(tflite_event));
       }
+      run_summarizer_->ProcessProfiles(tflite_ptr_events, *interpreter_);
       profiler_.Reset();
       profiler_.StartProfiling();
     }
@@ -394,8 +423,12 @@ class BenchmarkLiteRtModel : public BenchmarkModel {
   std::unique_ptr<std::vector<litert::TensorBuffer>> input_buffers_;
   std::unique_ptr<std::vector<litert::TensorBuffer>> output_buffers_;
   litert::Profiler profiler_;
-  BenchmarkLoggingListener log_output_;
+  std::unique_ptr<BenchmarkLoggingListener> log_output_;
   std::unique_ptr<ModelRuntimeInfoListener> model_runtime_info_listener_;
+
+  // TFLite Interpreter is needed for run_summarizer_
+  ::tflite::Interpreter* interpreter_ = nullptr;
+  std::unique_ptr<tflite::profiling::ProfileSummarizer> run_summarizer_;
 };
 
 }  // namespace benchmark