[Profiler] Add disabled and global methods to ProfilerConfig. (pytorch#83891)

`ProfilerState::Disabled` and `ProfilerState::KINETO_ONDEMAND` have special semantics. The former is somewhat intuitive, but the degree of behavior branching on the latter (and why the branching is necessary) is less clear. By factoring the enum checks into methods, we can both clairify intent and future proof in case we ever add other global profiling contexts.

Differential Revision: [D38917980](https://our.internmc.facebook.com/intern/diff/D38917980/)
Pull Request resolved: pytorch#83891
Approved by: https://github.com/slgong-fb

Author: Taylor Robie

torch/csrc/autograd/profiler_kineto.cpp

@@ -226,7 +226,7 @@ struct KinetoThreadLocalState : public ProfilerThreadLocalStateBase {
       int64_t total_allocated,
       int64_t total_reserved,
       c10::Device device) override {
-    if (config_.profile_memory && config_.state != ProfilerState::Disabled) {
+    if (config_.profile_memory && !config_.disabled()) {
       record_queue_.getSubqueue()->emplace_allocation_event(
           torch::profiler::impl::getApproximateTime(),
           ptr,
@@ -243,7 +243,7 @@ struct KinetoThreadLocalState : public ProfilerThreadLocalStateBase {
       int64_t total_allocated,
       int64_t total_reserved,
       c10::Device device) override {
-    if (config_.profile_memory && config_.state != ProfilerState::Disabled) {
+    if (config_.profile_memory && !config_.disabled()) {
       record_queue_.getSubqueue()->emplace_ooms_event(
           torch::profiler::impl::getApproximateTime(),
           alloc_size,
@@ -558,13 +558,18 @@ void enableProfiler(
 
   TORCH_CHECK(
       config.state == ProfilerState::KINETO ||
-      config.state == ProfilerState::KINETO_GPU_FALLBACK ||
-      config.state == ProfilerState::KINETO_ONDEMAND);
+      config.state == ProfilerState::KINETO_GPU_FALLBACK || config.global());
   TORCH_CHECK(
       !activities.empty(), "No activities specified for Kineto profiler");
 
-  if (config.state == ProfilerState::KINETO ||
-      config.state == ProfilerState::KINETO_GPU_FALLBACK) {
+  if (config.global()) {
+    KinetoTLSGlobalStateManager::init(config, activities);
+
+    TORCH_INTERNAL_ASSERT(
+        activities.count(ActivityType::CPU),
+        "Ondemand profiling must enable CPU tracing");
+    pushProfilingCallbacks<true>(scopes);
+  } else {
     auto state = std::make_shared<KinetoThreadLocalState>(config, activities);
     c10::ThreadLocalDebugInfo::_push(c10::DebugInfoKind::PROFILER_STATE, state);
 
@@ -573,15 +578,6 @@ void enableProfiler(
     }
     torch::profiler::impl::kineto::startTrace();
   }
-
-  if (config.state == ProfilerState::KINETO_ONDEMAND) {
-    KinetoTLSGlobalStateManager::init(config, activities);
-
-    TORCH_INTERNAL_ASSERT(
-        activities.count(ActivityType::CPU),
-        "Ondemand profiling must enable CPU tracing");
-    pushProfilingCallbacks<true>(scopes);
-  }
 }
 
 std::unique_ptr<ProfilerResult> disableProfiler() {

torch/csrc/autograd/profiler_legacy.cpp

@@ -193,7 +193,7 @@ thread_event_lists ProfilerLegacyThreadLocalState::consolidate() {
 }
 
 void ProfilerLegacyThreadLocalState::mark(std::string name, bool include_cuda) {
-  if (config_.state == torch::profiler::impl::ProfilerState::Disabled) {
+  if (config_.disabled()) {
     return;
   }
   if (config_.state == torch::profiler::impl::ProfilerState::NVTX) {
@@ -225,7 +225,7 @@ void ProfilerLegacyThreadLocalState::pushRange(
     const at::RecordFunction& fn,
     const bool record_cuda,
     std::vector<std::vector<int64_t>>&& shapes) {
-  if (config_.state == torch::profiler::impl::ProfilerState::Disabled) {
+  if (config_.disabled()) {
     return;
   }
   if (config_.state == torch::profiler::impl::ProfilerState::NVTX) {
@@ -273,7 +273,7 @@ void ProfilerLegacyThreadLocalState::pushRange(
 void ProfilerLegacyThreadLocalState::popRange(
     const at::RecordFunction& fn,
     const bool record_cuda) {
-  if (config_.state == torch::profiler::impl::ProfilerState::Disabled) {
+  if (config_.disabled()) {
     return;
   }
   if (config_.state == torch::profiler::impl::ProfilerState::NVTX) {
@@ -300,8 +300,7 @@ void ProfilerLegacyThreadLocalState::reportMemoryUsage(
     int64_t /* total_allocated, unused for legacy */,
     int64_t /* total_reserved, unused for legacy */,
     c10::Device device) {
-  if (config_.profile_memory &&
-      config_.state != torch::profiler::impl::ProfilerState::Disabled) {
+  if (config_.profile_memory && !config_.disabled()) {
     uint64_t thread_id = at::RecordFunction::currentThreadId();
     LegacyEvent evt(
         EventKind::MemoryAlloc,
@@ -372,9 +371,7 @@ void pushProfilingCallbacksLegacy() {
           [](const at::RecordFunction& fn)
               -> std::unique_ptr<at::ObserverContext> {
             auto state_ptr = ProfilerLegacyThreadLocalState::getTLS();
-            if (!state_ptr ||
-                state_ptr->config().state ==
-                    torch::profiler::impl::ProfilerState::Disabled) {
+            if (!state_ptr || state_ptr->config().disabled()) {
               return nullptr;
             }
             bool record_cuda = state_ptr->config().state ==
@@ -396,9 +393,7 @@ void pushProfilingCallbacksLegacy() {
           },
           [](const at::RecordFunction& fn, at::ObserverContext*) {
             auto state_ptr = ProfilerLegacyThreadLocalState::getTLS();
-            if (!state_ptr ||
-                state_ptr->config().state ==
-                    torch::profiler::impl::ProfilerState::Disabled) {
+            if (!state_ptr || state_ptr->config().disabled()) {
               return;
             }
             bool record_cuda = state_ptr->config().state ==
@@ -454,9 +449,7 @@ thread_event_lists disableProfilerLegacy(
 
   auto state_ptr = static_cast<ProfilerLegacyThreadLocalState*>(state.get());
   TORCH_CHECK(
-      state_ptr &&
-          state_ptr->config().state !=
-              torch::profiler::impl::ProfilerState::Disabled,
+      state_ptr && !state_ptr->config().disabled(),
       "Can't disable profiler when it's not running");
 
   if (cleanupTLSState) {

torch/csrc/profiler/api.cpp

@@ -14,6 +14,14 @@ ExperimentalConfig::ExperimentalConfig(
   return !profiler_metrics.empty();
 }
 
+bool ProfilerConfig::disabled() const {
+  return state == torch::profiler::impl::ProfilerState::Disabled;
+}
+
+bool ProfilerConfig::global() const {
+  return state == torch::profiler::impl::ProfilerState::KINETO_ONDEMAND;
+}
+
 namespace {
 enum ProfilerIValueIdx {
   STATE = 0,
@@ -52,9 +60,7 @@ ProfilerConfig ProfilerConfig::fromIValue(
 
 bool profilerEnabled() {
   auto state_ptr = ProfilerThreadLocalStateBase::getTLS();
-  return state_ptr &&
-      state_ptr->config().state !=
-      torch::profiler::impl::ProfilerState::Disabled;
+  return state_ptr && !state_ptr->config().disabled();
 }
 
 TORCH_API ActiveProfilerType profilerType() {

torch/csrc/profiler/api.h

@@ -66,6 +66,10 @@ struct TORCH_API ProfilerConfig {
         with_flops(with_flops),
         with_modules(with_modules) {}
   ~ProfilerConfig() = default;
+
+  bool disabled() const;
+  bool global() const;
+
   ProfilerState state;
   ExperimentalConfig experimental_config;
   bool report_input_shapes;

torch/csrc/profiler/collection.cpp

@@ -836,7 +836,7 @@ trace_ptr_t addKinetoEvents(
   passEventsToKineto(results, start_time_us, end_time_us);
 
   // In on demand mode kineto is directly controlled by other machinery.
-  if (config.state == ProfilerState::KINETO_ONDEMAND) {
+  if (config.global()) {
     return nullptr;
   }