LocalCallbackManager Class — pytorch Architecture

Source Code

aten/src/ATen/record_function.cpp lines 167–207

class LocalCallbackManager {
 public:
  static LocalCallbackManager& get(); // Singleton

 private:
  LocalCallbackManager();

 public:
  const RecordFunctionTLS& getTLS() const;
  StepCallbacks getActiveCallbacks(const RecordScope scope);
  std::optional<StepCallbacks> getActiveCallbacksUnlessEmpty(
      const RecordScope scope);

  void setTLS(const RecordFunctionTLS& tls);
  void seed(uint32_t seed);
  CallbackHandle addCallback(RecordFunctionCallback callback);
  bool setCallbackEnabled(CallbackHandle handle, bool enabled);
  bool removeCallback(CallbackHandle handle);
  void clearCallbacks();

 private:
  void rebuildActiveCallbacksIfNeeded();

  void rebuild_all(const GlobalCallbackManager::snapshot_t& global_snapshot);

  void rebuild_callback_scopes(
      const GlobalCallbackManager::snapshot_t& global_snapshot,
      const RecordFunctionCallback& callback);

  void rebuild_scope(
      const GlobalCallbackManager::snapshot_t& global_snapshot,
      const RecordScope scope);

  // Source of truth.
  RecordFunctionTLS registered_callbacks_;

  // Runtime cache.
  size_t global_version_{GlobalCallbackManager::NoVersion};
  std::array<CacheEntry, NumRecordScopes> active_callbacks_;
  std::mt19937 generator_;
};