Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/fixed_divisor.h

@@ -0,0 +1,137 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef CAFFE2_UTILS_FIXED_DIVISOR_H_
+#define CAFFE2_UTILS_FIXED_DIVISOR_H_
+
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+
+// See Note [hip-clang differences to hcc]
+
+#if defined(__CUDA_ARCH__) || defined(__HIP_ARCH__) || defined(__HIP__) || \
+    (defined(__clang__) && defined(__CUDA__))
+#define FIXED_DIVISOR_DECL inline __host__ __device__
+#else
+#define FIXED_DIVISOR_DECL inline
+#endif
+
+namespace caffe2 {
+
+// Utility class for quickly calculating quotients and remainders for
+// a known integer divisor
+template <typename T>
+class FixedDivisor {};
+
+// Works for any positive divisor, 1 to INT_MAX. One 64-bit
+// multiplication and one 64-bit shift is used to calculate the
+// result.
+template <>
+class FixedDivisor<std::int32_t> {
+ public:
+  FixedDivisor() = default;
+
+  explicit FixedDivisor(const std::int32_t d) : d_(d) {
+#if !defined(USE_ROCM)
+    CalcSignedMagic();
+#endif // USE_ROCM
+  }
+
+  FIXED_DIVISOR_DECL std::int32_t d() const {
+    return d_;
+  }
+
+#if !defined(USE_ROCM)
+  FIXED_DIVISOR_DECL std::uint64_t magic() const {
+    return magic_;
+  }
+
+  FIXED_DIVISOR_DECL int shift() const {
+    return shift_;
+  }
+#endif // USE_ROCM
+
+  /// Calculates `q = n / d`.
+  FIXED_DIVISOR_DECL std::int32_t Div(const std::int32_t n) const {
+#if defined(USE_ROCM)
+    return n / d_;
+#else // USE_ROCM
+    // In lieu of a mulhi instruction being available, perform the
+    // work in uint64
+    return (int32_t)((magic_ * (uint64_t)n) >> shift_);
+#endif // USE_ROCM
+  }
+
+  /// Calculates `r = n % d`.
+  FIXED_DIVISOR_DECL std::int32_t Mod(const std::int32_t n) const {
+    return n - d_ * Div(n);
+  }
+
+  /// Calculates `q = n / d` and `r = n % d` together.
+  FIXED_DIVISOR_DECL void
+  DivMod(const std::int32_t n, std::int32_t* q, int32_t* r) const {
+    *q = Div(n);
+    *r = n - d_ * *q;
+  }
+
+ private:
+#if !defined(USE_ROCM)
+  // Calculates magic multiplicative value and shift amount for calculating `q =
+  // n / d` for signed 32-bit integers.
+  // Implementation taken from Hacker's Delight section 10.
+  void CalcSignedMagic() {
+    if (d_ == 1) {
+      magic_ = UINT64_C(0x1) << 32;
+      shift_ = 32;
+      return;
+    }
+
+    const std::uint32_t two31 = UINT32_C(0x80000000);
+    const std::uint32_t ad = std::abs(d_);
+    const std::uint32_t t = two31 + ((uint32_t)d_ >> 31);
+    const std::uint32_t anc = t - 1 - t % ad; // Absolute value of nc.
+    std::uint32_t p = 31; // Init. p.
+    std::uint32_t q1 = two31 / anc; // Init. q1 = 2**p/|nc|.
+    std::uint32_t r1 = two31 - q1 * anc; // Init. r1 = rem(2**p, |nc|).
+    std::uint32_t q2 = two31 / ad; // Init. q2 = 2**p/|d|.
+    std::uint32_t r2 = two31 - q2 * ad; // Init. r2 = rem(2**p, |d|).
+    std::uint32_t delta = 0;
+    do {
+      ++p;
+      q1 <<= 1; // Update q1 = 2**p/|nc|.
+      r1 <<= 1; // Update r1 = rem(2**p, |nc|).
+      if (r1 >= anc) { // (Must be an unsigned
+        ++q1; // comparison here).
+        r1 -= anc;
+      }
+      q2 <<= 1; // Update q2 = 2**p/|d|.
+      r2 <<= 1; // Update r2 = rem(2**p, |d|).
+      if (r2 >= ad) { // (Must be an unsigned
+        ++q2; // comparison here).
+        r2 -= ad;
+      }
+      delta = ad - r2;
+    } while (q1 < delta || (q1 == delta && r1 == 0));
+    std::int32_t magic = q2 + 1;
+    if (d_ < 0) {
+      magic = -magic;
+    }
+    shift_ = p;
+    magic_ = (std::uint64_t)(std::uint32_t)magic;
+  }
+#endif // USE_ROCM
+
+  std::int32_t d_ = 1;
+
+#if !defined(USE_ROCM)
+  std::uint64_t magic_;
+  int shift_;
+#endif // USE_ROCM
+};
+
+} // namespace caffe2
+
+#endif // CAFFE2_UTILS_FIXED_DIVISOR_H_
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/proto_wrap.h

