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

This commit is contained in:
Kolp
2026-09-24 13:22:23 +07:00
commit 642cc11a9f
18968 changed files with 5683248 additions and 0 deletions
@@ -0,0 +1,66 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#ifndef CAFFE2_CORE_COMMON_H_
#define CAFFE2_CORE_COMMON_H_
#include <algorithm>
#include <cmath>
#include <map>
#include <memory>
#include <numeric>
#include <set>
#include <sstream>
#include <string>
#include <type_traits>
#include <vector>
#ifdef __APPLE__
#include <TargetConditionals.h>
#endif
#if defined(_MSC_VER)
#include <io.h>
#else
#include <unistd.h>
#endif
// Macros used during the build of this caffe2 instance. This header file
// is automatically generated by the cmake script during build.
#include "caffe2/core/macros.h"
#include <c10/macros/Macros.h>
namespace caffe2 {
// Using statements for common classes that we refer to in caffe2 very often.
// Note that we only place it inside caffe2 so the global namespace is not
// polluted.
/* using override */
using std::set;
using std::string;
using std::unique_ptr;
using std::vector;
// Define alignment macro that is cross platform
#if (defined _MSC_VER && !defined NOMINMAX)
#define NOMINMAX
#endif
using std::make_unique;
#if defined(__ANDROID__) && !defined(__NDK_MAJOR__)
using ::round;
#else
using std::round;
#endif // defined(__ANDROID__) && !defined(__NDK_MAJOR__)
// Returns which setting Caffe2 was configured and built with (exported from
// CMake)
TORCH_API const std::map<string, string>& GetBuildOptions();
} // namespace caffe2
#endif // CAFFE2_CORE_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)
@@ -0,0 +1,77 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// Automatically generated header file for caffe2 macros. These
// macros are used to build the Caffe2 binary, and if you are
// building a dependent library, they will need to be set as well
// for your program to link correctly.
#pragma once
#define CAFFE2_BUILD_SHARED_LIBS
/* #undef CAFFE2_FORCE_FALLBACK_CUDA_MPI */
/* #undef CAFFE2_HAS_MKL_DNN */
/* #undef CAFFE2_HAS_MKL_SGEMM_PACK */
#define CAFFE2_PERF_WITH_AVX
#define CAFFE2_PERF_WITH_AVX2
/* #undef CAFFE2_THREADPOOL_MAIN_IMBALANCE */
/* #undef CAFFE2_THREADPOOL_STATS */
/* #undef CAFFE2_USE_ACCELERATE */
/* #undef CAFFE2_USE_CUDNN */
/* #undef CAFFE2_USE_EIGEN_FOR_BLAS */
/* #undef CAFFE2_USE_FBCODE */
/* #undef CAFFE2_USE_GOOGLE_GLOG */
/* #undef CAFFE2_USE_LITE_PROTO */
#define CAFFE2_USE_MKL
#define USE_MKLDNN
/* #undef CAFFE2_USE_NVTX */
/* #undef CAFFE2_USE_ITT */
#ifndef EIGEN_MPL2_ONLY
#define EIGEN_MPL2_ONLY
#endif
// Useful build settings that are recorded in the compiled binary
// torch.__config__.show()
#define CAFFE2_BUILD_STRINGS { \
{"TORCH_VERSION", "2.12.0"}, \
{"CXX_COMPILER", "/opt/rh/gcc-toolset-13/root/usr/bin/c++"}, \
{"CXX_FLAGS", " -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DLIBKINETO_NOROCTRACER -DLIBKINETO_NOXPUPTI=ON -DUSE_FBGEMM -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -DC10_NODEPRECATED -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=range-loop-construct -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-unknown-pragmas -Wno-unused-parameter -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wsuggest-override -Wno-psabi -Wno-error=old-style-cast -faligned-new -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-dangling-reference -Wno-error=dangling-reference -Wno-stringop-overflow"}, \
{"CUDA_FLAGS", " -DLIBCUDACXX_ENABLE_SIMPLIFIED_COMPLEX_OPERATIONS -Xfatbin -compress-all -Wno-deprecated-gpu-targets --expt-extended-lambda -DCUB_WRAPPED_NAMESPACE=at_cuda_detail -DDISABLE_CUSPARSE_DEPRECATED -DCUDA_HAS_FP16=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -DC10_NODEPRECATED"}, \
{"BUILD_TYPE", "Release"}, \
{"BLAS_INFO", "mkl"}, \
{"LAPACK_INFO", "mkl"}, \
{"USE_CUDA", "0"}, \
{"USE_ROCM", "OFF"}, \
{"CUDA_VERSION", ""}, \
{"ROCM_VERSION", ""}, \
{"USE_CUDNN", "OFF"}, \
{"COMMIT_SHA", "7661cd9c6b841b62b7f411aa52ec51f05457263b"}, \
{"CUDNN_VERSION", ""}, \
{"USE_NCCL", "OFF"}, \
{"USE_MPI", "OFF"}, \
{"USE_GFLAGS", "OFF"}, \
{"USE_GLOG", "OFF"}, \
{"USE_GLOO", "ON"}, \
{"USE_NNPACK", "ON"}, \
{"USE_OPENMP", "ON"}, \
{"FORCE_FALLBACK_CUDA_MPI", ""}, \
{"HAS_MKL_DNN", ""}, \
{"HAS_MKL_SGEMM_PACK", ""}, \
{"PERF_WITH_AVX", "1"}, \
{"PERF_WITH_AVX2", "1"}, \
{"USE_ACCELERATE", ""}, \
{"USE_EIGEN_FOR_BLAS", ""}, \
{"USE_LITE_PROTO", ""}, \
{"USE_MKL", "ON"}, \
{"USE_MKLDNN", "ON"}, \
{"USE_NVTX", ""}, \
{"USE_ITT", ""}, \
{"USE_ROCM_KERNEL_ASSERT", "OFF"}, \
{"USE_CUSPARSELT", "OFF"}, \
{"USE_XPU", "OFF"}, \
{"USE_XCCL", "OFF"}, \
{"SYCL_COMPILER_VERSION", ""}, \
}
#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)
@@ -0,0 +1,53 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#ifndef CAFFE2_CORE_TIMER_H_
#define CAFFE2_CORE_TIMER_H_
#include <chrono>
#include "caffe2/core/common.h"
namespace caffe2 {
/**
* @brief A simple timer object for measuring time.
*
* This is a minimal class around a std::chrono::high_resolution_clock that
* serves as a utility class for testing code.
*/
class Timer {
public:
typedef std::chrono::high_resolution_clock clock;
typedef std::chrono::nanoseconds ns;
Timer() { Start(); }
/**
* @brief Starts a timer.
*/
inline void Start() { start_time_ = clock::now(); }
inline float NanoSeconds() {
return static_cast<float>(
std::chrono::duration_cast<ns>(clock::now() - start_time_).count());
}
/**
* @brief Returns the elapsed time in milliseconds.
*/
inline float MilliSeconds() { return NanoSeconds() / 1000000.f; }
/**
* @brief Returns the elapsed time in microseconds.
*/
inline float MicroSeconds() { return NanoSeconds() / 1000.f; }
/**
* @brief Returns the elapsed time in seconds.
*/
inline float Seconds() { return NanoSeconds() / 1000000000.f; }
protected:
std::chrono::time_point<clock> start_time_;
C10_DISABLE_COPY_AND_ASSIGN(Timer);
};
}
#endif // CAFFE2_CORE_TIMER_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)
@@ -0,0 +1,145 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// !!!! PLEASE READ !!!!
// Minimize (transitively) included headers from _avx*.cc because some of the
// functions defined in the headers compiled with platform dependent compiler
// options can be reused by other translation units generating illegal
// instruction run-time error.
// Common utilities for writing performance kernels and easy dispatching of
// different backends.
/*
The general workflow shall be as follows, say we want to
implement a functionality called void foo(int a, float b).
In foo.h, do:
void foo(int a, float b);
In foo_avx512.cc, do:
void foo__avx512(int a, float b) {
[actual avx512 implementation]
}
In foo_avx2.cc, do:
void foo__avx2(int a, float b) {
[actual avx2 implementation]
}
In foo_avx.cc, do:
void foo__avx(int a, float b) {
[actual avx implementation]
}
In foo.cc, do:
// The base implementation should *always* be provided.
void foo__base(int a, float b) {
[base, possibly slow implementation]
}
decltype(foo__base) foo__avx512;
decltype(foo__base) foo__avx2;
decltype(foo__base) foo__avx;
void foo(int a, float b) {
// You should always order things by their preference, faster
// implementations earlier in the function.
AVX512_DO(foo, a, b);
AVX2_DO(foo, a, b);
AVX_DO(foo, a, b);
BASE_DO(foo, a, b);
}
*/
// Details: this functionality basically covers the cases for both build time
// and run time architecture support.
//
// During build time:
// The build system should provide flags CAFFE2_PERF_WITH_AVX512,
// CAFFE2_PERF_WITH_AVX2, and CAFFE2_PERF_WITH_AVX that corresponds to the
// __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX2__, and __AVX__ flags the
// compiler provides. Note that we do not use the compiler flags but rely on
// the build system flags, because the common files (like foo.cc above) will
// always be built without __AVX512F__, __AVX512DQ__, __AVX512VL__, __AVX2__
// and __AVX__.
// During run time:
// we use cpuinfo to identify cpu support and run the proper functions.
#pragma once
#if defined(CAFFE2_PERF_WITH_SVE) || defined(CAFFE2_PERF_WITH_AVX512) || \
defined(CAFFE2_PERF_WITH_AVX2) || defined(CAFFE2_PERF_WITH_AVX)
#include <cpuinfo.h>
#endif
// DO macros: these should be used in your entry function, similar to foo()
// above, that routes implementations based on CPU capability.
#define BASE_DO(funcname, ...) return funcname##__base(__VA_ARGS__);
#ifdef CAFFE2_PERF_WITH_SVE
#define SVE_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && cpuinfo_has_arm_sve(); \
if (isDo) { \
return funcname##__sve(__VA_ARGS__); \
} \
}
#else // CAFFE2_PERF_WITH_SVE
#define SVE_DO(funcname, ...)
#endif // CAFFE2_PERF_WITH_SVE
#ifdef CAFFE2_PERF_WITH_AVX512
#define AVX512_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && \
cpuinfo_has_x86_avx512f() && cpuinfo_has_x86_avx512dq() && \
cpuinfo_has_x86_avx512vl(); \
if (isDo) { \
return funcname##__avx512(__VA_ARGS__); \
} \
}
#else // CAFFE2_PERF_WITH_AVX512
#define AVX512_DO(funcname, ...)
#endif // CAFFE2_PERF_WITH_AVX512
#ifdef CAFFE2_PERF_WITH_AVX2
#define AVX2_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx2(); \
if (isDo) { \
return funcname##__avx2(__VA_ARGS__); \
} \
}
#define AVX2_FMA_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx2() && \
cpuinfo_has_x86_fma3(); \
if (isDo) { \
return funcname##__avx2_fma(__VA_ARGS__); \
} \
}
#else // CAFFE2_PERF_WITH_AVX2
#define AVX2_DO(funcname, ...)
#define AVX2_FMA_DO(funcname, ...)
#endif // CAFFE2_PERF_WITH_AVX2
#ifdef CAFFE2_PERF_WITH_AVX
#define AVX_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx(); \
if (isDo) { \
return funcname##__avx(__VA_ARGS__); \
} \
}
#define AVX_F16C_DO(funcname, ...) \
{ \
static const bool isDo = cpuinfo_initialize() && cpuinfo_has_x86_avx() && \
cpuinfo_has_x86_f16c(); \
if (isDo) { \
return funcname##__avx_f16c(__VA_ARGS__); \
} \
}
#else // CAFFE2_PERF_WITH_AVX
#define AVX_DO(funcname, ...)
#define AVX_F16C_DO(funcname, ...)
#endif // CAFFE2_PERF_WITH_AVX
#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)
@@ -0,0 +1,62 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <cstdint>
namespace caffe2 {
// clang-format off
/**
* Embedding lookup with reduction.
*
* `input` of size data_size * block_size
* `indices` of size index_size
* `offsets` of size output_size
* `weights` nullptr or array of size index_size
* `out` of size output_size * block_size
*
* Behavior is roughly equivalent to pseudocode:
*
* pos = 0
* for (i = 0..output_size-1)
* for (k = 0..block_size-1)
* out[i*block_size + k] = 0
* start_offset = offsets[i]
* end_offset = offsets[i+1]
* length = end_offset - start_offset
* for (j = start_offset..end_offset-1)
* for (k = 0..block_size-1)
* out[i*block_size + k] += input[indices[pos]*block_size + k] *
* (weights ? weights[IS_WEIGHT_POSITIONAL ? j - start_offset : pos] : 1.0)
* pos += 1
* if (normalize_weights && length > 0)
* for (k = 0..block_size-1)
* out[i*block_size + k] /= length
*
* TODO: make this API also take "offsets" rather than "lengths" to match the
* API for PyTorch's EmbeddingBag
*/
// clang-format on
template <
typename IndexType,
typename InType,
typename OutType,
bool IS_WEIGHT_POSITIONAL = false>
void EmbeddingLookupIdx(
const std::int64_t block_size,
const std::int64_t output_size,
const std::int64_t index_size,
const std::int64_t data_size,
const InType* input,
const IndexType* indices,
const IndexType* offsets,
const float* weights, // optional, can be null for non-weighted sum
const float* scale_bias, // optional scale & bias params for uint8 input
bool normalize_by_lengths,
OutType* out);
} // 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)
@@ -0,0 +1,40 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <c10/macros/Macros.h>
#include <fstream>
#include <memory>
#include "caffe2/serialize/istream_adapter.h"
#include "caffe2/serialize/read_adapter_interface.h"
namespace caffe2::serialize {
class TORCH_API FileAdapter final : public ReadAdapterInterface {
public:
C10_DISABLE_COPY_AND_ASSIGN(FileAdapter);
explicit FileAdapter(const std::string& file_name);
size_t size() const override;
size_t read(uint64_t pos, void* buf, size_t n, const char* what = "")
const override;
~FileAdapter() override;
private:
// An RAII Wrapper for a FILE pointer. Closes on destruction.
struct RAIIFile {
FILE* fp_;
explicit RAIIFile(const std::string& file_name);
~RAIIFile();
};
RAIIFile file_;
// The size of the opened file in bytes
uint64_t size_;
};
} // namespace caffe2::serialize
#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)
@@ -0,0 +1,35 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <caffe2/serialize/read_adapter_interface.h>
#include <cstring>
namespace caffe2 {
namespace serialize {
class MemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface {
public:
explicit MemoryReadAdapter(const void* data, off_t size)
: data_(data), size_(size) {}
size_t size() const override {
return size_;
}
size_t read(uint64_t pos, void* buf, size_t n, const char* what = "")
const override {
(void)what;
memcpy(buf, (int8_t*)(data_) + pos, n);
return n;
}
private:
const void* data_;
off_t size_;
};
} // namespace serialize
} // 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)
@@ -0,0 +1,314 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <cerrno>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <istream>
#include <mutex>
#include <ostream>
#include <unordered_set>
#include <c10/core/Allocator.h>
#include <c10/core/Backend.h>
#include "caffe2/serialize/istream_adapter.h"
#include "caffe2/serialize/read_adapter_interface.h"
#include "caffe2/serialize/versions.h"
extern "C" {
typedef struct mz_zip_archive mz_zip_archive;
}
// PyTorch containers are a special zip archive with the following layout
// archive_name.zip contains:
// archive_name/
// version # a file with a single decimal number written in ascii,
// # used to establish the version of the archive format
// model.json # overall model description, this is a json output of
// # ModelDef from torch.proto
// # the following names are by convention only, model.json will
// # refer to these files by full names
// tensors/
// 0 # flat storage for tensor data, meta-data about shapes, etc. is
// # in model.json
// 1
// ...
// # code entries will only exist for modules that have methods attached
// code/
// archive_name.py # serialized torch script code (python syntax, using
// PythonPrint) archive_name_my_submodule.py # submodules have separate
// files
//
// The PyTorchStreamWriter also ensures additional useful properties for these
// files
// 1. All files are stored uncompressed.
// 2. All files in the archive are aligned to 64 byte boundaries such that
// it is possible to mmap the entire file and get an aligned pointer to
// tensor data.
// 3. We universally write in ZIP64 format for consistency.
// The PyTorchStreamReader also provides additional properties:
// 1. It can read zip files that are created with common
// zip tools. This means that even though our writer doesn't compress files,
// the reader can still read files that were compressed.
// 2. It provides a getRecordOffset function which returns the offset into the
// raw file where file data lives. If the file was written with
// PyTorchStreamWriter it is guaranteed to be 64 byte aligned.
// PyTorchReader/Writer handle checking the version number on the archive format
// and ensure that all files are written to a archive_name directory so they
// unzip cleanly.
// When developing this format we want to pay particular attention to the
// following use cases:
//
// -- Reading --
// 1) Reading with full random access
// a) Reading with file api's such as fread()
// b) mmaping the file and jumping around the mapped region
// 2) Reading with 1-pass sequential access
// -> A reader will need to build up a data structure of parsed structures
// as it reads
//
// -- Writing --
// 1) Writing with full random access
// 2) Writing with 1-pass sequential access
// -> We must take care not to require updating values that have already
// been written. We place the variable-length index at the end and do
// not put any index into the header to fulfill this constraint.
// The model.json, which contains all the metadata information,
// should be written as the last file. One reason is that the size of tensor
// data is usually stable. As long as the shape and type of the tensor do not
// change, the size of the data won't change. On the other sied, the size of the
// serialized model is likely to change, so we store it as the last record, and
// we don't need to move previous records when updating the model data.
// The zip format is sufficiently flexible to handle the above use-case.
// it puts its central directory at the end of the archive and we write
// model.json as the last file when writing after we have accumulated all
// other information.
namespace caffe2::serialize {
static constexpr const char* kSerializationIdRecordName =
".data/serialization_id";
struct MzZipReaderIterWrapper;
class TORCH_API ChunkRecordIterator {
public:
~ChunkRecordIterator();
// Read at most `chunkSize` into `buf`. Return the number of actual bytes
// read.
size_t next(void* buf);
size_t recordSize() const {
return recordSize_;
}
private:
ChunkRecordIterator(
size_t recordSize,
size_t chunkSize,
std::unique_ptr<MzZipReaderIterWrapper> iter);
const size_t recordSize_;
const size_t chunkSize_;
size_t offset_;
std::unique_ptr<MzZipReaderIterWrapper> iter_;
friend class PyTorchStreamReader;
};
class TORCH_API PyTorchStreamReader final {
public:
explicit PyTorchStreamReader(const std::string& file_name);
explicit PyTorchStreamReader(std::istream* in);
explicit PyTorchStreamReader(std::shared_ptr<ReadAdapterInterface> in);
// return dataptr, size
// set allocator to override default cpu allocator
std::tuple<at::DataPtr, size_t> getRecord(
const std::string& name,
std::optional<at::Allocator*> allocator = std::nullopt);
// multi-thread getRecord
std::tuple<at::DataPtr, size_t> getRecord(
const std::string& name,
std::vector<std::shared_ptr<ReadAdapterInterface>>& additionalReaders,
std::optional<at::Allocator*> allocator = std::nullopt);
// inplace memory writing
size_t getRecord(const std::string& name, void* dst, size_t n);
// inplace memory writing, multi-threads.
// When additionalReaders is empty, the default behavior is call
// getRecord(name, dst, n) with default reader This approach can be used for
// reading large tensors.
size_t getRecord(
const std::string& name,
void* dst,
size_t n,
std::vector<std::shared_ptr<ReadAdapterInterface>>& additionalReaders);
size_t getRecord(
const std::string& name,
void* dst,
size_t n,
size_t chunk_size,
void* buf,
const std::function<void(void*, const void*, size_t)>& memcpy_func =
nullptr);
// Concurrent reading records with multiple readers.
// additionalReaders are additional clients to access the underlying record at
// different offsets and write to different trunks of buffers. If the overall
// size of the tensor is 10, and size of additionalReader is 2. The default
// thread will read [0,4), the additional reader will read [4,8). The default
// reader will read [8,10). The default reader will write to buffer[0,4), the
// additional reader will write to buffer[4,8), the additional reader will
// write to buffer[8,10). When additionalReaders is empty, the default
// behavior is call getRecord(name) with default reader This approach can be
// used for reading large tensors.
size_t getRecordMultiReaders(
const std::string& name,
std::vector<std::shared_ptr<ReadAdapterInterface>>& additionalReaders,
void* dst,
size_t n);
size_t getRecordSize(const std::string& name);
size_t getRecordHeaderOffset(const std::string& name);
size_t getRecordOffset(const std::string& name);
size_t getRecordOffsetNoRead(
size_t cursor,
std::string filename,
size_t size,
uint64_t alignment);
bool hasRecord(const std::string& name);
std::vector<std::string> getAllRecords();
ChunkRecordIterator createChunkReaderIter(
const std::string& name,
const size_t recordSize,
const size_t chunkSize);
~PyTorchStreamReader();
uint64_t version() const {
return version_;
}
const std::string& serializationId() {
return serialization_id_;
}
void setShouldLoadDebugSymbol(bool should_load_debug_symbol) {
load_debug_symbol_ = should_load_debug_symbol;
}
void setAdditionalReaderSizeThreshold(const size_t& size) {
additional_reader_size_threshold_ = size;
}
private:
void init();
size_t read(uint64_t pos, char* buf, size_t n);
void valid(const char* what, const char* info = "");
size_t getRecordID(const std::string& name);
friend size_t
istream_read_func(void* pOpaque, uint64_t file_ofs, void* pBuf, size_t n);
std::unique_ptr<mz_zip_archive> ar_;
std::string archive_name_;
std::string archive_name_plus_slash_;
std::shared_ptr<ReadAdapterInterface> in_;
int64_t version_;
std::mutex reader_lock_;
bool load_debug_symbol_ = true;
std::string serialization_id_;
size_t additional_reader_size_threshold_;
};
class TORCH_API PyTorchStreamWriter final {
public:
explicit PyTorchStreamWriter(
const std::string& archive_name,
bool compute_crc32 = true,
uint64_t alignment = 64);
explicit PyTorchStreamWriter(
const std::function<size_t(const void*, size_t)> writer_func,
bool compute_crc32 = true,
uint64_t alignment = 64);
void setMinVersion(const uint64_t version);
void writeRecord(
const std::string& name,
const void* data,
size_t size,
bool compress = false);
void writeEndOfFile();
const std::unordered_set<std::string>& getAllWrittenRecords();
bool finalized() const {
return finalized_;
}
const std::string& archiveName() {
return archive_name_;
}
const std::string& serializationId() {
return serialization_id_;
}
~PyTorchStreamWriter();
private:
void setup(const std::string& file_name);
void valid(const char* what, const char* info = "");
void writeSerializationId();
size_t current_pos_ = 0;
std::unordered_set<std::string> files_written_;
std::unique_ptr<mz_zip_archive> ar_;
std::string archive_name_;
std::string archive_name_plus_slash_;
std::string padding_;
std::ofstream file_stream_;
std::function<size_t(const void*, size_t)> writer_func_;
uint64_t combined_uncomp_crc32_ = 0;
std::string serialization_id_;
bool compute_crc32_;
uint64_t alignment_;
// This number will be updated when the model has operators
// that have valid upgraders.
uint64_t version_ = kMinProducedFileFormatVersion;
bool finalized_ = false;
bool err_seen_ = false;
friend size_t ostream_write_func(
void* pOpaque,
uint64_t file_ofs,
const void* pBuf,
size_t n);
};
namespace detail {
// Returns a record to be appended to the local user extra data entry in order
// to make data beginning aligned at kFieldAlignment bytes boundary.
size_t getPadding(
size_t cursor,
size_t filename_size,
size_t size,
std::string& padding_buf,
uint64_t alignment);
std::tuple<size_t, size_t>
getOffset(size_t cursor, size_t filename_size, size_t size, uint64_t alignment);
} // namespace detail
} // namespace caffe2::serialize
#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)
@@ -0,0 +1,31 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <istream>
#include "c10/macros/Macros.h"
#include "caffe2/serialize/read_adapter_interface.h"
namespace caffe2::serialize {
// this is a reader implemented by std::istream
class TORCH_API IStreamAdapter final : public ReadAdapterInterface {
public:
C10_DISABLE_COPY_AND_ASSIGN(IStreamAdapter);
explicit IStreamAdapter(std::istream* istream);
size_t size() const override;
size_t read(uint64_t pos, void* buf, size_t n, const char* what = "")
const override;
~IStreamAdapter() override;
private:
std::istream* istream_;
void validate(const char* what) const;
};
} // namespace caffe2::serialize
#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)
@@ -0,0 +1,27 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <cstddef>
#include <cstdint>
#include "c10/macros/Macros.h"
namespace caffe2::serialize {
// this is the interface for the (file/stream/memory) reader in
// PyTorchStreamReader. with this interface, we can extend the support
// besides standard istream
class TORCH_API ReadAdapterInterface {
public:
virtual size_t size() const = 0;
virtual size_t read(uint64_t pos, void* buf, size_t n, const char* what = "")
const = 0;
virtual ~ReadAdapterInterface();
};
} // namespace caffe2::serialize
#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)
@@ -0,0 +1,137 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <cstdint>
namespace caffe2::serialize {
constexpr uint64_t kMinSupportedFileFormatVersion = 0x1L;
constexpr uint64_t kMaxSupportedFileFormatVersion = 0xAL;
// Versions (i.e. why was the version number bumped?)
// Note [Dynamic Versions and torch.jit.save vs. torch.save]
//
// Our versioning scheme has a "produced file format version" which
// describes how an archive is to be read. The version written in an archive
// is at least this current produced file format version, but may be greater
// if it includes certain symbols. We refer to these conditional versions
// as "dynamic," since they are identified at runtime.
//
// Dynamic versioning is useful when an operator's semantics are updated.
// When using torch.jit.save we want those semantics to be preserved. If
// we bumped the produced file format version on every change, however,
// then older versions of PyTorch couldn't read even simple archives, like
// a single tensor, from newer versions of PyTorch. Instead, we
// assign dynamic versions to these changes that override the
// produced file format version as needed. That is, when the semantics
// of torch.div changed it was assigned dynamic version 4, and when
// torch.jit.saving modules that use torch.div those archives also have
// (at least) version 4. This prevents earlier versions of PyTorch
// from accidentally performing the wrong kind of division. Modules
// that don't use torch.div or other operators with dynamic versions
// can write the produced file format version, and these programs will
// run as expected on earlier versions of PyTorch.
//
// While torch.jit.save attempts to preserve operator semantics,
// torch.save does not. torch.save is analogous to pickling Python, so
// a function that uses torch.div will have different behavior if torch.saved
// and torch.loaded across PyTorch versions. From a technical perspective,
// torch.save ignores dynamic versioning.
// 1. Initial version
// 2. Removed op_version_set version numbers
// 3. Added type tags to pickle serialization of container types
// 4. (Dynamic) Stopped integer division using torch.div
// (a versioned symbol preserves the historic behavior of versions 1--3)
// 5. (Dynamic) Stops torch.full inferring a floating point dtype
// when given bool or integer fill values.
// 6. Write version string to `./data/version` instead of `version`.
// [12/15/2021]
// kProducedFileFormatVersion is set to 7 from 3 due to a different
// interpretation of what file format version is.
// Whenever there is new upgrader introduced,
// this number should be bumped.
// The reasons that version is bumped in the past:
// 1. aten::div is changed at version 4
// 2. aten::full is changed at version 5
// 3. torch.package uses version 6
// 4. Introduce new upgrader design and set the version number to 7
// mark this change
// --------------------------------------------------
// We describe new operator version bump reasons here:
// 1) [01/24/2022]
// We bump the version number to 8 to update aten::linspace
// and aten::linspace.out to error out when steps is not
// provided. (see: https://github.com/pytorch/pytorch/issues/55951)
// 2) [01/30/2022]
// Bump the version number to 9 to update aten::logspace and
// and aten::logspace.out to error out when steps is not
// provided. (see: https://github.com/pytorch/pytorch/issues/55951)
// 3) [02/11/2022]
// Bump the version number to 10 to update aten::gelu and
// and aten::gelu.out to support the new approximate kwarg.
// (see: https://github.com/pytorch/pytorch/pull/61439)
constexpr uint64_t kProducedFileFormatVersion = 0xAL;
// Absolute minimum version we will write packages. This
// means that every package from now on will always be
// greater than this number.
constexpr uint64_t kMinProducedFileFormatVersion = 0x3L;
// The version we write when the archive contains bytecode.
// It must be higher or eq to kProducedFileFormatVersion.
// Because torchscript changes is likely introduce bytecode change.
// If kProducedFileFormatVersion is increased, kProducedBytecodeVersion
// should be increased too. The relationship is:
// kMaxSupportedFileFormatVersion >= (most likely ==) kProducedBytecodeVersion
// >= kProducedFileFormatVersion
// If a format change is forward compatible (still readable by older
// executables), we will not increment the version number, to minimize the
// risk of breaking existing clients. TODO: A better way would be to allow
// the caller that creates a model to specify a maximum version that its
// clients can accept.
// Versions:
// 0x1L: Initial version
// 0x2L: (Comment missing)
// 0x3L: (Comment missing)
// 0x4L: (update) Added schema to function tuple. Forward-compatible change.
// 0x5L: (update) Update bytecode is sharing constant tensor files from
// torchscript, and only serialize extra tensors that are not in the
// torchscript constant table. Also update tensor storage schema adapting to
// the unify format, the root key of tensor storage is updated from {index} to
// {the_pointer_value_the_tensor.storage}, for example:
// `140245072983168.storage` Forward-compatibility change.
// 0x6L: Implicit opereator versioning using number of specified argument.
// Refer to the summary of https://github.com/pytorch/pytorch/pull/56845 for
// details.
// 0x7L: Enable support for operators with default arguments plus out
// arguments. Refer. See https://github.com/pytorch/pytorch/pull/63651 for
// details.
// 0x8L: Emit promoted operators as instructions. See
// https://github.com/pytorch/pytorch/pull/71662 for details.
// 0x9L: Change serialization format from pickle to format This version is to
// serve migration. v8 pickle and v9 flatbuffer are the same. Refer to the
// summary of https://github.com/pytorch/pytorch/pull/75201 for more details.
constexpr uint64_t kProducedBytecodeVersion = 0x8L;
// static_assert(
// kProducedBytecodeVersion >= kProducedFileFormatVersion,
// "kProducedBytecodeVersion must be higher or equal to
// kProducedFileFormatVersion.");
// Introduce kMinSupportedBytecodeVersion and kMaxSupportedBytecodeVersion
// for limited backward/forward compatibility support of bytecode. If
// kMinSupportedBytecodeVersion <= model_version <= kMaxSupportedBytecodeVersion
// (in loader), we should support this model_version. For example, we provide a
// wrapper to handle an updated operator.
constexpr uint64_t kMinSupportedBytecodeVersion = 0x4L;
constexpr uint64_t kMaxSupportedBytecodeVersion = 0x9L;
} // namespace caffe2::serialize
#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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)