xref: /aosp_15_r20/external/pytorch/test/edge/kernel_runtime_context.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include "event_tracer.h"
4 
5 namespace torch {
6 namespace executor {
7 
8 /**
9  * Bucket type abstraction that contains many elements of runtime state that
10  * a kernel author may want available, but would otherwise be unable to access.
11  *
12  * Forwarded along to all operators when running in lean mode. NOTE: Will not be
13  * forwarded to operators if running in ATen mode as those operators do not
14  * expect to receive a KernelRuntimeContext and would not use it.
15  *
16  * This includes things like setting an error state, a scratch allocator for
17  * operators that need more then constant space, and a TensorResizer for dynamic
18  * shape tensors allowing programs to be more flexible with Tensor shape.
19  */
20 class KernelRuntimeContext {
21   public:
22   /**
23    * Construct a new kernel runtime context along with an optional event tracer.
24    */
25   KernelRuntimeContext(EventTracer* event_tracer = nullptr)
event_tracer_(event_tracer)26       : event_tracer_(event_tracer) {}
27 
28   /**
29    * INTERNAL ONLY
30    *
31    * Returns a pointer to an instance of EventTracer to do profiling/debugging
32    * logging inside the codegen layer. This is only for internal usage inside
33    * the codegen layer and users should not be accessing this.
34    */
internal_event_tracer()35   EventTracer* internal_event_tracer() {
36     return event_tracer_;
37   }
38 
39   private:
40   EventTracer* event_tracer_;
41 };
42 
43 } // namespace executor
44 } // namespace torch
45