@@ -0,0 +1,42 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef CAFFE2_UTILS_PROTO_WRAP_H_
+#define CAFFE2_UTILS_PROTO_WRAP_H_
+
+#include <c10/util/Logging.h>
+
+namespace caffe2 {
+
+// A wrapper function to shut down protobuf library (this is needed in ASAN
+// testing and valgrind cases to avoid protobuf appearing to "leak" memory).
+TORCH_API void ShutdownProtobufLibrary();
+
+// Caffe2 wrapper functions for protobuf's GetEmptyStringAlreadyInited()
+// function used to avoid duplicated global variable in the case when protobuf
+// is built with hidden visibility.
+TORCH_API const ::std::string& GetEmptyStringAlreadyInited();
+} // namespace caffe2
+
+namespace ONNX_NAMESPACE {
+
+// ONNX wrapper functions for protobuf's GetEmptyStringAlreadyInited() function
+// used to avoid duplicated global variable in the case when protobuf
+// is built with hidden visibility.
+TORCH_API const ::std::string& GetEmptyStringAlreadyInited();
+
+} // namespace ONNX_NAMESPACE
+
+namespace torch {
+
+// Caffe2 wrapper functions for protobuf's GetEmptyStringAlreadyInited()
+// function used to avoid duplicated global variable in the case when protobuf
+// is built with hidden visibility.
+TORCH_API const ::std::string& GetEmptyStringAlreadyInited();
+
+void ShutdownProtobufLibrary();
+
+} // namespace torch
+#endif // CAFFE2_UTILS_PROTO_WRAP_H_
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/string_utils.h

@@ -0,0 +1,56 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <c10/macros/Export.h>
+
+namespace caffe2 {
+
+TORCH_API std::vector<std::string>
+split(char separator, const std::string& string, bool ignore_empty = false);
+
+TORCH_API std::string trim(const std::string& str);
+
+TORCH_API size_t editDistance(
+    const std::string& s1,
+    const std::string& s2,
+    size_t max_distance = 0);
+
+TORCH_API inline bool StartsWith(
+    const std::string& str,
+    const std::string& prefix) {
+  return str.length() >= prefix.length() &&
+      std::mismatch(prefix.begin(), prefix.end(), str.begin()).first ==
+      prefix.end();
+}
+
+TORCH_API inline bool EndsWith(
+    const std::string& full,
+    const std::string& ending) {
+  if (full.length() >= ending.length()) {
+    return (
+        0 ==
+        full.compare(full.length() - ending.length(), ending.length(), ending));
+  } else {
+    return false;
+  }
+}
+
+TORCH_API int32_t editDistanceHelper(
+    const char* s1,
+    size_t s1_len,
+    const char* s2,
+    size_t s2_len,
+    std::vector<size_t>& current,
+    std::vector<size_t>& previous,
+    std::vector<size_t>& previous1,
+    size_t max_distance);
+} // namespace caffe2
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/ThreadPool.h

