Add support for event_tracer in KernelContext and codegen layer

Differential Revision: D48975975

fbshipit-source-id: 033fa80537557245f59806232960efa6d0222b69

Author: tarun292

runtime/core/event_tracer_hooks.h

@@ -114,17 +114,16 @@ inline void event_tracer_create_event_block(
 inline EventTracerEntry event_tracer_begin_profiling_event(
     EventTracer* event_tracer,
     char const* name) {
-  EventTracerEntry event_tracer_entry;
 #ifdef ET_EVENT_TRACER_ENABLED
   if (event_tracer) {
-    event_tracer_entry = event_tracer->start_profiling(name);
+    return event_tracer->start_profiling(name);
   }
 #else //! ET_EVENT_TRACER_ENABLED
   (void)event_tracer;
   (void)name;
 #endif
   // There is no active tracer; this value will be ignored.
-  return event_tracer_entry;
+  return EventTracerEntry();
 }
 
 /**

runtime/executor/method.cpp

@@ -905,7 +905,7 @@ Error Method::execute_instruction() {
       EXECUTORCH_SCOPE_PROF("OPERATOR_CALL");
       // TODO(T147221312): Also expose the temp allocator and tensor resizer
       // via the context.
-      KernelRuntimeContext context;
+      KernelRuntimeContext context(event_tracer_);
       auto args = chain.argument_lists_[step_state_.instr_idx];
       chain.kernels_[step_state_.instr_idx](context, args.data());
       Error err = context.failure_state();

runtime/kernel/kernel_runtime_context.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/event_tracer_hooks.h>
 #include <executorch/runtime/platform/compiler.h>
 
 namespace torch {
@@ -22,6 +23,11 @@ namespace executor {
  */
 class KernelRuntimeContext {
  public:
+  /**
+   * Construct a new kernel runtime context along with an optional event tracer.
+   */
+  KernelRuntimeContext(EventTracer* event_tracer = nullptr)
+      : event_tracer_(event_tracer) {}
   /**
    * Tells the runtime that the kernel call has failed. Prefer this over
    * ET_CHECK_*(), which fatally panics the process/system.
@@ -43,11 +49,23 @@ class KernelRuntimeContext {
     return failure_state_;
   }
 
+  /**
+   * INTERNAL ONLY
+   *
+   * Returns a pointer to an instance of EventTracer to do profiling/debugging
+   * logging inside the codegen layer. This is only for internal usage inside
+   * the codegen layer and users should not be accessing this.
+   */
+  EventTracer* internal_event_tracer() {
+    return event_tracer_;
+  }
+
   // TODO(T147221312): Add a way to allocate temporary memory.
 
   // TODO(T147221312): Add a way to resize a tensor.
 
  private:
+  EventTracer* event_tracer_;
   Error failure_state_ = Error::Ok;
 };