CacheEntry Class — pytorch Architecture

Source Code

aten/src/ATen/record_function.cpp lines 127–165

class CacheEntry {
 public:
  CacheEntry() = default;
  CacheEntry(std::mt19937* generator, RecordScope scope);

  // The caller is expected to check `GlobalCallbackManager::get().version()'
  // and call CacheEntry::update() if necessary.
  StepCallbacks getActiveCallbacks();
  std::optional<StepCallbacks> getActiveCallbacksUnlessEmpty();

  // Full rebuild. (E.g. during registration)
  void update(const std::vector<RecordFunctionCallback>& callbacks);

 private:
  struct CallbackAndCounter {
    RecordFunctionCallback callback_;

    // `-1` indicates that a callback is not sampled.
    int tries_left_{-1};
  };

  C10_ALWAYS_INLINE void getActiveCallbacksImpl();

  void rebuildActiveCallbacks();
  int sampleTries(double p) const;

  // std::mt19937 is quite large, so all scopes share the same generator.
  std::mt19937* generator_{nullptr};

  // Includes sampling callbacks which are waiting to run.
  c10::SmallVector<CallbackAndCounter, kSoftLimitCallbacks> callbacks_;
  RecordScope scope_{RecordScope::FUNCTION};

  StepCallbacks active_callbacks_;

  // For managing sampling callbacks
  int sampling_countdown_{0};
  int steps_for_this_update_{0};
};