@@ -0,0 +1,84 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef CAFFE2_UTILS_THREADPOOL_H_
+#define CAFFE2_UTILS_THREADPOOL_H_
+
+#include "ThreadPoolCommon.h"
+
+#include <atomic>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <vector>
+
+#include "c10/util/Flags.h"
+#include "caffe2/core/common.h"
+
+//
+// A work-stealing threadpool loosely based off of pthreadpool
+//
+
+namespace caffe2 {
+
+struct Task;
+class WorkersPool;
+
+constexpr size_t kCacheLineSize = 64;
+
+// A threadpool with the given number of threads.
+// NOTE: the kCacheLineSize alignment is present only for cache
+// performance, and is not strictly enforced (for example, when
+// the object is created on the heap). Thus, in order to avoid
+// misaligned intrinsics, no SSE instructions shall be involved in
+// the ThreadPool implementation.
+// Note: alignas is disabled because some compilers do not deal with
+// TORCH_API and alignas annotations at the same time.
+class TORCH_API /*alignas(kCacheLineSize)*/ ThreadPool {
+ public:
+  static ThreadPool* createThreadPool(int numThreads);
+  static std::unique_ptr<ThreadPool> defaultThreadPool();
+  virtual ~ThreadPool() = default;
+  // Returns the number of threads currently in use
+  virtual int getNumThreads() const = 0;
+  virtual void setNumThreads(size_t numThreads) = 0;
+
+  // Sets the minimum work size (range) for which to invoke the
+  // threadpool; work sizes smaller than this will just be run on the
+  // main (calling) thread
+  void setMinWorkSize(size_t size) {
+    std::lock_guard<std::mutex> guard(executionMutex_);
+    minWorkSize_ = size;
+  }
+
+  size_t getMinWorkSize() const {
+    return minWorkSize_;
+  }
+  virtual void run(const std::function<void(int, size_t)>& fn, size_t range) = 0;
+
+  // Run an arbitrary function in a thread-safe manner accessing the Workers
+  // Pool
+  virtual void withPool(const std::function<void(WorkersPool*)>& fn) = 0;
+
+ protected:
+  static size_t defaultNumThreads_;
+  mutable std::mutex executionMutex_;
+  size_t minWorkSize_;
+};
+
+size_t getDefaultNumThreads();
+} // namespace caffe2
+
+C10_DECLARE_bool(caffe2_threadpool_force_inline);
+
+// Whether or not threadpool caps apply to Android
+C10_DECLARE_int(caffe2_threadpool_android_cap);
+
+// Whether or not threadpool caps apply to iOS and MacOS
+C10_DECLARE_int(caffe2_threadpool_ios_cap);
+C10_DECLARE_int(caffe2_threadpool_macos_cap);
+
+C10_DECLARE_int(pthreadpool_size);
+#endif // CAFFE2_UTILS_THREADPOOL_H_
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/ThreadPoolCommon.h

@@ -0,0 +1,25 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#ifndef CAFFE2_UTILS_THREADPOOL_COMMON_H_
+#define CAFFE2_UTILS_THREADPOOL_COMMON_H_
+
+#ifdef __APPLE__
+#include <TargetConditionals.h>
+#endif
+
+// caffe2 depends upon NNPACK, which depends upon this threadpool, so
+// unfortunately we can't reference core/common.h here
+
+// This is copied from core/common.h's definition of C10_MOBILE
+// Define enabled when building for iOS or Android devices
+#if defined(__ANDROID__)
+#define C10_ANDROID 1
+#elif (defined(__APPLE__) &&                                            \
+       (TARGET_IPHONE_SIMULATOR || TARGET_OS_SIMULATOR || TARGET_OS_IPHONE))
+#define C10_IOS 1
+#endif // ANDROID / IOS
+
+#endif  // CAFFE2_UTILS_THREADPOOL_COMMON_H_
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/WorkersPool.h

