Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
@@ -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)
|
||||
+62
@@ -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)
|
||||
File diff suppressed because it is too large
Load Diff
+40
@@ -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)
|
||||
+35
@@ -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)
|
||||
+314
@@ -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)
|
||||
+31
@@ -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)
|
||||
+27
@@ -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)
|
||||
+84
@@ -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)
|
||||
+25
@@ -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)
|
||||
+383
@@ -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)
|
||||
+60
@@ -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)
|
||||
+198
@@ -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)
|
||||
+28
@@ -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)
|
||||
Reference in New Issue
Block a user