Add support for an operator level thread local observer (pytorch#108822)

Summary: Add support for an operator level thread local observer

Test Plan: Verified the interception as part of a pytorch model evaluation with static runtime.

Differential Revision: D49082250

Pull Request resolved: pytorch#108822
Approved by: https://github.com/davidberard98

Author: rohithmenon

torch/csrc/jit/runtime/static/impl.cpp

@@ -760,6 +760,31 @@ size_t StaticModule::prepareStaticNodeInfos(
   return node_idx - node_start;
 }
 
+#ifdef FBCODE_CAFFE2
+thread_local SROperatorObserver* tlsOpObserver = nullptr;
+
+void SROperatorObserver::setCurrentThreadObserver(
+    SROperatorObserver* observer) {
+  tlsOpObserver = observer;
+}
+
+SROperatorObserver* SROperatorObserver::getCurrentThreadObserver() {
+  return tlsOpObserver;
+}
+
+void SROperatorObserver::onStart(const Node* node) {
+  if (tlsOpObserver != nullptr && tlsOpObserver->startCb != nullptr) {
+    tlsOpObserver->startCb(node);
+  }
+}
+
+void SROperatorObserver::onEnd(const Node* node) {
+  if (tlsOpObserver != nullptr && tlsOpObserver->endCb != nullptr) {
+    tlsOpObserver->endCb(node);
+  }
+}
+#endif // FBCODE_CAFFE2
+
 BlockInfo::BlockInfo(uint32_t input_idx, Block& block)
     : input_idx_(input_idx), block_(block) {}
 
@@ -2052,6 +2077,9 @@ std::vector<IValue> ProcessedNode::inputs_ivalue_vec() const {
 }
 
 void ProcessedNode::run() {
+#ifdef FBCODE_CAFFE2
+  SROperatorObserver::onStart(node());
+#endif
 #ifndef PYTORCH_DISABLE_PER_OP_PROFILING
   auto step_callbacks =
       at::getStepCallbacksUnlessEmpty(at::RecordScope::STATIC_RUNTIME_OP);
@@ -2085,6 +2113,9 @@ void ProcessedNode::run() {
     DCHECK(verify_no_memory_overlap());
   }
 #endif
+#ifdef FBCODE_CAFFE2
+  SROperatorObserver::onEnd(node());
+#endif
 }
 
 static bool checkNoMemoryOverlap(const at::Tensor& a, const at::Tensor& b) {

torch/csrc/jit/runtime/static/impl.h

@@ -248,6 +248,19 @@ class StaticRuntime;
 
 using SROperator = std::function<void(ProcessedNode*)>;
 
+#ifdef FBCODE_CAFFE2
+struct TORCH_API SROperatorObserver {
+  using OperatorCallback = void (*)(const Node*);
+  OperatorCallback startCb = nullptr;
+  OperatorCallback endCb = nullptr;
+
+  static void setCurrentThreadObserver(SROperatorObserver* observer);
+  static SROperatorObserver* getCurrentThreadObserver();
+  static void onStart(const Node* name);
+  static void onEnd(const Node* name);
+};
+#endif
+
 // A `BlockInfo` instance stores all of the shared state that each
 // `BlockRunner` will need to access. Most of this information is
 // read-only and shared between threads.