@@ -0,0 +1,383 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <atomic>
+#include <condition_variable>
+#include <thread>
+#include "c10/util/thread_name.h"
+#include <c10/util/irange.h>
+#include <c10/util/Logging.h>
+
+#if defined(_MSC_VER)
+#include <intrin.h>
+#endif
+
+namespace caffe2 {
+
+// Uses code derived from gemmlowp,
+// https://github.com/google/gemmlowp/blob/6c91e1ed0c2eff1182d804310b92911fe9c18019/internal/multi_thread_gemm.h
+// Changes:
+// - allocation-free execute()
+// - Use RAII where possible.
+// - Run the first task on the main thread (since that is the largest task).
+// - removed custom allocator.
+// - Removed some ifdef's
+// - cache-line align Worker.
+// - use std::atomic instead of volatile and custom barriers.
+// - use std::mutex/std::condition_variable instead of raw pthreads.
+
+constexpr size_t kGEMMLOWPCacheLineSize = 64;
+
+template <typename T>
+struct AllocAligned {
+  // Allocate a T aligned at an `align` byte address
+  template <typename... Args>
+  static T* alloc(Args&&... args) {
+    void* p = nullptr;
+
+#if defined(__ANDROID__)
+    p = memalign(kGEMMLOWPCacheLineSize, sizeof(T));
+#elif defined(_MSC_VER)
+    p = _aligned_malloc(sizeof(T), kGEMMLOWPCacheLineSize);
+#else
+    auto res = posix_memalign(&p, kGEMMLOWPCacheLineSize, sizeof(T));
+    (void)res;
+#endif
+
+    if (p) {
+      return new (p) T(std::forward<Args>(args)...);
+    }
+
+    return nullptr;
+  }
+
+  // Free a T previously allocated via AllocAligned<T>::alloc()
+  static void release(T* p) {
+    if (p) {
+      p->~T();
+#if defined(_MSC_VER)
+      _aligned_free((void*)p);
+#else
+      free((void*)p);
+#endif
+    }
+  }
+};
+
+// Deleter object for unique_ptr for an aligned object
+template <typename T>
+struct AlignedDeleter {
+  void operator()(T* p) const { AllocAligned<T>::release(p); }
+};
+
+// make_unique that guarantees alignment
+template <typename T>
+struct MakeAligned {
+  template <typename... Args>
+  static std::unique_ptr<T, AlignedDeleter<T>> make(Args&&... args) {
+    return std::unique_ptr<T, AlignedDeleter<T>>(
+        AllocAligned<T>::alloc(std::forward<Args>(args)...));
+  }
+};
+
+const int kMaxBusyWaitNOPs = 32 * 1000 * 1000;
+
+#if defined(_MSC_VER)
+#define GEMMLOWP_NOP __nop();
+#else
+#define GEMMLOWP_NOP "nop\n"
+#endif
+
+#define GEMMLOWP_STRING_CONCAT_4(X) X X X X
+#define GEMMLOWP_NOP4 GEMMLOWP_STRING_CONCAT_4(GEMMLOWP_NOP)
+#define GEMMLOWP_NOP16 GEMMLOWP_STRING_CONCAT_4(GEMMLOWP_NOP4)
+#define GEMMLOWP_NOP64 GEMMLOWP_STRING_CONCAT_4(GEMMLOWP_NOP16)
+
+inline int Do256NOPs() {
+#if defined(_MSC_VER)
+  GEMMLOWP_NOP64;
+#else
+  asm volatile(GEMMLOWP_NOP64);
+#endif
+  return 64;
+}
+
+#undef GEMMLOWP_STRING_CONCAT_4
+#undef GEMMLOWP_NOP256
+#undef GEMMLOWP_NOP64
+#undef GEMMLOWP_NOP16
+#undef GEMMLOWP_NOP4
+#undef GEMMLOWP_NOP
+
+// Waits until *var != initial_value.
+//
+// Returns the new value of *var. The guarantee here is that
+// the return value is different from initial_value, and that that
+// new value has been taken by *var at some point during the
+// execution of this function. There is no guarantee that this is
+// still the value of *var when this function returns, since *var is
+// not assumed to be guarded by any lock.
+//
+// First does some busy-waiting for a fixed number of no-op cycles,
+// then falls back to passive waiting for the given condvar, guarded
+// by the given mutex.
+//
+// The idea of doing some initial busy-waiting is to help get
+// better and more consistent multithreading benefits for small GEMM sizes.
+// Busy-waiting help ensuring that if we need to wake up soon after having
+// started waiting, then we can wake up quickly (as opposed to, say,
+// having to wait to be scheduled again by the OS). On the other hand,
+// we must still eventually revert to passive waiting for longer waits
+// (e.g. worker threads having finished a GEMM and waiting until the next GEMM)
+// so as to avoid permanently spinning.
+//
+template <typename T>
+T WaitForVariableChange(std::atomic<T>* var,
+                        T initial_value,
+                        std::condition_variable* cond,
+                        std::mutex* mutex) {
+  // If we are on a platform that supports it, spin for some time.
+  {
+    int nops = 0;
+    // First, trivial case where the variable already changed value.
+    T new_value = var->load(std::memory_order_relaxed);
+    if (new_value != initial_value) {
+      std::atomic_thread_fence(std::memory_order_acquire);
+      return new_value;
+    }
+    // Then try busy-waiting.
+    while (nops < kMaxBusyWaitNOPs) {
+      nops += Do256NOPs();
+      new_value = var->load(std::memory_order_relaxed);
+      if (new_value != initial_value) {
+        std::atomic_thread_fence(std::memory_order_acquire);
+        return new_value;
+      }
+    }
+  }
+
+  // Finally, do real passive waiting.
+  {
+    std::unique_lock<std::mutex> g(*mutex);
+    T new_value = var->load(std::memory_order_relaxed);
+    // Handle spurious wakeups.
+    cond->wait(g, [&]() {
+      new_value = var->load(std::memory_order_relaxed);
+      return new_value != initial_value;
+    });
+    TORCH_DCHECK_NE(static_cast<size_t>(new_value), static_cast<size_t>(initial_value));
+    return new_value;
+  }
+}
+
+// A BlockingCounter lets one thread to wait for N events to occur.
+// This is how the master thread waits for all the worker threads
+// to have finished working.
+class BlockingCounter {
+ public:
+  // Sets/resets the counter; initial_count is the number of
+  // decrementing events that the Wait() call will be waiting for.
+  void Reset(std::size_t initial_count) {
+    std::lock_guard<std::mutex> g(mutex_);
+    TORCH_DCHECK_EQ(count_, 0);
+    count_ = initial_count;
+  }
+
+  // Decrements the counter; if the counter hits zero, signals
+  // the thread that was waiting for that, and returns true.
+  // Otherwise (if the decremented count is still nonzero),
+  // returns false.
+  bool DecrementCount() {
+    const auto count_value = count_.fetch_sub(1, std::memory_order_relaxed) - 1;
+    if (count_value == 0) {
+      std::lock_guard<std::mutex> g(mutex_);
+      cond_.notify_one();
+    }
+    bool retval = count_value == 0;
+    return retval;
+  }
+
+  // Waits for the N other threads (N having been set by Reset())
+  // to hit the BlockingCounter.
+  void Wait() {
+    while (size_t count_value = count_.load(std::memory_order_relaxed)) {
+      WaitForVariableChange(&count_, count_value, &cond_, &mutex_);
+    }
+  }
+
+ private:
+  std::condition_variable cond_;
+  std::mutex mutex_;
+  std::atomic<std::size_t> count_{0};
+};
+
+// A workload for a worker.
+struct Task {
+  Task() = default;
+  virtual ~Task() = default;
+  virtual void Run() = 0;
+};
+
+// A worker thread.
+class alignas(kGEMMLOWPCacheLineSize) Worker {
+ public:
+  enum class State : uint8_t {
+    ThreadStartup, // The initial state before the thread main loop runs.
+    Ready, // Is not working, has not yet received new work to do.
+    HasWork, // Has work to do.
+    ExitAsSoonAsPossible // Should exit at earliest convenience.
+  };
+
+  explicit Worker(BlockingCounter* counter_to_decrement_when_ready)
+      : task_(nullptr),
+        state_(State::ThreadStartup),
+        counter_to_decrement_when_ready_(counter_to_decrement_when_ready) {
+    thread_ = std::make_unique<std::thread>([this]() {
+      c10::setThreadName("pt_thread_pool");
+      this->ThreadFunc();
+    });
+  }
+
+  ~Worker() {
+    ChangeState(State::ExitAsSoonAsPossible);
+    thread_->join();
+  }
+
+  // Changes State; may be called from either the worker thread
+  // or the master thread; however, not all state transitions are legal,
+  // which is guarded by assertions.
+  void ChangeState(State new_state) {
+    std::lock_guard<std::mutex> g(state_mutex_);
+    DCHECK(new_state != state_.load(std::memory_order_relaxed));
+    switch (state_.load(std::memory_order_relaxed)) {
+    case State::ThreadStartup:
+      DCHECK(new_state == State::Ready);
+      break;
+    case State::Ready:
+      DCHECK(new_state == State::HasWork || new_state == State::ExitAsSoonAsPossible);
+      break;
+    case State::HasWork:
+      DCHECK(new_state == State::Ready || new_state == State::ExitAsSoonAsPossible);
+      break;
+    case State::ExitAsSoonAsPossible:
+    default:
+      abort();
+    }
+    state_.store(new_state, std::memory_order_relaxed);
+    state_cond_.notify_one();
+    if (new_state == State::Ready) {
+      counter_to_decrement_when_ready_->DecrementCount();
+    }
+  }
+
+  // Thread entry point.
+  void ThreadFunc() {
+    c10::setThreadName("CaffeWorkersPool");
+    ChangeState(State::Ready);
+
+    // Thread main loop
+    while (true) {
+      // Get a state to act on
+      // In the 'Ready' state, we have nothing to do but to wait until
+      // we switch to another state.
+      State state_to_act_upon =
+          WaitForVariableChange(&state_, State::Ready, &state_cond_, &state_mutex_);
+
+      // We now have a state to act on, so act.
+      switch (state_to_act_upon) {
+      case State::HasWork:
+        // Got work to do! So do it, and then revert to 'Ready' state.
+        DCHECK(task_.load());
+        (*task_).Run();
+        task_ = nullptr;
+        ChangeState(State::Ready);
+        break;
+      case State::ExitAsSoonAsPossible:
+        return;
+      case State::Ready:
+      case State::ThreadStartup:
+      default:
+        abort();
+      }
+    }
+  }
+
+  static void* ThreadFunc(void* arg) {
+    static_cast<Worker*>(arg)->ThreadFunc();
+    return nullptr;
+  }
+
+  // Called by the master thread to give this worker work to do.
+  // It is only legal to call this if the worker
+  void StartWork(Task* task) {
+    DCHECK(!task_.load());
+    task_ = task;
+    DCHECK(state_.load(std::memory_order_acquire) == State::Ready);
+    ChangeState(State::HasWork);
+  }
+
+ private:
+  // The underlying thread.
+  std::unique_ptr<std::thread> thread_;
+
+  // The task to be worked on.
+  std::atomic<Task*> task_;
+
+  // The condition variable and mutex guarding state changes.
+  std::condition_variable state_cond_;
+  std::mutex state_mutex_;
+
+  // The state enum tells if we're currently working, waiting for work, etc.
+  std::atomic<State> state_;
+
+  // pointer to the master's thread BlockingCounter object, to notify the
+  // master thread of when this worker switches to the 'Ready' state.
+  BlockingCounter* const counter_to_decrement_when_ready_;
+};
+
+class WorkersPool {
+ public:
+  WorkersPool() = default;
+
+  void Execute(const std::vector<std::shared_ptr<Task>>& tasks) {
+    CAFFE_ENFORCE_GE(tasks.size(), 1);
+    // One of the tasks will be run on the current thread.
+    int workers_count = tasks.size() - 1;
+    CreateWorkers(workers_count);
+    TORCH_DCHECK_LE(workers_count, (int)workers_.size());
+    counter_to_decrement_when_ready_.Reset(workers_count);
+    for (const auto task : c10::irange(1, tasks.size())) {
+      workers_[task - 1]->StartWork(tasks[task].get());
+    }
+    // Execute the remaining workload immediately on the current thread.
+    auto& task = tasks.front();
+    task->Run();
+    // Wait for the workers submitted above to finish.
+    counter_to_decrement_when_ready_.Wait();
+  }
+
+ private:
+  // Ensures that the pool has at least the given count of workers.
+  // If any new worker has to be created, this function waits for it to
+  // be ready.
+  void CreateWorkers(std::size_t workers_count) {
+    if (workers_.size() >= workers_count) {
+      return;
+    }
+    counter_to_decrement_when_ready_.Reset(workers_count - workers_.size());
+    while (workers_.size() < workers_count) {
+      workers_.push_back(MakeAligned<Worker>::make(&counter_to_decrement_when_ready_));
+    }
+    counter_to_decrement_when_ready_.Wait();
+  }
+
+  C10_DISABLE_COPY_AND_ASSIGN(WorkersPool);
+  std::vector<std::unique_ptr<Worker, AlignedDeleter<Worker>>> workers_;
+  // The BlockingCounter used to wait for the workers.
+  BlockingCounter counter_to_decrement_when_ready_;
+};
+} // namespace caffe2
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/pthreadpool-cpp.h

