Add is_hidden_event method to KinetoEvent Python interface (#155214)

Fixes #155213

Pull Request resolved: https://github.com/pytorch/pytorch/pull/155214
Approved by: https://github.com/sraikund16

torch/_C/_autograd.pyi

@@ -77,6 +77,7 @@ class _KinetoEvent:
     def cuda_elapsed_us(self) -> int: ...
     def privateuse1_elapsed_us(self) -> int: ...
     def is_user_annotation(self) -> bool: ...
+    def is_hidden_event(self) -> bool: ...
 
 class _ProfilerResult:
     def events(self) -> list[_KinetoEvent]: ...

torch/autograd/profiler.py

@@ -582,7 +582,10 @@ class profile:
         device_corr_map: dict[int, list[FunctionEvent]] = {}
         max_evt_id = 0
         for kineto_event in result.events():
-            if _filter_name(kineto_event.name()):
+            if (
+                _filter_name(kineto_event.name())
+                or getattr(kineto_event, "is_hidden_event", lambda: False)()
+            ):
                 continue
             rel_start_ns = kineto_event.start_ns() - trace_start_ns
             rel_end_ns = kineto_event.end_ns() - trace_start_ns

torch/csrc/autograd/init.cpp

@@ -307,7 +307,11 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject* unused) {
                 e.activityType() ==
                 (uint8_t)libkineto::ActivityType::GPU_USER_ANNOTATION;
           })
-      .def("nbytes", [](const KinetoEvent& e) { return e.nBytes(); });
+      .def("nbytes", [](const KinetoEvent& e) { return e.nBytes(); })
+      // whether the event is hidden
+      .def("is_hidden_event", [](const KinetoEvent& e) {
+        return e.isHiddenEvent();
+      });
 
   m.def("_soft_assert_raises", &setSoftAssertRaises);
   m.def("_get_sequence_nr", &at::sequence_number::peek);

torch/csrc/autograd/profiler_kineto.cpp

@@ -936,6 +936,10 @@ bool KinetoEvent::hasKwinputs() const {
   return !kwinputs_.empty();
 }
 
+bool KinetoEvent::isHiddenEvent() const {
+  return result_ && result_->hidden_;
+}
+
 const std::unordered_map<std::string, c10::IValue> KinetoEvent::kwinputs()
     const {
   return kwinputs_;

torch/csrc/autograd/profiler_kineto.h

@@ -37,6 +37,7 @@ struct TORCH_API KinetoEvent {
   bool hasConcreteInputs() const;
   const c10::ArrayRef<c10::IValue> concreteInputs() const;
   bool hasKwinputs() const;
+  bool isHiddenEvent() const;
   const std::unordered_map<std::string, c10::IValue> kwinputs() const;
   uint64_t flops() const;
   int64_t sequenceNr() const;

torch/csrc/profiler/collection.cpp

@@ -1016,6 +1016,12 @@ class TransferEvents {
     }
   }
 
+  bool isHiddenEvent(const itrace_t* activity) const {
+    TORCH_INTERNAL_ASSERT(activity != nullptr);
+    // Kineto uses "hidden" metadata to mark events that should be hidden.
+    return activity->getMetadataValue("hidden") == "1";
+  }
+
   std::shared_ptr<Result> resultFromActivity(const itrace_t* activity) {
     TORCH_INTERNAL_ASSERT(activity != nullptr);
 
@@ -1036,7 +1042,7 @@ class TransferEvents {
             {/*id=*/static_cast<uint32_t>(activity->flowId()),
              /*type=*/static_cast<uint32_t>(activity->flowType()),
              /*start=*/activity->flowStart()}});
-
+    event->hidden_ = isHiddenEvent(activity);
     // NB: It's tempting to set `event->kineto_activity_`; however we can only
     // guarantee that the events we passed to Kineto are of type
     // `GenericTraceActivity`. Others may derive from ITraceActivity and thus

torch/csrc/profiler/collection.h

@@ -421,7 +421,7 @@ struct TORCH_API Result : public std::enable_shared_from_this<Result> {
   std::weak_ptr<Result> parent_;
   std::vector<std::shared_ptr<Result>> children_;
   bool finished_{false};
-
+  bool hidden_{false};
   const torch::profiler::impl::kineto::activity_t* kineto_activity_{nullptr};
 
  private: