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