@@ -0,0 +1,60 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#ifdef USE_PTHREADPOOL
+
+#ifdef USE_INTERNAL_PTHREADPOOL_IMPL
+#include <caffe2/utils/threadpool/pthreadpool.h>
+#else
+#include <pthreadpool.h>
+#endif
+
+#include <functional>
+#include <memory>
+#include <mutex>
+
+namespace caffe2 {
+
+class PThreadPool final {
+ public:
+  explicit PThreadPool(size_t thread_count);
+  ~PThreadPool() = default;
+
+  PThreadPool(const PThreadPool&) = delete;
+  PThreadPool& operator=(const PThreadPool&) = delete;
+
+  PThreadPool(PThreadPool&&) = delete;
+  PThreadPool& operator=(PThreadPool&&) = delete;
+
+  size_t get_thread_count() const;
+  void set_thread_count(size_t thread_count);
+
+  // Run, in parallel, function fn(task_id) over task_id in range [0, range).
+  // This function is blocking.  All input is processed by the time it returns.
+  void run(const std::function<void(size_t)>& fn, size_t range);
+
+ private:
+  friend pthreadpool_t pthreadpool_();
+
+ private:
+  mutable std::mutex mutex_;
+  std::unique_ptr<pthreadpool, decltype(&pthreadpool_destroy)> threadpool_;
+};
+
+// Return a singleton instance of PThreadPool for ATen/TH multithreading.
+PThreadPool* pthreadpool();
+PThreadPool* pthreadpool(size_t thread_count);
+
+// Exposes the underlying implementation of PThreadPool.
+// Only for use in external libraries so as to unify threading across
+// internal (i.e. ATen, etc.) and external (e.g. NNPACK, QNNPACK, XNNPACK)
+// use cases.
+pthreadpool_t pthreadpool_();
+
+} // namespace caffe2
+
+#endif /* USE_PTHREADPOOL */
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/pthreadpool.h

