Implement alternate profiler for MicroInterpreter.

Enable use of alternate profiler by decompression code.
Enable use of alternate profiler by Generic Benchmark application.

tensorflow/lite/micro/fake_micro_context.cc

@@ -181,4 +181,13 @@ const CompressionTensorData* FakeMicroContext::GetTensorCompressionData(
 
 #endif  // USE_TFLM_COMPRESSION
 
+TfLiteStatus FakeMicroContext::SetAlternateProfiler(
+    MicroProfilerInterface* alt_profiler) {
+  return kTfLiteError;
+}
+
+MicroProfilerInterface* FakeMicroContext::GetAlternateProfiler() const {
+  return nullptr;
+}
+
 }  // namespace tflite

tensorflow/lite/micro/fake_micro_context.h

@@ -73,6 +73,21 @@ class FakeMicroContext : public MicroContext {
 
 #endif  // USE_TFLM_COMPRESSION
 
+  // Set the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  TfLiteStatus SetAlternateProfiler(
+      MicroProfilerInterface* alt_profiler) override;
+
+  // Get the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  MicroProfilerInterface* GetAlternateProfiler() const override;
+
  private:
   static constexpr int kNumScratchBuffers_ = 12;
 

tensorflow/lite/micro/kernels/decompress.cc

@@ -22,8 +22,6 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/micro/micro_common.h"
-#include "tensorflow/lite/micro/micro_log.h"
-#include "tensorflow/lite/micro/micro_profiler.h"
 
 namespace tflite {
 

tensorflow/lite/micro/kernels/decompress.h

@@ -32,7 +32,7 @@ struct DecompressionState {
                      const size_t count_indices,
                      const CompressionTensorData& comp_data,
                      const size_t num_channels,
-                     MicroProfiler* profiler = nullptr)
+                     MicroProfilerInterface* profiler = nullptr)
       : compressed_indices_(compressed_indices),
         count_indices_(count_indices),
         comp_data_(comp_data),
@@ -79,7 +79,7 @@ struct DecompressionState {
       comp_data_.data.lut_data->use_alternate_axis
           ? 1
           : count_indices_ / num_channels_;
-  MicroProfiler* micro_profiler_;
+  MicroProfilerInterface* micro_profiler_;
 };
 
 #endif  // USE_TFLM_COMPRESSION

tensorflow/lite/micro/micro_context.cc

@@ -96,8 +96,7 @@ void* MicroContext::DecompressTensorToBuffer(
   }
 
   DecompressionState ds(static_cast<uint8_t*>(tensor.data.data), count,
-                        compression_data, num_channels,
-                        static_cast<MicroProfiler*>(external_context()));
+                        compression_data, num_channels, GetAlternateProfiler());
 
   switch (tensor.type) {
     case kTfLiteBool: {

tensorflow/lite/micro/micro_context.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/micro/micro_graph.h"
+#include "tensorflow/lite/micro/micro_profiler_interface.h"
 
 #ifdef USE_TFLM_COMPRESSION
 
@@ -125,6 +126,21 @@ class MicroContext {
 
 #endif  // USE_TFLM_COMPRESSION
 
+  // Set the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  virtual TfLiteStatus SetAlternateProfiler(
+      MicroProfilerInterface* alt_profiler) = 0;
+
+  // Get the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  virtual MicroProfilerInterface* GetAlternateProfiler() const = 0;
+
  private:
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };

tensorflow/lite/micro/micro_interpreter.cc

@@ -334,4 +334,9 @@ TfLiteStatus MicroInterpreter::SetMicroExternalContext(
   return micro_context_.set_external_context(external_context_payload);
 }
 
+TfLiteStatus MicroInterpreter::SetAlternateProfiler(
+    MicroProfilerInterface* alt_profiler) {
+  return micro_context_.SetAlternateProfiler(alt_profiler);
+}
+
 }  // namespace tflite

tensorflow/lite/micro/micro_interpreter.h

@@ -146,6 +146,14 @@ class MicroInterpreter {
     return allocator_.preserves_all_tensor();
   }
 
+  // Set the alternate MicroProfilerInterface.
+  // This value is passed through to the MicroContext.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  TfLiteStatus SetAlternateProfiler(MicroProfilerInterface* alt_profiler);
+
  protected:
   const MicroAllocator& allocator() const { return allocator_; }
   const TfLiteContext& context() const { return context_; }

tensorflow/lite/micro/micro_interpreter_context.cc

@@ -206,4 +206,14 @@ void* MicroInterpreterContext::DecompressTensorToBuffer(
 
 #endif  // USE_TFLM_COMPRESSION
 
+TfLiteStatus MicroInterpreterContext::SetAlternateProfiler(
+    tflite::MicroProfilerInterface* alt_profiler) {
+  alt_profiler_ = alt_profiler;
+  return kTfLiteOk;
+}
+
+MicroProfilerInterface* MicroInterpreterContext::GetAlternateProfiler() const {
+  return alt_profiler_;
+}
+
 }  // namespace tflite

tensorflow/lite/micro/micro_interpreter_context.h

@@ -130,6 +130,21 @@ class MicroInterpreterContext : public MicroContext {
 
 #endif  // USE_TFLM_COMPRESSION
 
+  // Set the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  TfLiteStatus SetAlternateProfiler(
+      MicroProfilerInterface* alt_profiler) override;
+
+  // Get the alternate MicroProfilerInterface.
+  // This can be used to profile subsystems simultaneously with the profiling
+  // of kernels during the Eval phase.  See (b/379584353).
+  // The alternate MicroProfilerInterface is currently used by the tensor
+  // decompression subsystem.
+  MicroProfilerInterface* GetAlternateProfiler() const override;
+
  private:
   MicroAllocator& allocator_;
   MicroInterpreterGraph& graph_;
@@ -138,6 +153,7 @@ class MicroInterpreterContext : public MicroContext {
 
   ScratchBufferHandle* scratch_buffer_handles_ = nullptr;
   void* external_context_payload_ = nullptr;
+  MicroProfilerInterface* alt_profiler_ = nullptr;
 
   TF_LITE_REMOVE_VIRTUAL_DELETE
 };

tensorflow/lite/micro/tools/benchmarking/generic_model_benchmark.cc

@@ -212,7 +212,7 @@ int Benchmark(const uint8_t* model_data, tflite::PrettyPrintType print_type) {
   profiler.ClearEvents();
 
   if (using_compression) {
-    TF_LITE_ENSURE_STATUS(interpreter.SetMicroExternalContext(&profiler2));
+    TF_LITE_ENSURE_STATUS(interpreter.SetAlternateProfiler(&profiler2));
   }
 
   MicroPrintf("");  // null MicroPrintf serves as a newline.