@@ -0,0 +1,198 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+// pthreadpool header from https://github.com/Maratyszcza/pthreadpool
+// for NNPACK
+#ifndef CAFFE2_UTILS_PTHREADPOOL_H_
+#define CAFFE2_UTILS_PTHREADPOOL_H_
+
+#include "ThreadPoolCommon.h"
+
+#include <stddef.h> // for size_t
+#include <stdint.h> // for uint32_t
+
+#if defined(USE_PTHREADPOOL)
+// This is a hack.
+// Mainly introduced here because
+// 1. NNPACK can be compiled to use internal legacy threadpool implementation because much of C2 depends on that.
+// 2. Then if we want to use NNPACK in PyTorch, which uses new pthreadpool, then we will supply new pthreadpool pointer
+//    to NNPACK. This will not work if NNPACK is compiled with internal legacy threadpool. Thus this guard
+//    along with changes in pthreadpool_impl.cc allows us to override that behavior.
+//    It enables us to use NNPACK from pytorch using `caffe2::pthreadpool_()`
+namespace caffe2 {
+class WithCastToNewThreadPool {
+  public:
+    explicit WithCastToNewThreadPool(bool use_new_threadpool);
+    ~WithCastToNewThreadPool();
+  private:
+    bool use_new_threadpool_;
+};
+}
+#endif
+
+typedef struct pthreadpool* legacy_pthreadpool_t;
+
+typedef void (*legacy_pthreadpool_function_1d_t)(void*, size_t);
+typedef void (*legacy_pthreadpool_function_1d_tiled_t)(void*, size_t, size_t);
+typedef void (*legacy_pthreadpool_function_2d_t)(void*, size_t, size_t);
+typedef void (*legacy_pthreadpool_function_2d_tiled_t)(void*, size_t, size_t, size_t, size_t);
+typedef void (*legacy_pthreadpool_function_3d_tiled_t)(
+    void*,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t);
+typedef void (*legacy_pthreadpool_function_4d_tiled_t)(
+    void*,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t,
+    size_t);
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Creates a thread pool with the specified number of threads.
+ *
+ * @param[in]  threads_count  The number of threads in the thread pool.
+ *    A value of 0 has special interpretation: it creates a thread for each
+ *    processor core available in the system.
+ *
+ * @returns  A pointer to an opaque thread pool object.
+ *    On error the function returns NULL and sets errno accordingly.
+ */
+
+// Returns internal threadpool impl.
+legacy_pthreadpool_t legacy_pthreadpool_create(size_t threads_count);
+
+/**
+ * Queries the number of threads in a thread pool.
+ *
+ * @param[in]  threadpool  The thread pool to query.
+ *
+ * @returns  The number of threads in the thread pool.
+ */
+size_t legacy_pthreadpool_get_threads_count(legacy_pthreadpool_t threadpool);
+
+/**
+ * Processes items in parallel using threads from a thread pool.
+ *
+ * When the call returns, all items have been processed and the thread pool is
+ * ready for a new task.
+ *
+ * @note If multiple threads call this function with the same thread pool, the
+ *    calls are serialized.
+ *
+ * @param[in]  threadpool  The thread pool to use for parallelisation.
+ * @param[in]  function    The function to call for each item.
+ * @param[in]  argument    The first argument passed to the @a function.
+ * @param[in]  items       The number of items to process. The @a function
+ *    will be called once for each item.
+ */
+void legacy_pthreadpool_compute_1d(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_1d_t function,
+    void* argument,
+    size_t range);
+
+void legacy_pthreadpool_parallelize_1d(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_1d_t function,
+    void* argument,
+    size_t range,
+    uint32_t flags);
+
+void legacy_pthreadpool_compute_1d_tiled(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_1d_tiled_t function,
+    void* argument,
+    size_t range,
+    size_t tile);
+
+void legacy_pthreadpool_compute_2d(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_2d_t function,
+    void* argument,
+    size_t range_i,
+    size_t range_j);
+
+void legacy_pthreadpool_compute_2d_tiled(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_2d_tiled_t function,
+    void* argument,
+    size_t range_i,
+    size_t range_j,
+    size_t tile_i,
+    size_t tile_j);
+
+void legacy_pthreadpool_compute_3d_tiled(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_3d_tiled_t function,
+    void* argument,
+    size_t range_i,
+    size_t range_j,
+    size_t range_k,
+    size_t tile_i,
+    size_t tile_j,
+    size_t tile_k);
+
+void legacy_pthreadpool_compute_4d_tiled(
+    legacy_pthreadpool_t threadpool,
+    legacy_pthreadpool_function_4d_tiled_t function,
+    void* argument,
+    size_t range_i,
+    size_t range_j,
+    size_t range_k,
+    size_t range_l,
+    size_t tile_i,
+    size_t tile_j,
+    size_t tile_k,
+    size_t tile_l);
+
+/**
+ * Terminates threads in the thread pool and releases associated resources.
+ *
+ * @warning  Accessing the thread pool after a call to this function constitutes
+ *    undefined behaviour and may cause data corruption.
+ *
+ * @param[in,out]  threadpool  The thread pool to destroy.
+ */
+void legacy_pthreadpool_destroy(legacy_pthreadpool_t threadpool);
+
+#ifdef USE_INTERNAL_PTHREADPOOL_IMPL
+
+#define pthreadpool_t legacy_pthreadpool_t
+#define pthreadpool_function_1d_t legacy_pthreadpool_function_1d_t
+#define pthreadpool_function_1d_tiled_t legacy_pthreadpool_function_1d_tiled_t
+#define pthreadpool_function_2d_t legacy_pthreadpool_function_2d_t
+#define pthreadpool_function_2d_tiled_t legacy_pthreadpool_function_2d_tiled_t
+#define pthreadpool_function_3d_tiled_t legacy_pthreadpool_function_3d_tiled_t
+#define pthreadpool_function_4d_tiled_t legacy_pthreadpool_function_4d_tiled_t
+#define pthreadpool_create legacy_pthreadpool_create
+#define pthreadpool_destroy legacy_pthreadpool_destroy
+#define pthreadpool_get_threads_count legacy_pthreadpool_get_threads_count
+#define pthreadpool_compute_1d legacy_pthreadpool_compute_1d
+#define pthreadpool_parallelize_1d legacy_pthreadpool_parallelize_1d
+#define pthreadpool_compute_1d_tiled legacy_pthreadpool_compute_1d_tiled
+#define pthreadpool_compute_2d legacy_pthreadpool_compute_2d
+#define pthreadpool_compute_2d_tiled legacy_pthreadpool_compute_2d_tiled
+#define pthreadpool_compute_3d_tiled legacy_pthreadpool_compute_3d_tiled
+#define pthreadpool_compute_4d_tiled legacy_pthreadpool_compute_4d_tiled
+
+#endif /* USE_INTERNAL_PTHREADPOOL_IMPL */
+
+#ifdef __cplusplus
+} /* extern "C" */
+#endif
+
+#endif // CAFFE2_UTILS_PTHREADPOOL_H_
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)

kronos-venv/lib/python3.12/site-packages/torch/include/caffe2/utils/threadpool/thread_pool_guard.h

@@ -0,0 +1,28 @@
+#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
+#pragma once
+
+#include <c10/macros/Macros.h>
+
+namespace caffe2 {
+
+// A RAII, thread local (!) guard that enables or disables grad mode upon
+// construction, and sets it back to the original value upon destruction.
+struct TORCH_API _NoPThreadPoolGuard {
+  static bool is_enabled();
+  static void set_enabled(bool enabled);
+
+  _NoPThreadPoolGuard(): prev_mode_(_NoPThreadPoolGuard::is_enabled()) {
+      _NoPThreadPoolGuard::set_enabled(true);
+  }
+  ~_NoPThreadPoolGuard() {
+      _NoPThreadPoolGuard::set_enabled(prev_mode_);
+  }
+  private:
+    bool prev_mode_;
+};
+
+}
+
+#else
+#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
+#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)