kolandr/gridbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/ConvUtils.h
+195
View File
@@ -0,0 +1,195 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <array>
#include <stdexcept>
#include <string>
#include <type_traits>
namespace fbgemm {
template <int N, int... Vals>
constexpr std::enable_if_t<N == sizeof...(Vals), std::array<int, N>>
array_of_ones() {
return std::array<int, N>{{Vals...}};
}
template <int N, int... Vals>
constexpr std::enable_if_t<N != sizeof...(Vals), std::array<int, N>>
array_of_ones() {
return array_of_ones<N, Vals..., 1>();
}
template <int N, int... Vals>
constexpr std::enable_if_t<N == sizeof...(Vals), std::array<int, N>>
array_of_zeroes() {
return std::array<int, N>{{Vals...}};
}
template <int N, int... Vals>
constexpr std::enable_if_t<N != sizeof...(Vals), std::array<int, N>>
array_of_zeroes() {
return array_of_zeroes<N, Vals..., 0>();
}
/**
* @brief A struct to conveniently store all convolution parameters.
*/
template <int SPATIAL_DIM = 2>
struct conv_param_t {
int MB; ///< Mini Batch size
int IC; ///< Number of Input Channels
int OC; ///< Number of Output Channels
std::array<int, SPATIAL_DIM> IN_DIM; ///< Input Image Dimension
int G; ///< Number of Groups
std::array<int, SPATIAL_DIM> K; ///< Filter (Kernel) dimensions
std::array<int, SPATIAL_DIM> stride; //< Strides
std::array<int, SPATIAL_DIM * 2>
pad; //< Padding (first SPATIAL_DIM is for prev/top/left padding, second
// SPATIAL_DIM is for next/bottom/right padding)
std::array<int, SPATIAL_DIM> dilation; //< Kernel dilation
// The following are derived parameters
std::array<int, SPATIAL_DIM> OUT_DIM; //< Output Image Dimension
std::array<int, SPATIAL_DIM> IN_DIMP; //< Input Image Dimension Padded
// The following is for tranposed convolution
std::array<int, SPATIAL_DIM>
output_pad; //< Padding (next/bottom/right padding in output buffer)
bool transposed;
/**
* @brief Constructor for initializing the convolution parameters.
*/
conv_param_t(
int mb,
int ic,
int oc,
std::array<int, SPATIAL_DIM> in_dim,
int g,
std::array<int, SPATIAL_DIM> k,
std::array<int, SPATIAL_DIM> strd,
std::array<int, SPATIAL_DIM * 2> pd,
std::array<int, SPATIAL_DIM> dilations = array_of_ones<SPATIAL_DIM>(),
std::array<int, SPATIAL_DIM> otpt_pd = array_of_zeroes<SPATIAL_DIM>(),
bool transposed = false)
: MB(mb),
IC(ic),
OC(oc),
IN_DIM(in_dim),
G(g),
K(k),
stride(strd),
pad(pd),
dilation(dilations),
output_pad(otpt_pd),
transposed(transposed) {
if (ic % g != 0) {
throw std::runtime_error(
"groups = " + std::to_string(g) +
" does not divide number of input channels = " + std::to_string(ic));
}
if (oc % g != 0) {
throw std::runtime_error(
"groups = " + std::to_string(g) +
" does not divide number of output channels = " + std::to_string(oc));
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
if (transposed) {
this->IN_DIMP[d] = this->IN_DIM[d] +
(this->dilation[d] * (this->K[d] - 1) - this->pad[d]) +
(this->dilation[d] * (this->K[d] - 1) - this->pad[SPATIAL_DIM + d]);
this->OUT_DIM[d] = (this->IN_DIM[d] - 1) * this->stride[d] -
this->pad[d] - this->pad[SPATIAL_DIM + d] +
this->dilation[d] * (this->K[d] - 1) + output_pad[d] + 1;
} else {
IN_DIMP[d] = IN_DIM[d] + pad[d] + pad[SPATIAL_DIM + d];
OUT_DIM[d] =
(IN_DIMP[d] - dilation[d] * (K[d] - 1) - 1) / stride[d] + 1;
}
}
}
/**
* @brief Helper function to get convolution parameters as string.
*/
std::string toString() const {
std::string dim_string[3] = {"T", "H", "W"};
std::string out;
out += "MB:" + std::to_string(MB) + ", ";
out += "IC:" + std::to_string(IC) + ", ";
out += "OC:" + std::to_string(OC) + ", ";
if constexpr (SPATIAL_DIM <= 3) {
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "I" + dim_string[3 - SPATIAL_DIM + d] + ":" +
std::to_string(IN_DIM[d]) + ", ";
}
} else {
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "I" + std::to_string(d) + ":" + std::to_string(IN_DIM[d]) + ", ";
}
}
out += "G:" + std::to_string(G) + ", ";
if constexpr (SPATIAL_DIM <= 3) {
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "K" + dim_string[3 - SPATIAL_DIM + d] + ":" +
std::to_string(K[d]) + ", ";
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "stride_" + dim_string[3 - SPATIAL_DIM + d] + ":" +
std::to_string(stride[d]) + ", ";
}
for (int d = 0; d < SPATIAL_DIM * 2; ++d) {
out += "pad_" + dim_string[3 - SPATIAL_DIM + (d % SPATIAL_DIM)] + ":" +
std::to_string(pad[d]) + ", ";
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "dilation_" + dim_string[3 - SPATIAL_DIM + d] + ":" +
std::to_string(dilation[d]);
if (d < SPATIAL_DIM - 1) {
out += ", ";
}
}
} else {
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "K" + std::to_string(d) + ":" + std::to_string(K[d]) + ", ";
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "stride_" + std::to_string(d) + ":" + std::to_string(stride[d]) +
", ";
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "pad_" + std::to_string(d) + ":" + std::to_string(pad[d]);
if (d < SPATIAL_DIM * 2 - 1) {
out += ", ";
}
}
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "dilation_" + std::to_string(d) + ":" +
std::to_string(dilation[d]) + ", ";
}
}
if (transposed) {
for (int d = 0; d < SPATIAL_DIM; ++d) {
out += "output_padding_" + std::to_string(d) + ":" +
std::to_string(output_pad[d]) + ", ";
}
}
return out;
}
};
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/Fbgemm.h
+1515
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmBuild.h
+116
View File
@@ -0,0 +1,116 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
// For details about dllexport/dllimport, checkout the following SO question
// https://stackoverflow.com/questions/57999/what-is-the-difference-between-dllexport-and-dllimport
#if !defined(FBGEMM_API)
#if defined(FBGEMM_STATIC)
#define FBGEMM_API
#define FBGEMM_ENUM_CLASS_API
#elif defined _WIN32 || defined __CYGWIN__
#if (__GNUC__ || __clang__) && !(__MINGW64__ || __MINGW32__)
#if defined(FBGEMM_EXPORTS)
#define FBGEMM_API __attribute__((__dllexport__))
#else
#define FBGEMM_API __attribute__((__dllimport__))
#endif
#else
#if defined(FBGEMM_EXPORTS)
#define FBGEMM_API __declspec(dllexport)
#else
#define FBGEMM_API __declspec(dllimport)
#endif
#endif
#define FBGEMM_ENUM_CLASS_API
#else
#if __clang__ || __GNUC__ || __INTEL_COMPILER
#define FBGEMM_API __attribute__((__visibility__("default")))
#else
#define FBGEMM_API
#endif
// Currently, enum classes need to be declaredly explicitly for shared build on
// macos
#if __clang__
#define FBGEMM_ENUM_CLASS_API __attribute__((__visibility__("default")))
#else
#define FBGEMM_ENUM_CLASS_API
#endif
#endif
#endif
// Use this to indicate to not inline functions
#if __clang__ || __GNUC__ || __INTEL_COMPILER
#define NOINLINE __attribute__((noinline))
#elif _MSC_VER
#define NOINLINE __declspec(noinline)
#else
#define NOINLINE
#endif
// Use this to indicate always inline functions
#if __clang__ || __GNUC__ || __INTEL_COMPILER
#define ALWAYS_INLINE inline __attribute__((__always_inline__))
#elif _MSC_VER
// commenting out because __forceinline takes too long time in MSVC
#define ALWAYS_INLINE // __forceinline
#else
#define ALWAYS_INLINE inline
#endif
// Use the C++11 keyword "alignas" if you can
#if _MSC_VER
#define ALIGNAS(byte_alignment) __declspec(align(byte_alignment))
#else
#define ALIGNAS(byte_alignment) __attribute__((aligned(byte_alignment)))
#endif
// Sanitizers annotations
#if defined(__has_attribute)
#if __has_attribute(no_sanitize)
#define NO_SANITIZE(what) __attribute__((no_sanitize(what)))
#endif
#endif
#if !defined(NO_SANITIZE)
#define NO_SANITIZE(what)
#endif
// Ignore __builtin_assume() when not supported by compiler.
#ifndef __has_builtin
#define __has_builtin(x) 0
#endif
#if !__has_builtin(__builtin_assume)
#define __builtin_assume(x) (static_cast<void>(0))
#endif
// Macro for silencing warnings
#if __clang__ || __GNUC__
// clang-format off
#define FBGEMM_PUSH_WARNING _Pragma("GCC diagnostic push")
#define FBGEMM_DISABLE_WARNING_INTERNAL2(warningName) #warningName
#define FBGEMM_DISABLE_WARNING(warningName) \
_Pragma( \
FBGEMM_DISABLE_WARNING_INTERNAL2(GCC diagnostic ignored warningName))
#define FBGEMM_PUSH_WARNING_AND_DISABLE(warningName) \
_Pragma("GCC diagnostic push") \
_Pragma( \
FBGEMM_DISABLE_WARNING_INTERNAL2(GCC diagnostic ignored warningName))
#define FBGEMM_POP_WARNING _Pragma("GCC diagnostic pop")
// clang-format on
#else
#define FBGEMM_PUSH_WARNING
#define FBGEMM_DISABLE_WARNING(NAME)
#define FBGEMM_PUSH_WARNING_AND_DISABLE(NAME)
#define FBGEMM_POP_WARNING
#endif
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmConvert.h
+205
View File
@@ -0,0 +1,205 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstddef>
#include <cstdint>
#include "fbgemm/FbgemmBuild.h"
#include "fbgemm/Types.h"
namespace fbgemm {
/**
* @ Transform all entries in a matrix from fp32 to bfloat16: reference
* implementation.
*
*/
FBGEMM_API void
FloatToBfloat16_ref(const float* src, bfloat16* dst, size_t size);
/**
* @ Transform all entries in a matrix from bfloat16 to fp32: reference
* implementation.
*
*/
FBGEMM_API void
Bfloat16ToFloat_ref(const bfloat16* src, float* dst, size_t size);
/**
* @ Transform all entries in a matrix from fp32 to bfloat16: simd
* implementation.
*
*/
FBGEMM_API void
FloatToBfloat16_simd(const float* src, bfloat16* dst, size_t size);
/**
* @ Transform all entries in a matrix from bfloat16 to fp32: simd
* implementation.
*
*/
FBGEMM_API void
Bfloat16ToFloat_simd(const bfloat16* src, float* dst, size_t size);
#if !defined(__aarch64__)
/**
* @brief AVX2 implementation to convert fp32 numbers to bf16 numbers.
*
*/
FBGEMM_API void
FloatToBfloat16_avx2(const float* src, bfloat16* dst, size_t size);
/**
* @brief AVX512 implementation to convert fp32 numbers to bf16 numbers.
*
*/
FBGEMM_API void
FloatToBfloat16_avx512(const float* src, bfloat16* dst, size_t size);
/**
* @brief AVX2 implementation to convert bf16 numbers to fp32 numbers.
*
*/
FBGEMM_API void
Bfloat16ToFloat_avx2(const bfloat16* src, float* dst, size_t size);
/**
* @brief AVX512 implementation to convert bf16 numbers to fp32 numbers.
*
*/
FBGEMM_API void
Bfloat16ToFloat_avx512(const bfloat16* src, float* dst, size_t size);
#endif
/**
* @ Transform all entries in a matrix from fp32 to float16: reference
* implementation.
*
* @param do_clip if true we saturate to fp16 min and max instead of generating
* infinities.
*/
FBGEMM_API void FloatToFloat16_ref(
const float* src,
float16* dst,
size_t size,
bool do_clip = false);
/**
* @ Transform all entries in a matrix from float16 to fp32: reference
* implementation.
*
*/
FBGEMM_API void Float16ToFloat_ref(const float16* src, float* dst, size_t size);
/**
* @ Transform all entries in a matrix from fp32 to float16: simd
* implementation.
*
* @param do_clip if true we saturate to fp16 min and max instead of generating
* infinities.
*/
FBGEMM_API void FloatToFloat16_simd(
const float* src,
float16* dst,
size_t size,
bool do_clip = false);
/**
* @ Transform all entries in a matrix from float16 to fp32: simd
* implementation.
*
*/
FBGEMM_API void
Float16ToFloat_simd(const float16* src, float* dst, size_t size);
/**
* @brief AVX2 implementation to convert fp32 numbers to fp16 numbers.
*
*/
#if !defined(__aarch64__)
FBGEMM_API void FloatToFloat16_avx2(
const float* src,
float16* dst,
size_t size,
bool do_clip = false);
/**
* @brief AVX512 implementation to convert fp32 numbers to fp16 numbers.
*
*/
FBGEMM_API void FloatToFloat16_avx512(
const float* src,
float16* dst,
size_t size,
bool do_clip = false);
#endif
/**
* @brief SVE2 implementation to convert fp32 numbers to fp16 numbers.
*
*/
FBGEMM_API void FloatToFloat16_sve2(
const float* src,
float16* dst,
size_t size,
bool do_clip = false);
#if !defined(__aarch64__)
/**
* @brief AVX2 implementation to convert fp16 numbers to fp32 numbers.
*
*/
FBGEMM_API void
Float16ToFloat_avx2(const float16* src, float* dst, size_t size);
/**
* @brief AVX512 implementation to convert fp16 numbers to fp32 numbers.
*
*/
FBGEMM_API void
Float16ToFloat_avx512(const float16* src, float* dst, size_t size);
#endif
/**
* @brief Transform all entries in a matrix from fp32 to float16 and back to
* fp32.
*/
FBGEMM_API void RoundToFloat16(
const float* input,
float* output,
size_t size,
bool clamp = false,
bool clamp_denorms = false);
/**
* @brief Quantize float32 to float8. The code is a copy of float_to_hfp8() in
* fbgemm_gpu/quantize_ops_utils.h
*/
FBGEMM_API void FloatToFloat8_ref(
float input,
uint8_t* output,
int exponent_bits,
int exponent_bias);
/**
* @brief Dequantize float8 to float32. The code is a copy of hf8_to_float() in
* fbgemm_gpu/quantize_ops_utils.h
*/
FBGEMM_API void Float8ToFloat_ref(
uint8_t input,
float* output,
int exponent_bits,
int exponent_bias);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmEmbedding.h
+383
View File
@@ -0,0 +1,383 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include <functional>
#include "fbgemm/FbgemmBuild.h"
namespace fbgemm {
template <
typename InType,
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float>
class EmbeddingSpMDMKernelSignature {
public:
/**
* Behavior is as the follow pseudocode
* (when use_offsets == true, lengths[i] == offsets[i + 1] - offsets[i])
* (when is_weight_positional == true, use weights[j - offsets[i]] instead of
* weights[j])
*
* for i in range(output_size):
* out[i * block_size : (i + 1) * block_size] = 0
* for j in range(offsets[i], offsets[i + 1]):
* for k in range(block_size):
* out[i * block_size + k] += input[indices[j] * block_size + k] *
* weights ? weights[j] : 1;
* if normalize_weights and lengths[i] > 0:
* out[i * block_size : (i + 1) * block_size] /= lengths[i]
*
* @param data_size the number of rows in embedding table
*/
using Type = std::function<bool(
std::int64_t output_size,
std::int64_t index_size,
std::int64_t data_size,
const InType* input,
const IndexType* indices,
const OffsetType* offsets_or_lengths,
const float* weights, // optional, can be null for non-weighted sum
OutType* out)>;
};
/**
* @tparam InType can be float, float16, or uint8_t
* @tparam IndexType can be int32_t or int64_t
* @tparam IndexType can be int32_t or int64_t
*
* @param use_offsets If true, the generated code assumes we will pass offsets
* instead of lengths that confirms PyTorch EmbeddingBag
* interface. In this case, the length of offsets array
* should be output_size + 1 and offsets[output_size] should
* be index_size.
* If false, the generate code assumes we will pass lengths
* that confirms Caffe2 SparseLengthsSum interface.
*/
template <
typename InType,
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float,
bool THREAD_LOCAL = false>
FBGEMM_API typename EmbeddingSpMDMKernelSignature<
InType,
IndexType,
OffsetType,
OutType>::Type
GenerateEmbeddingSpMDM(
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true,
bool is_bf16_out = false,
bool is_bf16_in = false);
/**
* @param output_stride If -1, output_stride is same as block_size
* @param input_stride If -1, input_stride is same as block_size
* @param scale_bias_last if false, scale and bias appear at the beginning
* of each row and are in fp16 for table batched embedding (TBE)
* in FBGEMM_GPU. If false, it can also take -1 indices (output from
* pruned embedding id mapping)
*/
template <
typename InType,
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float,
bool THREAD_LOCAL = false>
FBGEMM_API typename EmbeddingSpMDMKernelSignature<
InType,
IndexType,
OffsetType,
OutType>::Type
GenerateEmbeddingSpMDMWithStrides(
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true,
std::int64_t output_stride = -1,
std::int64_t input_stride = -1,
bool scale_bias_last = true,
bool no_bag = false,
bool is_bf16_out = false,
bool is_bf16_in = false);
/**
* @tparam IndexType can be int32_t or int64_t
* @tparam OffsetType can be int32_t or int64_t
* @param bit_rate can be 2 or 4
*/
template <
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float>
FBGEMM_API typename EmbeddingSpMDMKernelSignature<
std::uint8_t,
IndexType,
OffsetType,
OutType>::Type
GenerateEmbeddingSpMDMNBit(
int bit_rate,
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true);
/**
* @param output_stride If -1, output_stride is same as block_size
* @param input_stride in Bytes. If -1, input_stride is same as
* block_size / num_elem_per_byte + 2 * sizeof(float16)
* @param scale_bias_last if false, scale and bias appear at the beginning
* of each row and are in fp16 for table batched embedding (TBE)
* in FBGEMM_GPU. If false, it can also take -1 indices (output from
* pruned embedding id mapping)
*/
template <
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float,
bool THREAD_LOCAL = false>
FBGEMM_API typename EmbeddingSpMDMKernelSignature<
std::uint8_t,
IndexType,
OffsetType,
OutType>::Type
GenerateEmbeddingSpMDMNBitWithStrides(
const int input_bit_rate,
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true,
std::int64_t output_stride = -1,
std::int64_t input_stride = -1,
bool scale_bias_last = true,
const bool is_bf16_out = false,
const bool no_bag = false,
int output_bit_rate = -1);
/**
* @param output_stride If -1, output_stride is same as block_size
* @param input_stride in Bytes. If -1, input_stride is same as
* block_size / num_elem_per_byte + 2 * sizeof(float16)
* @param exponent_bits is the number of exponent bits in the FP8 encode
* (normally 4 or 5)
* @param exponent_bias is subtracted from the exponent to obtain the actual
* exponent for the floating-point number
*/
template <
typename IndexType,
typename OffsetType = std::int32_t,
typename OutType = float>
FBGEMM_API typename EmbeddingSpMDMKernelSignature<
std::uint8_t,
IndexType,
OffsetType,
OutType>::Type
GenerateEmbeddingSpMDMFP8WithStrides(
const std::int64_t block_size,
bool normalize_by_lengths,
bool is_weight_positional = false,
bool use_offsets = true,
std::int64_t output_stride = -1,
std::int64_t input_stride = -1,
int exponent_bits = 4,
int exponent_bias = 7,
bool is_bf16_out = false);
template <
typename InType,
typename IndexType,
typename OffsetType = std::int32_t>
class EmbeddingSpMDMRowWiseSparseKernelSignature {
public:
using Type = std::function<bool(
std::int64_t output_size,
std::int64_t index_size,
std::int64_t uncompressed_data_size,
// TODO: add compressed_data_size and check array bound
const InType* input,
const IndexType* indices,
const OffsetType* offsets_or_lengths,
const float* weights, // optional, can be null for non-weighted sum
float* out,
const std::int32_t* compressed_indices_table)>;
};
/**
* @tparam InType can be float, float16, or uint8_t
* @tparam IndexType can be int32_t or int64_t
* @tparam OffsetType can be int32_t or int64_t
*/
template <
typename InType,
typename IndexType,
typename OffsetType = std::int32_t>
FBGEMM_API typename EmbeddingSpMDMRowWiseSparseKernelSignature<
InType,
IndexType,
OffsetType>::Type
GenerateEmbeddingSpMDMRowWiseSparse(
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true);
/**
* @tparam IndexType can be int32_t or int64_t
* @tparam OffsetType can be int32_t or int64_t
* @param bit_rate can be 2 or 4
*/
template <typename IndexType, typename OffsetType = std::int32_t>
FBGEMM_API typename EmbeddingSpMDMRowWiseSparseKernelSignature<
std::uint8_t,
IndexType,
OffsetType>::Type
GenerateEmbeddingSpMDMNBitRowWiseSparse(
int bit_rate,
const std::int64_t block_size,
bool has_weight,
bool normalize_by_lengths,
int prefetch = 16,
bool is_weight_positional = false,
bool use_offsets = true);
/**
* @return The number of rows processed. If smaller than num_rows, an error
* must have happened at the last row processed.
*/
template <typename IndexType>
class SparseAdaGradSignature {
public:
using Type = std::function<int(
int num_rows, // number of rows reading
std::uint64_t param_size, // total number of parameters
float* w, // input/output parameters
const float* g, // input gradients
float* h, // input/output momentums
const IndexType* indices, // indices of each row
float epsilon,
float lr,
float weight_decay,
const double* counter, // used for weight_decay adjusted for frequency
// nullptr when frequency adjustment is not used.
// ignored when the kernel is generated with
// use_weight_decay = false.
std::int64_t counter_halflife)>; // frequency adjust happens only after
};
template <typename IndexType>
FBGEMM_API typename SparseAdaGradSignature<IndexType>::Type
GenerateSparseAdaGrad(
int block_size, // number of parameters per row
bool rowwise = false,
int prefetch = 16,
bool use_weight_decay = false);
// RowWiseSparseAdaGrad fused with SLS gradient
// Weights can be either float or float16
template <
typename IndexType,
typename OffsetType = std::int32_t,
typename DataType = float>
class RowWiseSparseAdaGradFusedSignature {
public:
using Type = std::function<bool(
std::int64_t output_size,
std::int64_t index_size,
std::int64_t data_size, // number of rows in w
DataType* w, // input/output parameters
const float* g, // input gradients
float* h, // input/output momentums
const IndexType* indices, // indices of each row
const OffsetType* offsets_or_lengths,
float epsilon,
float lr)>;
};
/**
* @param grad_stride If -1, grad_stride is same as block size
*/
template <
typename IndexType,
typename OffsetType = std::int32_t,
typename DataType = float>
FBGEMM_API typename RowWiseSparseAdaGradFusedSignature<
IndexType,
OffsetType,
DataType>::Type
GenerateRowWiseSparseAdaGradFused(
int block_size, // number of parameters per row
int prefetch = 16,
bool use_offsets = true,
bool use_stochastic_rounding = true,
int grad_stride = -1);
namespace internal {
// Specialization for block size 1 internally called by GenerateEmbeddingSpMDM
template <typename InType, typename IndexType, typename OffsetType>
FBGEMM_API bool EmbeddingSpMDMBlockSize1_(
const std::int64_t output_size,
const std::int64_t index_size,
const std::int64_t data_size, // the number of rows in input
const InType* input,
const IndexType* indices,
const OffsetType* offsets_or_lengths,
const float* weights, // optional, can be null for non-weighted sum
bool normalize_by_lengths,
float* out,
bool is_weight_positional = false,
bool use_offsets = true,
bool is_bf16 = false);
#if !defined(__aarch64__)
template <typename IndexType, bool HAS_WEIGHTS>
void compressed_indices_remap_avx512(
std::int32_t offsets_numel,
const IndexType* indices,
const int32_t* compressed_indices_mapping,
const IndexType* offsets,
const float* weights, // optional, can be null,
IndexType* out_indices,
IndexType* out_offsets,
float* out_weights);
#endif
} // namespace internal
template <typename IndexType>
FBGEMM_API void compressed_indices_remap(
std::int32_t offsets_numel,
const IndexType* indices,
const int32_t* compressed_indices_mapping,
const IndexType* offsets,
const float* weights, // optional, can be null,
IndexType* out_indices,
IndexType* out_offsets,
float* out_weights);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmFP16.h
+60
View File
@@ -0,0 +1,60 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
// WARNING: this is a legacy fp16 fbgemm implementation and will soon be
// upgraded to match with new fbgemm interface.
#include <cpuinfo.h>
#include "./FbgemmPackMatrixB.h" // @manual
#include "./FloatConversion.h" // @manual
#include "./Types.h" // @manual
#include "./Utils.h" // @manual
namespace fbgemm {
template <>
struct TypeConverter<float16> {
float16 operator()(float src) const {
constexpr float FP16_MAX = 65504.f;
const float fp16 = std::max(-FP16_MAX, std::min(src, FP16_MAX));
return cpu_float2half(fp16);
}
};
using PackedGemmMatrixFP16 = PackedGemmMatrixB<float16>;
template <typename T>
FBGEMM_API void cblas_gemm_compute(
const matrix_op_t transa,
const int m,
const float* A,
const PackedGemmMatrixB<T>& Bp,
const float beta,
float* C,
int thread_id = 0,
int num_threads = 1);
extern template void cblas_gemm_compute<float16>(
const matrix_op_t transa,
const int m,
const float* A,
const PackedGemmMatrixFP16& Bp,
const float beta,
float* C,
int thread_id,
int num_threads);
}; // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmFP32.h
+54
View File
@@ -0,0 +1,54 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
#pragma once
// WARNING: this is a legacy fp16 fbgemm implementation and will soon be
// upgraded to match with new fbgemm interface.
#include <cpuinfo.h>
#include "fbgemm/FbgemmFPCommon.h"
#include "fbgemm/FbgemmPackMatrixB.h"
#include "fbgemm/Utils.h"
namespace fbgemm {
template <>
struct TypeConverter<float> {
float operator()(float src) const {
return src;
}
};
using GemmParamsFP32 = GemmParams<float>;
using PackedGemmMatrixFP32 = PackedGemmMatrixB<float>;
template <typename T, int _kernel_ncol_blocks, int _brow>
void cblas_gemm_compute(
const matrix_op_t transa,
const int m,
const float* A,
const PackedGemmMatrixB<T>& Bp,
const float beta,
float* C,
int thread_id = 0,
int num_threads = 1);
extern template void cblas_gemm_compute(
const matrix_op_t transa,
const int m,
const float* A,
const PackedGemmMatrixFP32& Bp,
const float beta,
float* C,
int thread_id,
int num_threads);
template <>
const isa_descriptor<float>& getIsaHandlers(inst_set_t isa);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmFPCommon.h
+319
View File
@@ -0,0 +1,319 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* Copyright 2024-2025 Arm Limited and/or its affiliates
* <open-source-office@arm.com> All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <fbgemm/FbgemmPackMatrixB.h>
#include <fbgemm/Types.h>
#include <fbgemm/Utils.h>
#include <array>
#include <memory>
#if defined(FBGEMM_FP16_FALLBACK_TO_REF_KERNEL) || \
defined(FBGEMM_FP32_FALLBACK_TO_REF_KERNEL)
#if defined(__APPLE__) && defined(__aarch64__)
#define FBGEMM_USE_REF_KERNEL
#endif
#endif
namespace fbgemm {
using partition_array_t = std::array<std::array<std::array<int, 2>, 2>, 121>;
extern partition_array_t partition_avx2;
extern partition_array_t partition_avx512;
extern partition_array_t partition_sve128;
#ifdef FBGEMM_ENABLE_KLEIDIAI
extern partition_array_t partition_neon;
#endif
template <typename T>
struct GemmParams {
uint64_t k;
float* A;
const T* B;
float beta;
float* C;
uint64_t ldc;
uint64_t b_block_cols;
uint64_t b_block_size;
};
template <>
struct GemmParams<float16> {
uint64_t k;
float* A;
const float16* B;
float beta;
float* C;
uint64_t ldc;
uint64_t b_block_cols;
#ifdef FBGEMM_ENABLE_KLEIDIAI
uint64_t lda;
#else
uint64_t b_block_size;
#endif
};
template <>
struct GemmParams<float> {
uint64_t k;
float* A;
const float* B;
float beta;
float* C;
uint64_t ldc;
uint64_t b_block_cols;
#ifdef FBGEMM_ENABLE_KLEIDIAI
uint64_t lda;
#else
uint64_t b_block_size;
#endif
};
template <typename T>
using funcptr_t = void (*)(GemmParams<T>*);
template <typename T>
using kernel_array_t = std::array<funcptr_t<T>, 15>;
template <typename T>
using isa_descriptor = std::tuple<kernel_array_t<T>, partition_array_t>;
template <typename T>
extern const isa_descriptor<T>& getIsaHandlers(inst_set_t isa);
void PackA(int nrow, int ncol, const float* from, int ldim, float* to);
// define fp16/fp32 kernels using a reference C implementation
#if defined(FBGEMM_FP16_FALLBACK_TO_REF_KERNEL) || \
defined(FBGEMM_FP32_FALLBACK_TO_REF_KERNEL)
template <typename T>
FBGEMM_API void ref_kernel(
int kernel_nrows,
GemmParams<T>* gp,
const float* C_base,
int m_total,
int n_total,
int vlen);
#endif
template <typename T>
FBGEMM_API void cblas_gemm_compute(
const matrix_op_t transa,
const int m,
const float* A,
const PackedGemmMatrixB<T>& Bp,
const float beta,
float* C,
int thread_id = 0,
int num_threads = 1);
#if defined(FBGEMM_EXPORTS)
// autotuned kernel splits for various cases m = 1:mb_max
template <typename T>
void cblas_gemm_compute(
const matrix_op_t transa [[maybe_unused]],
const int m,
const float* A,
const PackedGemmMatrixB<T>& Bp,
const float beta,
float* C,
int thread_id,
int num_threads) {
// ground truth
assert(cpuinfo_initialize());
#ifndef __aarch64__
assert(cpuinfo_has_x86_fma3());
assert(cpuinfo_has_x86_f16c());
#endif
assert(transa == matrix_op_t::NoTranspose);
// private scratchpad storage
static thread_local std::unique_ptr<std::array<float, 256 * 1024>> scratchpad(
new std::array<float, 256 * 1024>());
// constants
const int n = Bp.numCols(), k = Bp.numRows(), ldc = n;
const int mb_max = 120;
#if defined(FBGEMM_USE_REF_KERNEL) && defined(__APPLE__)
const auto& [_, partition] = getIsaHandlers<float16>(inst_set_t::sve);
#else
const auto iset = fbgemmInstructionSet();
const auto& [kernels, partition] = getIsaHandlers<T>(iset);
#endif
#ifdef FBGEMM_USE_REF_KERNEL
// By some reason, if packed B is using packing layout for avx2, we just use
// avx2 even if avx512 is available.
const int simd_width =
#ifndef __aarch64__
(iset == inst_set_t::avx512 || iset == inst_set_t::avx512_vnni) &&
(Bp.blockColSize() == 16 * Bp.kernelNumColBlocks())
? simd_info<inst_set_t::avx512>::WIDTH_32BIT_ELEMS
: simd_info<inst_set_t::avx2>::WIDTH_32BIT_ELEMS;
#else
simd_info<inst_set_t::sve>::WIDTH_32BIT_ELEMS;
#endif
#endif
GemmParams<T> gp;
int i_begin = 0, i_end = 0;
i_begin = 0;
i_end = m;
for (auto m0 = i_begin; m0 < i_end; m0 += mb_max) {
int mb = std::min(mb_max, i_end - m0);
assert(mb < static_cast<int64_t>(partition.size()));
for (auto k_ind = 0; k_ind < k; k_ind += Bp.blockRowSize()) {
// set up proper accumulation to avoid "Nan" problem
// accumulate of beta != 0.0
// do not!!! accumulate otherwise
float beta_ = beta;
if (k_ind != 0) {
// always accumulate with beta_ = 1.0f
beta_ = 1.0f;
}
const int kb = std::min(Bp.blockRowSize(), Bp.numRows() - k_ind);
auto m1 = m0;
auto const num_cycles = partition[mb].size();
for (size_t c = 0; c < num_cycles; ++c) {
auto kernel_nrows = partition[mb][c][0];
auto nkernel_nrows = partition[mb][c][1];
auto m_start = m1;
auto m_end = m1 + kernel_nrows * nkernel_nrows;
for (auto m2 = m_start; m2 < m_end; m2 += kernel_nrows) {
assert(kernel_nrows * kb < static_cast<int64_t>(scratchpad->size()));
if (m != 1) {
#ifdef FBGEMM_ENABLE_KLEIDIAI
if constexpr (
std::is_same<T, float16>::value ||
std::is_same<T, float>::value) {
gp.A = const_cast<float*>(&A[m2 * k + k_ind]);
} else {
#endif
PackA(
kernel_nrows, kb, &A[m2 * k + k_ind], k, scratchpad->data());
gp.A = scratchpad->data();
#ifdef FBGEMM_ENABLE_KLEIDIAI
}
#endif
} else {
// When m == 1, it is actually vector matrix multiplication. We
// don't need to do the transposition for packA here. Instead, we
// can just pass the pointer of the original A matrix buffer to the
// packed A buffer.
gp.A = const_cast<float*>(&A[k_ind]);
}
int nbcol = n / Bp.blockColSize();
gp.k = kb;
gp.B = &(Bp(k_ind, 0));
gp.beta = beta_;
gp.C = &C[m2 * ldc];
gp.ldc = ldc * sizeof(C[0]);
gp.b_block_cols = nbcol;
#ifdef FBGEMM_ENABLE_KLEIDIAI
if constexpr (
std::is_same<T, float16>::value ||
std::is_same<T, float>::value) {
gp.lda = k * sizeof(A[0]);
} else {
#endif
gp.b_block_size = gp.k * Bp.blockColSize() * sizeof(gp.B[0]);
#ifdef FBGEMM_ENABLE_KLEIDIAI
}
#endif
if ((n % Bp.blockColSize()) == 0) {
int64_t jb_begin = 0, jb_end = 0;
fbgemmPartition1D(
thread_id, num_threads, gp.b_block_cols, jb_begin, jb_end);
gp.B += gp.k * Bp.blockColSize() * jb_begin;
gp.C += Bp.blockColSize() * jb_begin;
gp.b_block_cols = jb_end - jb_begin;
if (gp.b_block_cols) {
#ifdef FBGEMM_USE_REF_KERNEL
ref_kernel<T>(kernel_nrows, &gp, C, m, n, simd_width);
#else
kernels[kernel_nrows](&gp);
#endif
}
} else {
int last_blk_col = nbcol * Bp.blockColSize();
if (nbcol) {
int64_t jb_begin = 0, jb_end = 0;
fbgemmPartition1D(
thread_id, num_threads, gp.b_block_cols, jb_begin, jb_end);
gp.B += gp.k * Bp.blockColSize() * jb_begin;
gp.C += Bp.blockColSize() * jb_begin;
gp.b_block_cols = jb_end - jb_begin;
if (gp.b_block_cols) {
#ifdef FBGEMM_USE_REF_KERNEL
ref_kernel(kernel_nrows, &gp, C, m, n, simd_width);
#else
kernels[kernel_nrows](&gp);
#endif
}
}
// use one thread to handle the fringe cases
if (thread_id == num_threads - 1) {
// leftover
const int rem [[maybe_unused]] = n - last_blk_col;
assert(rem < Bp.blockColSize());
// small temporary buffer: the size should be larger than the
// required kernel_nrow x kernel_ncols elements computed in the
// registers.
std::array<float, 14 * 32> c_tmp{0.f};
assert(
static_cast<int64_t>(c_tmp.size()) >=
kernel_nrows * Bp.blockColSize());
gp.B = &(Bp(k_ind, last_blk_col));
gp.C = c_tmp.data();
gp.ldc = Bp.blockColSize() * sizeof(C[0]);
gp.b_block_cols = 1;
#ifdef FBGEMM_USE_REF_KERNEL
ref_kernel<T>(
kernel_nrows, &gp, c_tmp.data(), 14, 32, simd_width);
#else
kernels[kernel_nrows](&gp);
#endif
for (int i = 0; i < kernel_nrows; i++) {
// Todo: use assembly
for (int j = last_blk_col; j < n; j++) {
assert(
i * Bp.blockColSize() + (j - last_blk_col) <
static_cast<int64_t>(sizeof(c_tmp) / sizeof(c_tmp[0])));
if (beta_ == 0.f) {
C[(m2 + i) * ldc + j] =
c_tmp[i * Bp.blockColSize() + (j - last_blk_col)];
} else {
C[(m2 + i) * ldc + j] = beta_ * C[(m2 + i) * ldc + j] +
c_tmp[i * Bp.blockColSize() + (j - last_blk_col)];
}
}
}
}
}
}
m1 += kernel_nrows * nkernel_nrows;
}
}
}
}
#endif
#undef FBGEMM_USE_REF_KERNEL
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmI64.h
+36
View File
@@ -0,0 +1,36 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include "fbgemm/Utils.h"
namespace fbgemm {
FBGEMM_API void cblas_gemm_i64_i64acc(
matrix_op_t transa,
matrix_op_t transb,
int M,
int N,
int K,
const std::int64_t* A,
int lda,
const std::int64_t* B,
int ldb,
bool accumulate,
std::int64_t* C,
int ldc);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmI8DepthwiseAvx2.h
+117
View File
@@ -0,0 +1,117 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include "fbgemm/ConvUtils.h"
#include "fbgemm/FbgemmBuild.h"
#include "fbgemm/UtilsAvx2.h"
namespace fbgemm {
class FBGEMM_API PackedDepthWiseConvMatrix {
public:
/**
* @param IC the number of input channels (same as the number of groups
* because depth-wise convolution has one input channel per group)
* @param OC the number of output channels
* @param kernel_prod the product of all kernels. For example, kernel_prod =
* 9 for 3x3 conv, and 27 for 3x3x3 conv.
* @param smat the source unpacked weight in GRS layout
*/
PackedDepthWiseConvMatrix(int OC, int kernel_prod, const std::int8_t* smat);
PackedDepthWiseConvMatrix(const PackedDepthWiseConvMatrix&) = delete;
PackedDepthWiseConvMatrix(PackedDepthWiseConvMatrix&&) = delete;
PackedDepthWiseConvMatrix& operator=(const PackedDepthWiseConvMatrix&) =
delete;
PackedDepthWiseConvMatrix& operator=(PackedDepthWiseConvMatrix&&) = delete;
virtual ~PackedDepthWiseConvMatrix();
const std::int8_t* PackedMat() const {
return pmat_;
}
int GetKernelProduct() const {
return kernel_prod_;
}
/**
* @brief Unpacks pmat_ into unpack_data.
* Used for recovering the weight matrix into the original format
*/
void unpack(std::int8_t* unpacked_data);
/**
* @brief returns the index into pmat_ given the row and column for smat
*/
int addr(int r, int c);
private:
const int OC_; /**< the number of output channels */
const int kernel_prod_; /** the product of all kernel dims */
std::int8_t* pmat_; /** packed weight */
}; // PackedDepthWiseConvMatrix
/**
* Depth-wise convolution that results in the same output feature size as the
* input feature. That is PAD_T = PAD_B = (R - 1) / 2 and PAD_L = PAD_R =
* (S - 1) / 2. This function also does requantization.
* @param col_offsets nullptr if col_offsets are folded into bias
* @param act_times_w_scale Only used if BIAS_TYPE is float, i.e., bias is
* unquantized.
*/
template <QuantizationGranularity Q_GRAN, typename BIAS_TYPE = std::int32_t>
FBGEMM_API void depthwise_2d_same_pad(
int N,
int H,
int W,
int IC,
int OC,
int stride_h,
int stride_w,
std::int32_t A_zero_point,
const std::uint8_t* A,
const std::int32_t* B_zero_point,
const PackedDepthWiseConvMatrix& Bp,
const float* C_multiplier,
std::int32_t C_zero_point,
std::uint8_t* C,
const std::int32_t* col_offsets,
const BIAS_TYPE* bias,
bool fuse_relu = false,
const float* act_times_w_scale = nullptr,
int thread_id = 0,
int num_threads = 1);
/**
* @param col_offsets nullptr if col_offsets are folded into bias
*/
template <QuantizationGranularity Q_GRAN, typename BIAS_TYPE = std::int32_t>
FBGEMM_API void depthwise_3d_same_pad(
const conv_param_t<3>& conv_p,
std::int32_t A_zero_point,
const std::uint8_t* A,
const std::int32_t* B_zero_point,
const PackedDepthWiseConvMatrix& Bp,
const float* C_multiplier,
std::int32_t C_zero_point,
std::uint8_t* C,
const std::int32_t* col_offsets,
const BIAS_TYPE* bias,
bool fuse_relu = false,
const float* act_times_w_scale = nullptr,
int thread_id = 0,
int num_threads = 1);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmI8DirectconvAvx2.h
+69
View File
@@ -0,0 +1,69 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include <vector>
#include "fbgemm/ConvUtils.h"
#include "fbgemm/FbgemmBuild.h"
namespace fbgemm {
class FBGEMM_API PackedDirectConvMatrix {
public:
/**
* @param IC the number of input channels
* @param OC the number of output channels
* @param kernel_prod the product of all kernels. For example, kernel_prod =
* 9 for 3x3 conv, and 27 for 3x3x3 conv.
* @param smat the source unpacked weight in GRS layout
*/
PackedDirectConvMatrix(
int IC_per_G,
int OC_per_G,
int filter_prod,
const std::int8_t* smat);
PackedDirectConvMatrix(const PackedDirectConvMatrix&) = delete;
PackedDirectConvMatrix(PackedDirectConvMatrix&&) = delete;
PackedDirectConvMatrix& operator=(const PackedDirectConvMatrix&) = delete;
PackedDirectConvMatrix& operator=(PackedDirectConvMatrix&&) = delete;
virtual ~PackedDirectConvMatrix();
const std::int8_t* PackedMat() const {
return pmat_;
}
const bool& is_first_call() const {
return first_call;
}
/**
compute the column offsets of the weight matrix.
output of this function is the col_offsets vector
col_offses dimension is the same as conv_p.OUT_DIM
*/
template <int kSpatialDim>
FBGEMM_API void col_offsets_with_zero_pt_s8acc32_DirectConvT(
const fbgemm::conv_param_t<kSpatialDim>& conv_p,
std::int32_t* B_zero_point,
std::vector<int32_t>& col_offsets,
int ncols_per_quant_group);
private:
std::int8_t* pmat_; /** packed weight */
bool first_call{true};
};
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmI8Spmdm.h
+140
View File
@@ -0,0 +1,140 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include <vector>
#include "./ConvUtils.h" // @manual
#include "./FbgemmBuild.h" // @manual
#include "./Utils.h" // @manual
// #define FBGEMM_MEASURE_TIME_BREAKDOWN
#ifdef FBGEMM_MEASURE_TIME_BREAKDOWN
#include <chrono>
#include <iostream>
extern double spmdm_initial_time;
extern double spmdm_transpose_uint8_time;
extern double spmdm_transpose_32xN_time;
extern double spmdm_compute_time;
extern double spmdm_transpose_Nx32_time;
extern double spmdm_run_time;
extern double sconv_run_time;
#endif
namespace fbgemm {
/**
* @brief A class to represent a matrix in Compressed Sparse Column (CSC)
* format.
*
* The second input matrix of matrix multiplication is usually weight and can
* be sparse, and it's usually more efficient to use CSC format to represent
* the second input matrix.
*/
class FBGEMM_API CompressedSparseColumn {
public:
CompressedSparseColumn(int num_of_rows, int num_of_cols);
std::vector<std::int32_t>& ColPtr() {
return colptr_;
}
std::vector<std::int16_t>& RowIdx() {
return rowidx_;
}
std::vector<std::int8_t>& Values() {
return values_;
}
std::vector<std::int16_t>& KHs() {
return kh_;
}
std::vector<std::int16_t>& KWs() {
return kw_;
}
/**
* ICs include group: i.e. for ith input channels withint group g, ICs contain
* g*(groups_per_input_channels) + i
*/
std::vector<std::int16_t>& ICs() {
return ic_;
}
std::size_t NumOfRows() const {
return num_rows_;
}
std::size_t NumOfCols() const {
return colptr_.size() - 1;
}
std::int32_t NumOfNonZeros() const {
return colptr_.back();
}
/**
* @return Total number of non-zero elements as a fraction of total
* elements.
*/
double Density() const;
/**
* @return True if the number of non-zeros per row is smaller than a small
* threshold.
*/
bool IsHyperSparse() const;
/**
* @brief Perform dense-matrix * sparse matrix.
*
* C += A (dense matrix) * B (this CSC matrix) if accumulation = true \n
* C = A (dense matrix) * B (this CSC matrix) if accumulation = false
*/
void SpMDM(
const block_type_t& block,
const std::uint8_t* A,
int lda,
bool accumulation,
std::int32_t* C,
int ldc) const;
void SparseConv(
const conv_param_t<>& conv_p,
const block_type_t& block,
const std::uint8_t* A,
std::int32_t A_zero_point,
bool accumulation,
std::int32_t* C,
int ldc) const;
private:
const std::size_t num_rows_;
std::vector<std::int32_t> colptr_; // corresponds to out channels
std::vector<std::int8_t> values_;
// For SpMDM
std::vector<std::int16_t> rowidx_; // kh kw ic are flattened with im2col
// For direct sparse convolution
std::vector<std::int16_t> kh_;
std::vector<std::int16_t> kw_;
std::vector<std::int16_t> ic_; // in channels
// Cache IsHyperSparse to minimize its overhead.
mutable bool hyper_sparse_{false};
// Whether we can reuse the cached hyper_sparse_ is determined by checking
// if NumOfNonZeros() is same as old_nnz_ saved in previous invocation of
// IsHyperSparse call.
mutable std::int32_t old_nnz_{-1};
};
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmPackMatrixB.h
+339
View File
@@ -0,0 +1,339 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* Copyright 2024-2025 Arm Limited and/or its affiliates
* <open-source-office@arm.com> All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <assert.h>
#include <cpuinfo.h>
#include <stdexcept>
#include "SimdUtils.h" // @manual
#include "Types.h" // @manual
#include "Utils.h" // @manual
namespace fbgemm {
template <typename T>
struct TypeConverter {
template <typename F>
T operator()(F) const;
};
#define PMAT_ALIGNMENT 64
/// class that performs packing of matrix in
/// row-major format into
/// internal packed blocked-row major format
template <typename T, typename C = TypeConverter<T>>
class PackedGemmMatrixB {
public:
using value_type = T;
using size_type = uint64_t;
// takes smat input mamtrix in row-major format;
// packs it into gemm-friendly blocked format;
// allocate space and sets up all the internal variables;
// also premultiplies by alpha during packing.
// brow_ contains tile size along k dimension
// and also is # of fmas updates into int16 container
// before flushing into fp32.
// the smaller the brow_, the higher overhead
// of flushing is.
// kernel_ncol_blocks is the number of column blocks (in the size of 8 fp16,
// or 128 bit, or 1 xmm register size) in the kernel. Because the batch size
// can be dynamic and we need to prepack the weight matrix B, the internal
// packing layout of the weight matrix and kernel_ncol_blocks have to be
// fixed. We can choose kernel_ncol_blocks = 1 (with kernels of 1x1~14x1
// register layouts), 2 (with kernels of 1x2~6x2 register layout), or 3 (with
// kernels of 1x3~4x3 register layout).
PackedGemmMatrixB(
const matrix_op_t trans,
const int nrow,
const int ncol,
const float alpha,
const float* smat,
const int brow = 512)
: nrow_(nrow), ncol_(ncol), brow_(brow), kernel_ncol_blocks_(2) {
#ifdef FBGEMM_ENABLE_KLEIDIAI
if constexpr (std::is_same<T, float16>::value) {
kernel_ncol_blocks_ = 1;
}
#endif
initializeParam();
initializeMemory();
// copy source matrix into packed matrix
this->packFromSrc(trans, alpha, smat);
}
PackedGemmMatrixB(
const int nrow,
const int ncol,
const int brow,
const int last_brow,
const int bcol,
const int nbrow,
const int nbcol,
const uint64_t size)
: nrow_(nrow),
ncol_(ncol),
brow_(brow),
last_brow_(last_brow),
bcol_(bcol),
nbrow_(nbrow),
nbcol_(nbcol),
size_(size),
kernel_ncol_blocks_(2) {
#ifdef FBGEMM_ENABLE_KLEIDIAI
if constexpr (std::is_same<T, float16>::value) {
kernel_ncol_blocks_ = 1;
}
#endif
initializeMemory();
}
PackedGemmMatrixB(
const int nrow,
const int ncol,
const int brow,
const int last_brow,
const int bcol,
const int nbrow,
const int nbcol,
const uint64_t size,
const int kernel_ncol_blocks,
void* pmat)
: nrow_(nrow),
ncol_(ncol),
brow_(brow),
last_brow_(last_brow),
bcol_(bcol),
nbrow_(nbrow),
nbcol_(nbcol),
size_(size),
kernel_ncol_blocks_(kernel_ncol_blocks) {
#ifdef FBGEMM_ENABLE_KLEIDIAI
if constexpr (std::is_same<T, float16>::value) {
kernel_ncol_blocks_ = 1;
}
#endif
pmat_ = static_cast<T*>(pmat);
packed_ = true;
pmat_passed_in = true;
}
PackedGemmMatrixB(const PackedGemmMatrixB&) = delete;
PackedGemmMatrixB(PackedGemmMatrixB&&) = delete;
PackedGemmMatrixB& operator=(const PackedGemmMatrixB&) = delete;
PackedGemmMatrixB& operator=(PackedGemmMatrixB&&) = delete;
void initializeParam() {
if (!cpuinfo_initialize()) {
throw std::runtime_error("Failed to initialize cpuinfo!");
}
bcol_ = (isZmm(fbgemmInstructionSet())
? simd_info<inst_set_t::avx512>::WIDTH_32BIT_ELEMS
: simd_info<inst_set_t::avx2>::WIDTH_32BIT_ELEMS) *
kernelNumColBlocks();
// set up internal packing parameters
nbrow_ = (numRows() + blockRowSize() - 1) / blockRowSize();
last_brow_ = ((nrow_ % blockRowSize()) == 0) ? blockRowSize()
: (nrow_ % blockRowSize());
nbcol_ = (numCols() + blockColSize() - 1) / blockColSize();
if (numCols() != blockColSize() * nbcol_) {
#ifdef VLOG
VLOG(0) << "Packer warning: ncol(" << numCols()
<< ") is not a multiple of internal block size ("
<< blockColSize() << ")";
VLOG(0) << "lefover is not super optimized hence overhead will inccur";
#endif
}
}
void setPacked(bool p) {
packed_ = p;
}
bool packed() const {
return packed_;
}
void initializeMemory() {
// allocate and initialize packed memory
size_ = (blockRowSize() * nbrow_) * (blockColSize() * nbcol_);
pmat_ = static_cast<T*>(
fbgemmAlignedAlloc(PMAT_ALIGNMENT, matSize() * sizeof(T)));
memset(pmat_, 0, matSize() * sizeof(T));
}
~PackedGemmMatrixB() {
if (pmat_passed_in == false) {
fbgemmAlignedFree(pmat_);
}
}
void unpackFromSrc(const matrix_op_t trans, T* src_mat) {
bool tr = (trans == matrix_op_t::Transpose);
for (int i = 0; i < numRows(); i++) {
for (int j = 0; j < numCols(); j++) {
pmat_[tr ? i + numRows() * j : i * numCols() + j] = src_mat[addr(i, j)];
}
}
packed_ = false;
}
void unpack(T* origin_buf, const matrix_op_t trans) {
assert(packed_);
bool tr = (trans == matrix_op_t::Transpose);
for (int i = 0; i < numRows(); i++) {
for (int j = 0; j < numCols(); j++) {
origin_buf[tr ? i + numRows() * j : i * numCols() + j] =
pmat_[addr(i, j)];
}
}
}
// protected:
// blocked row-major format address arithmetic
uint64_t addr(const int r_, const int c_) const {
uint64_t r = (uint64_t)r_;
uint64_t c = (uint64_t)c_;
uint64_t block_row_id = r / blockRowSize();
uint64_t brow_offset =
(block_row_id * nbcol_) * (blockRowSize() * blockColSize());
uint64_t block_col_id = c / blockColSize();
uint64_t bcol_offset = block_col_id *
((static_cast<int64_t>(block_row_id) != nbrow_ - 1)
? (blockRowSize() * blockColSize())
: (last_brow_ * blockColSize()));
uint64_t block_offset = brow_offset + bcol_offset;
uint64_t inblock_offset =
r % blockRowSize() * blockColSize() + c % blockColSize();
uint64_t index = block_offset + inblock_offset;
assert(static_cast<int64_t>(index) < matSize());
return index;
}
void
packFromSrc(const matrix_op_t trans, const float alpha, const float* smat) {
bool tr = (trans == matrix_op_t::Transpose);
// pack
for (int i = 0; i < numRows(); i++) {
for (int j = 0; j < numCols(); j++) {
float src = alpha *
((tr == false) ? smat[i * numCols() + j] : smat[i + numRows() * j]);
pmat_[addr(i, j)] = C()(src);
}
}
packed_ = true;
}
// This function takes in an unpacked T matrix of the same size and
// packs it. There is no floating type conversion.
void packFromSrc(const matrix_op_t trans, const T* smat) {
bool tr = (trans == matrix_op_t::Transpose);
for (int i = 0; i < numRows(); ++i) {
for (int j = 0; j < numCols(); ++j) {
pmat_[addr(i, j)] = smat[tr ? i + numRows() * j : i * numCols() + j];
}
}
packed_ = true;
}
const T& operator()(const int r, const int c) const {
const auto a = addr(r, c);
assert(r < numRows());
assert(c < numCols());
assert(static_cast<int64_t>(a) < this->matSize());
return pmat_[a];
}
int matSize() const {
return size_;
}
int numRows() const {
return nrow_;
}
int numCols() const {
return ncol_;
}
int lastBrow() const {
return last_brow_;
}
int numBrow() const {
return nbrow_;
}
int numBcol() const {
return nbcol_;
}
T* pmat() const {
return pmat_;
}
int blockRowSize() const {
return brow_;
}
int blockColSize() const {
return bcol_;
}
int kernelNumColBlocks() const {
return kernel_ncol_blocks_;
}
const value_type* data() const {
return pmat_;
}
uint64_t size() const {
return size_ / sizeof(value_type);
}
int nrow_, ncol_;
int brow_, last_brow_, bcol_;
int nbrow_, nbcol_;
uint64_t size_;
int kernel_ncol_blocks_;
T* pmat_;
bool packed_{false};
bool pmat_passed_in{false};
};
#ifndef _M_X64
template <>
FBGEMM_API
PackedGemmMatrixB<float16, TypeConverter<float16>>::PackedGemmMatrixB(
const matrix_op_t trans,
const int nrow,
const int ncol,
const float alpha,
const float* smat,
const int brow);
template <>
FBGEMM_API
PackedGemmMatrixB<float16, TypeConverter<float16>>::PackedGemmMatrixB(
const int nrow,
const int ncol,
const int brow,
const int last_brow,
const int bcol,
const int nbrow,
const int nbcol,
const uint64_t size);
#endif
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FbgemmSparse.h
+230
View File
@@ -0,0 +1,230 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include <memory>
#include <vector>
#include "fbgemm/FbgemmBuild.h"
#include "fbgemm/UtilsAvx2.h"
#include "fbgemm/spmmUtilsAvx2.h"
namespace fbgemm {
template <typename T>
struct FBGEMM_API CSRMatrix {
std::vector<int> rowPtr;
std::vector<int> colIdx;
std::vector<T> values;
};
/**
* Tiled block CSR format
* Partial blocks are zero-filled
*
*/
template <typename T = std::int8_t, int ROW_BLOCK = 1, int COL_BLOCK = 4>
struct FBGEMM_API BCSRMatrix {
using DTYPE = T;
static constexpr int RB = ROW_BLOCK; // Block size for rows
static constexpr int CB = COL_BLOCK; // Block size for cols
// We only tile in column dimension currently
// COLTILE must be a multiple of COL_BLOCK
static constexpr int COLTILE = 4000;
std::vector<int> rowBPtr; // rowPtr for blocks
std::vector<int> colBIdx; // colIdx for blocks
std::vector<DTYPE> values;
// Sum of all elements in a row
std::vector<int32_t> row_offsets;
int R;
int C;
BCSRMatrix(int Rows, int Cols) {
R = Rows;
C = Cols;
row_offsets.resize(R, 0);
}
/**
* @brief pack from dense to tiled block CSR format
* @param R number of rows in the matrix
* @param C number of columns in the matrix
* @param src is the source matrix with data type DTYPE
* @param ld is the leading dimension
*/
void pack(const DTYPE* src, size_t ld);
/**
* @brief pack from dense to tiled block CSR format
* @param R number of rows in the matrix
* @param C number of columns in the matrix
* @param src is the source matrix with data type DTYPE
*
* leading dim of the matrix is assumed to be equal to C
*/
void pack(const DTYPE* src);
/**
* @brief unpack from tiled block CSR to dense
* @param dst should be able to hold R*C elements of type DTYPE
* @param ld is the leading dimension
*/
void unpack(DTYPE* dst, size_t ld);
/*
* @brief unpack from tiled block CSR to dense
* @param dst should be able to hold R*C elements of type DTYPE
*
* leading dimension of the matrix is assumed to be equal to C
*/
void unpack(DTYPE* dst);
};
template <typename T>
FBGEMM_API std::unique_ptr<CSRMatrix<T>>
fbgemmDenseToCSR(int R, int C, const T* inp, int ld);
template <typename T>
FBGEMM_API std::unique_ptr<CSRMatrix<T>>
fbgemmDenseToCSR(int R, int C, const T* inp);
template <typename T = std::int8_t, int RB = 1, int CB = 4>
FBGEMM_API std::unique_ptr<BCSRMatrix<T, RB, CB>>
fbgemmDenseToBCSR(int R, int C, const T* inp, int ld);
template <typename T = std::int8_t, int RB = 1, int CB = 4>
FBGEMM_API std::unique_ptr<BCSRMatrix<T, RB, CB>>
fbgemmDenseToBCSR(int R, int C, const T* inp);
/**
* @param accum Controls accumulation.
* 1 means we're accumulating to the C Matrix.
*
* Note on matrix order and layout:
* Unlike other fbgemm functions that follow PyTorch convention where A
* matrix is activation (so in uint8_t for quantized FC/Conv or fp32) and B
* matrix is weight (so in int8_t for quantized FC/Conv or fp32), here A is
* weight matrix. This is because we mostly target sparsity in weights and for
* row-major layout it's more efficient to have A as a sparse matrix: for each
* non-zero of A at ith row and kth column, we can access kth row of B, whose
* elements are contiguous in memory. If B matrix was sparse, for each non-zero
* of B at kth row and jth column, we would've needed to access kth column of A,
* whose elements are not contiguous in memory with C/C++'s row-major layout.
* Alternatively, we can call this function as if we're computing
* C^T = B^T * A^T while maintaining PyTorch's convention that the lefthand
* side matrix B is activation. If B matrix is in column-major layout, we don't
* need to do an extra transposition. The C matrix will be output in
* column-major layout, so if we have a back-to-back Sparse-Dense matrix-matrix
* multiplications, B matrices of subsequent matrices will be already in
* column-major layout. Refer to SparseDenseMMFP32Benchmark.cc for an example.
*
*/
FBGEMM_API void SparseDenseMM(
int M,
int N,
const int* row_ptr,
const int* col_idx,
const float* values,
const float* B,
int ldb,
float* C,
int ldc,
bool accum = false);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
FBGEMM_API void fbgemmSparseDenseInt8MM(
int N,
const std::unique_ptr<BCSRMatrix<>>& bcsr,
const uint8_t* B,
int ldb,
int32_t* C_i32,
uint8_t* C_u8,
int ldc,
trRequantizationParams_t& rParams,
bool accum = false,
int thread_id = 0,
int num_threads = 1);
namespace internal {
void SparseDenseMMAvx2(
int M,
int N,
const int* row_ptr,
const int* col_idx,
const float* values,
const float* B,
int ldb,
float* C,
int ldc,
bool accum = false);
#if !defined(__aarch64__)
void SparseDenseMMAvx512(
int M,
int N,
const int* row_ptr,
const int* col_idx,
const float* values,
const float* B,
int ldb,
float* C,
int ldc,
bool accum = false);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
void SparseDenseInt8MMAvx2(
int N,
const std::unique_ptr<BCSRMatrix<>>& bcsr,
const uint8_t* B,
int ldb,
int32_t* C_i32,
uint8_t* C_u8,
int ldc,
trRequantizationParams_t& rParams,
bool accum = false,
int thread_id = 0,
int num_threads = 1);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
void SparseDenseInt8MMAvx512(
int N,
const std::unique_ptr<BCSRMatrix<>>& bcsr,
const uint8_t* B,
int ldb,
int32_t* C_i32,
uint8_t* C_u8,
int ldc,
trRequantizationParams_t& rParams,
bool accum = false,
int thread_id = 0,
int num_threads = 1);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
void SparseDenseInt8MVAvx512(
const std::unique_ptr<BCSRMatrix<>>& bcsr,
const uint8_t* B,
int ldb,
int32_t* C_i32,
uint8_t* C_u8,
trRequantizationParams_t& rParams,
bool accum = false,
int thread_id = 0,
int num_threads = 1);
#endif
} // namespace internal
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/FloatConversion.h
+331
View File
@@ -0,0 +1,331 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <math.h>
#include <cassert>
#include <climits>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include "./Types.h" // @manual
#ifndef __is_identifier
#define __is_identifier(x) 1
#endif
#define __has_keyword(__x) !(__is_identifier(__x))
// TODO: we're disabling native fp16 on Windows to workaround test failures
// due to "undefined symbol __gnu_h2f_ieee" error. We should follup on this
// later.
#if __has_keyword(__fp16) && !defined(_WIN32)
#define HAS_NATIVE_FP16_TYPE
using native_fp16_t = __fp16;
#elif __has_keyword(_Float16) && !defined(_WIN32)
#define HAS_NATIVE_FP16_TYPE
using native_fp16_t = _Float16;
#else
using native_fp16_t = void;
#endif
namespace fbgemm {
namespace detail {
template <typename T, int ExponentBits, bool HasInfinity = true>
struct FloatFormat {
using value_type = T;
static constexpr int bits = sizeof(T) * CHAR_BIT;
static constexpr int exponent_bits = ExponentBits;
static constexpr int mantissa_bits = bits - exponent_bits - 1;
static constexpr int sign_bit_pos = bits - 1;
static constexpr int exponent_bias = (1 << (exponent_bits - 1)) - 1;
static constexpr int unbiased_exponent_min = -exponent_bias + 1;
static constexpr int unbiased_exponent_max =
HasInfinity ? exponent_bias : (exponent_bias + 1);
static constexpr T sign_bit = T{1} << sign_bit_pos;
static constexpr T exponent_mask = ((T{1} << exponent_bits) - 1)
<< mantissa_bits;
static constexpr T mantissa_mask = (T{1} << mantissa_bits) - 1;
// signaling/quiet encoding is unspecified by IEEE754. This mirrors x86/ARM.
static constexpr T quiet_nan_bit = T{1} << (mantissa_bits - 1);
static constexpr T nan = exponent_mask | mantissa_mask;
static constexpr T overflow_value = HasInfinity ? exponent_mask : nan;
static constexpr bool has_infinity = HasInfinity;
static constexpr bool has_nan_payload = HasInfinity;
};
using IEEE754Single = FloatFormat</*T=*/uint32_t, /*ExponentBits=*/8>;
using IEEE754Half = FloatFormat</*T=*/uint16_t, /*ExponentBits=*/5>;
// See https://arxiv.org/abs/1905.12322v3
using BFloat16 = FloatFormat</*T=*/uint16_t, /*ExponentBits=*/8>;
// See https://doi.org/10.48550/arXiv.2209.05433
using FP8_E5M2 = FloatFormat</*T=*/uint8_t, /*ExponentBits=*/5>;
// See https://doi.org/10.48550/arXiv.2209.05433
using FP8_E4M3FN = FloatFormat<
/*T=*/uint8_t,
/*ExponentBits=*/4,
/*HasInfinity=*/false>;
enum class RoundingMode {
ToNearestTiesToEven,
ToZero,
};
// Generic IEEE754 truncation algorithm.
template <typename Src, typename Tgt, RoundingMode RoundingMode>
[[gnu::always_inline]] inline typename Tgt::value_type ieee754_trunc(
typename Src::value_type value) {
static_assert(Src::exponent_bits >= Tgt::exponent_bits);
static_assert(Src::mantissa_bits > Tgt::mantissa_bits);
using ST = typename Src::value_type;
using TT = typename Tgt::value_type;
ST src_exponent = value & Src::exponent_mask;
ST src_mantissa = value & Src::mantissa_mask;
// Fast-path: If there is no difference in exponent sizes (e.g. fp32 -> bf16)
// and we round toward zero, then we can just drop the least significant bits.
if constexpr (
Src::exponent_bits == Tgt::exponent_bits && Src::has_infinity &&
Tgt::has_infinity && RoundingMode == RoundingMode::ToZero) {
TT result = value >> (Src::bits - Tgt::bits);
// Turn signaling NaN into quiet NaN. This also avoids that the mantissa
// is completely zero after truncation (which would be misinterpreted as
// INF).
if (src_exponent == Src::exponent_mask && src_mantissa != 0) {
result |= Tgt::quiet_nan_bit;
}
return result;
}
ST tgt_sign =
(value & Src::sign_bit) >> (Src::sign_bit_pos - Tgt::sign_bit_pos);
constexpr bool denormal_becomes_zero =
Tgt::unbiased_exponent_min - Src::unbiased_exponent_min >
Src::mantissa_bits - Tgt::mantissa_bits;
if constexpr (denormal_becomes_zero) {
// Fast-path for zero exponentbits: This means the number was zero or a
// denormal number that will turn into zero in the Tgt format.
if (src_exponent == 0) {
return tgt_sign; // tgt_exponent == 0, tgt_mantissa == 0
}
}
int unbiased_exponent =
(src_exponent >> Src::mantissa_bits) - Src::exponent_bias;
if (unbiased_exponent < Tgt::unbiased_exponent_min) {
int shift = Tgt::unbiased_exponent_min - unbiased_exponent;
if (shift <= Tgt::mantissa_bits + 1) {
// Result is denormal.
ST src_mantissa_one = src_mantissa;
// Add explicit one if the source was not denormal.
if (denormal_becomes_zero || src_exponent != 0) {
src_mantissa_one |= TT{1} << Src::mantissa_bits;
} else {
shift--;
}
TT tgt_mantissa =
src_mantissa_one >> (Src::mantissa_bits - Tgt::mantissa_bits + shift);
if constexpr (RoundingMode == RoundingMode::ToNearestTiesToEven) {
int half_pos = Src::mantissa_bits - Tgt::mantissa_bits + shift - 1;
ST half = 1 << half_pos;
ST remainder = src_mantissa_one & ((half << 1) - 1);
if (remainder > half ||
(remainder == half && (tgt_mantissa & 1) != 0)) {
tgt_mantissa += 1;
}
} else {
static_assert(RoundingMode == RoundingMode::ToZero);
}
return tgt_sign | tgt_mantissa; // tgt_exponent == 0
} else {
// Result is +/- zero
return tgt_sign; // tgt_exponent == 0, tgt_mantissa == 0
}
}
if (unbiased_exponent > Tgt::unbiased_exponent_max) {
if (unbiased_exponent == Src::exponent_bias + 1 && src_mantissa != 0) {
TT tgt_mantissa;
if constexpr (Tgt::has_nan_payload) {
// NaN; not a number
tgt_mantissa =
src_mantissa >> (Src::mantissa_bits - Tgt::mantissa_bits);
tgt_mantissa |= Tgt::quiet_nan_bit;
} else {
tgt_mantissa = Tgt::mantissa_mask;
}
return tgt_sign | Tgt::exponent_mask | tgt_mantissa;
} else {
if (RoundingMode == RoundingMode::ToZero &&
(!Src::has_infinity || src_exponent != Src::exponent_mask)) {
// Return largest finite number.
return tgt_sign | (Tgt::exponent_mask - Tgt::has_infinity) |
Tgt::mantissa_mask;
}
// Infinity or NaN for formats without infinity.
return tgt_sign | Tgt::overflow_value;
}
}
// Normal number.
TT tgt_mantissa = src_mantissa >> (Src::mantissa_bits - Tgt::mantissa_bits);
TT tgt_exponent = (unbiased_exponent + Tgt::exponent_bias)
<< Tgt::mantissa_bits;
if constexpr (RoundingMode == RoundingMode::ToNearestTiesToEven) {
ST half = 1 << (Src::mantissa_bits - Tgt::mantissa_bits - 1);
ST remainder = src_mantissa & ((half << 1) - 1);
if (remainder > half || (remainder == half && (tgt_mantissa & 1) != 0)) {
if (tgt_mantissa < Tgt::mantissa_mask) {
tgt_mantissa += 1;
} else {
// Mantissa overflowed, increment exponent.
// Normally we can just add to the exponent and will naturally end up
// on infinity on overflow. But we need special treatments for formats
// without infinity.
if (Tgt::has_infinity || tgt_exponent != Tgt::exponent_mask) {
tgt_mantissa = 0;
tgt_exponent += TT{1} << Tgt::mantissa_bits;
} else {
// Return NaN.
tgt_mantissa = Tgt::mantissa_mask;
}
}
}
} else {
static_assert(RoundingMode == RoundingMode::ToZero);
}
return tgt_sign | tgt_exponent | tgt_mantissa;
}
} // namespace detail
inline float16 cpu_float2half_rn(float f) {
uint32_t f_u32 = 0;
std::memcpy(&f_u32, &f, sizeof(f_u32));
return detail::ieee754_trunc<
/*Src=*/detail::IEEE754Single,
/*Tgt=*/detail::IEEE754Half,
detail::RoundingMode::ToNearestTiesToEven>(f_u32);
}
inline float16 cpu_float2half_rz(float f) {
uint32_t f_u32 = 0;
std::memcpy(&f_u32, &f, sizeof(f_u32));
return detail::ieee754_trunc<
/*Src=*/detail::IEEE754Single,
/*Tgt=*/detail::IEEE754Half,
detail::RoundingMode::ToZero>(f_u32);
}
// Converts a 16-bit unsigned integer representation of a IEEE754 half-precision
// float into an IEEE754 32-bit single-precision float
inline float cpu_half2float_ref(const float16 h) {
constexpr uint32_t f16_num_exponent_bits = 5;
constexpr uint32_t f16_num_mantissa_bits = 10;
constexpr uint32_t f16_num_non_sign_bits =
f16_num_exponent_bits + f16_num_mantissa_bits;
constexpr uint32_t f16_exponent_bias = 15;
constexpr uint32_t f16_exponent_mask = 0b1'1111;
constexpr uint32_t f16_mantissa_mask = 0b11'1111'1111;
constexpr uint32_t f32_num_exponent_bits = 8;
constexpr uint32_t f32_num_mantissa_bits = 23;
constexpr uint32_t f32_num_non_sign_bits =
f32_num_exponent_bits + f32_num_mantissa_bits;
constexpr uint32_t f32_exponent_bias = 127;
constexpr uint32_t f32_exponent_mask = 0b1111'1111;
constexpr uint32_t f32_mantissa_mask = 0x7F'FF'FF;
constexpr uint32_t f32_most_significant_bit = 1u << 22;
// Get sign and exponent alone by themselves
uint32_t sign_bit = (h >> f16_num_non_sign_bits) & 1;
uint32_t exponent = (h >> f16_num_mantissa_bits) & f16_exponent_mask;
// Shift mantissa so that it fills the most significant bits of a float32
uint32_t mantissa = (h & f16_mantissa_mask)
<< (f32_num_mantissa_bits - f16_num_mantissa_bits);
if (exponent == f16_exponent_mask) { // NaN or Inf
if (mantissa) {
mantissa = f32_mantissa_mask;
sign_bit = 0;
}
exponent = f32_exponent_mask;
} else if (!exponent) { // Denorm or Zero
if (mantissa) {
uint32_t msb = 0;
exponent = f32_exponent_bias - f16_exponent_bias + 1;
do {
msb = mantissa & f32_most_significant_bit;
mantissa <<= 1; // normalize
--exponent;
} while (!msb);
mantissa &= f32_mantissa_mask; // 1.mantissa is implicit
}
} else {
exponent += f32_exponent_bias - f16_exponent_bias;
}
const uint32_t i = (sign_bit << f32_num_non_sign_bits) |
(exponent << f32_num_mantissa_bits) | mantissa;
float ret = NAN;
std::memcpy(&ret, &i, sizeof(float));
return ret;
}
// Same as the previous function, but use the built-in fp16 to fp32
// conversion provided by the compiler
inline float cpu_half2float(const float16 h) {
#if defined(HAS_NATIVE_FP16_TYPE) && !defined(MISSING_GNU_F2H_IEEE)
__fp16 h_fp16 = NAN;
std::memcpy(&h_fp16, &h, sizeof(__fp16));
return h_fp16;
#else
return cpu_half2float_ref(h);
#endif
}
inline float16 cpu_float2half(const float f) {
#if defined(HAS_NATIVE_FP16_TYPE) && !defined(MISSING_GNU_F2H_IEEE)
__fp16 h = f;
float16 res = 0;
std::memcpy(&res, &h, sizeof(__fp16));
return res;
#else
return cpu_float2half_rn(f);
#endif
}
inline float cpu_bf162float(bfloat16 src) {
float ret = NAN;
uint32_t val_fp32 =
static_cast<uint32_t>(reinterpret_cast<const uint16_t*>(&src)[0]) << 16;
std::memcpy(&ret, &val_fp32, sizeof(float));
return ret;
}
inline bfloat16 cpu_float2bfloat16(float src) {
uint32_t temp = 0;
std::memcpy(&temp, &src, sizeof(uint32_t));
return (temp + (1u << 15)) >> 16;
}
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/OutputProcessing-inl.h
+320
View File
@@ -0,0 +1,320 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#include <math.h>
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
template <typename outT, typename inT, typename nextOPType>
template <inst_set_t instSet>
inline int memCopy<outT, inT, nextOPType>::f(
outT* out,
inT* inp,
const block_type_t& block,
int ld_out,
int ld_in) const {
static_assert(
std::is_same_v<outT, inT>,
"input and output data type must be of same type");
// only copy if destination is not the same as source
if (out + block.row_start * ld_out + block.col_start != inp) {
for (int i = block.row_start; i < block.row_start + block.row_size; ++i) {
memcpy(
out + block.col_start + i * ld_out,
inp + (i - block.row_start) * ld_in,
block.col_size * sizeof(inT));
}
}
return nextop_.template f<instSet>(out, out, block, ld_out, ld_out);
}
template <typename outT, typename inT, typename nextOPType>
template <inst_set_t instSet>
inline int DoSpmdmOnInpBuffer<outT, inT, nextOPType>::f(
outT* out,
inT* inp,
const block_type_t& block,
int ld_out,
int ld_in) const {
assert(B_csc_.NumOfCols() % groups_ == 0);
int n_per_group = B_csc_.NumOfCols() / groups_;
int g = block.col_start / n_per_group;
B_csc_.SpMDM(block, A_ + g * B_csc_.NumOfRows(), lda_, true, inp, ld_in);
return nextop_.template f<instSet>(out, inp, block, ld_out, ld_in);
}
template <typename outT, typename inT, typename nextOPType>
template <inst_set_t instSet>
inline int DoSConvOnInpBuffer<outT, inT, nextOPType>::f(
outT* out,
inT* inp,
const block_type_t& block,
int ld_out,
int ld_in) const {
B_csc_.SparseConv(conv_p_, block, A_, A_zero_point_, true, inp, ld_in);
return nextop_.template f<instSet>(out, inp, block, ld_out, ld_in);
}
template <
bool FUSE_RELU,
QuantizationGranularity Q_GRAN,
typename BIAS_TYPE,
typename outT,
typename inT,
typename nextOPType>
template <inst_set_t instSet>
inline int
ReQuantizeOutput<FUSE_RELU, Q_GRAN, BIAS_TYPE, outT, inT, nextOPType>::f(
outT* out,
const inT* inp,
const block_type_t& block,
int ld_out,
int ld_in) const {
static_assert(
std::is_same_v<inT, int32_t>, "input data type must be of int32_t type");
int ncol_per_group = ncols_ / groups_;
assert(
block.col_size <= ncol_per_group &&
"ReQuantizeOutput should be called at most 1 group at a time.");
if constexpr (
instSet == inst_set_t::anyarch || !std::is_same_v<outT, uint8_t>) {
for (int i = block.row_start; i < block.row_start + block.row_size; ++i) {
for (int j = block.col_start; j < block.col_start + block.col_size; ++j) {
inT raw = inp[(i - block.row_start) * ld_in + (j - block.col_start)];
if (Aq_zero_point_) {
raw -= Aq_zero_point_ * q_col_offsets_[j];
}
int Bq_zero_point_idx = 0;
if constexpr (Q_GRAN == QuantizationGranularity::TENSOR) {
Bq_zero_point_idx = 0;
} else if constexpr (Q_GRAN == QuantizationGranularity::GROUP) {
int g = block.col_start / ncol_per_group;
Bq_zero_point_idx = g;
} else {
static_assert(Q_GRAN == QuantizationGranularity::OUT_CHANNEL);
Bq_zero_point_idx = j;
}
if (q_row_offsets_) {
raw -= q_row_offsets_[i - block.row_start] *
Bq_zero_point_[Bq_zero_point_idx];
}
float raw_f = NAN;
if (bias_) {
if constexpr (std::is_same_v<BIAS_TYPE, float>) {
raw_f = raw;
raw_f += bias_[j] / act_times_w_scale_[Bq_zero_point_idx];
} else {
raw += bias_[j];
raw_f = raw;
}
} else {
raw_f = raw;
}
float ab = raw_f * C_multiplier_[Bq_zero_point_idx];
long rounded = std::lrintf(ab) + C_zero_point_;
out[i * ld_out + j] = std::max(
FUSE_RELU ? static_cast<long>(C_zero_point_) : 0l,
std::min(255l, rounded));
}
}
#if !defined(__aarch64__)
} else if constexpr (
instSet == inst_set_t::avx2 || instSet == inst_set_t::avx512) {
bool b_symmetric =
(Q_GRAN == QuantizationGranularity::TENSOR && Bq_zero_point_[0] == 0) ||
q_row_offsets_ == nullptr;
requantizationParams_t<BIAS_TYPE> r = {
Aq_zero_point_,
Bq_zero_point_,
C_zero_point_,
C_multiplier_,
q_row_offsets_,
q_col_offsets_,
bias_,
ncols_,
groups_,
act_times_w_scale_};
if (Aq_zero_point_ == 0) {
if (b_symmetric) {
if (bias_ == nullptr) {
requantizeOutputProcessingAvx2<true, true, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeOutputProcessingAvx2<true, true, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
} else {
if (bias_ == nullptr) {
requantizeOutputProcessingAvx2<true, false, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeOutputProcessingAvx2<true, false, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
}
} else {
if (b_symmetric) {
if (bias_ == nullptr) {
requantizeOutputProcessingAvx2<false, true, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeOutputProcessingAvx2<false, true, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
} else {
if (bias_ == nullptr) {
requantizeOutputProcessingAvx2<
false,
false,
Q_GRAN,
false,
FUSE_RELU>(out, inp, block, ld_out, ld_in, r);
} else {
requantizeOutputProcessingAvx2<false, false, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
}
}
#endif // __aarch64__
} else {
assert(0 && "Not supported yet");
}
return nextop_.template f<instSet>(out, out, block, ld_out, ld_out);
}
template <
bool FUSE_RELU,
QuantizationGranularity Q_GRAN,
typename outT,
typename inT,
typename nextOPType>
template <inst_set_t instSet>
inline int ReQuantizeForFloat<FUSE_RELU, Q_GRAN, outT, inT, nextOPType>::f(
outT* out,
inT* inp,
const block_type_t& block,
int ld_out,
int ld_in) const {
static_assert(
std::is_same_v<int32_t, inT>, "input data type is of not expected type");
static_assert(
std::is_same_v<float, outT>, "output data type is of not expected type");
int ncol_per_group = ncols_ / groups_;
assert(
block.col_size <= ncol_per_group &&
"ReQuantizeOutput should be called at most 1 group at a time.");
if constexpr (
instSet == inst_set_t::anyarch || !std::is_same_v<outT, float>) {
for (int i = block.row_start; i < block.row_start + block.row_size; ++i) {
for (int j = block.col_start; j < block.col_start + block.col_size; ++j) {
inT raw = inp[(i - block.row_start) * ld_in + j - block.col_start];
if (Aq_zero_point_) {
raw -= Aq_zero_point_ * q_col_offsets_[j];
}
int Bq_zero_point_idx = 0;
if constexpr (Q_GRAN == QuantizationGranularity::TENSOR) {
Bq_zero_point_idx = 0;
} else if constexpr (Q_GRAN == QuantizationGranularity::GROUP) {
int g = block.col_start / ncol_per_group;
Bq_zero_point_idx = g;
} else {
static_assert(Q_GRAN == QuantizationGranularity::OUT_CHANNEL);
Bq_zero_point_idx = j;
}
if (q_row_offsets_) {
raw -= q_row_offsets_[i - block.row_start] *
Bq_zero_point_[Bq_zero_point_idx];
}
float res = raw * Aq_scale_ * Bq_scale_[Bq_zero_point_idx];
if (bias_) {
res += bias_[j];
}
out[i * ld_out + j] = res;
if constexpr (FUSE_RELU) {
out[i * ld_out + j] = std::max<outT>(0.0f, out[i * ld_out + j]);
}
}
}
#if !defined(__aarch64__)
} else if constexpr (
instSet == inst_set_t::avx2 || instSet == inst_set_t::avx512) {
bool b_symmetric =
(Q_GRAN == QuantizationGranularity::TENSOR && Bq_zero_point_[0] == 0) ||
q_row_offsets_ == nullptr;
requantizationForFloatParams_t r = {
Aq_zero_point_,
Bq_zero_point_,
Aq_scale_,
Bq_scale_,
q_row_offsets_,
q_col_offsets_,
bias_,
ncols_,
groups_};
if (Aq_zero_point_ == 0) {
if (b_symmetric) {
if (bias_ == nullptr) {
requantizeForFloatAvx2<true, true, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeForFloatAvx2<true, true, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
} else {
if (bias_ == nullptr) {
requantizeForFloatAvx2<true, false, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeForFloatAvx2<true, false, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
}
} else {
if (b_symmetric) {
if (bias_ == nullptr) {
requantizeForFloatAvx2<false, true, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeForFloatAvx2<false, true, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
} else {
if (bias_ == nullptr) {
requantizeForFloatAvx2<false, false, Q_GRAN, false, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
} else {
requantizeForFloatAvx2<false, false, Q_GRAN, true, FUSE_RELU>(
out, inp, block, ld_out, ld_in, r);
}
}
}
#endif // __aarch64__
} else {
assert(0 && "Not supported yet");
}
return nextop_.template f<instSet>(out, out, block, ld_out, ld_out);
}
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/PackingTraits-inl.h
+541
View File
@@ -0,0 +1,541 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/*
* This file configures the important cache blocking parameters and registers
* blocking parameters for the matrix multiplication loops inside FBGEMM.
*
* ROW_INTERLEAVE: the number of interleaved rows to use vpmaddubsw instructions
* for packing B matrix. For 32-bit accumulation, ROW_INTERLEAVE = 4; For 16-bit
* accumulation, ROW_INTERLEAVE = 2.
*
* VLEN: the vector length of one SIMD register. For avx2, VLEN = 256; For
* avx512, VLEN = 512.
*
* NR: the register blocking parameters for N dimension. The total registers
* used in N dimension for C accumulations are NR * ROW_INTERLEAVE * 8 (int8) /
* VLEN.
*
* MR: the register blocking parameters for M dimension. The total number of
* registers used in M dimension for C accumulations is MR. This indicates the
* number of vpbroadcastw instructions for A.
*
* (MR) * (NR * ROW_INTERLEAVE * 8 (int8) / VLEN): the number of registers used
* for C accumulations. This number should be less than the maximum registers we
* can use for C accumulations (A max of 12 out of 16 ymm registers for avx2; a
* max of 28 out of 32 zmm registers for avx512 ). The remaining are used for A
* matrix loading, B matrix loading and as temp registers. C accumulation
* registers should be as large as possible to increase the register
* utilization.
*
* MCB: the cache blocking parameters for M dimension. MCB needs to be a
* multiple of MR.
*
* NCB: the cache blocking parameters for N dimension. NCB needs to be a
* multiple of NR.
*
* KCB: the cache blocking parameters for K dimension. KCB needs to be a
* multiple of ROW_INTERLEAVE.
*/
/**
* @brief Packing parameter specialization for accumulation into 32-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx2
*/
template <typename T>
struct PackingTraits<
T,
std::int32_t,
inst_set_t::avx2,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{12}; ///< Register block for M dimension.
static constexpr int NR_MIN{8}; ///< Minimum register block for N dimension.
///< 8 because 8*ROW_INTERLEAVE int8 elements
///< completely fill a 256-bit wide vector.
static constexpr int NR{8}; ///< Register block for N dimension.
///< NR = VLEN/8/ROW_INTERLEAVE = 256 / 8 / 4 = 8.
///< Total registers used for N dimension: NCB/NR.
///< Here we use 12 x 1 ymm register blocking for
///< the registers used for accumulation C.
static constexpr int ROW_INTERLEAVE{
4}; ///< 4 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
120}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
8}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{512}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for accumulation into 16-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx2.
*/
template <typename T>
struct PackingTraits<
T,
std::int16_t,
inst_set_t::avx2,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{3}; ///< Register block for M dimension.
static constexpr int NR_MIN{
16}; ///< Minimum register block for N dimension.
///< 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 256-bit wide vector.
static constexpr int NR{
16}; ///< Register block for N dimension;
///< NR = VLEN/8/ROW_INTERLEAVE = 256 / 8 / 2 = 16.
///< Total registers used for N dimension: NCB/NR.
///< Here we use 3 x 4 ymm register blocking for the
///< registers used for accumulation C.
static constexpr int ROW_INTERLEAVE{
2}; ///< 2 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
60}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
64}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{256}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for float input and float
* accumulation.
*
* This is picked when template paramtere T is of float type and instruction
* set is avx2.
*/
template <>
struct PackingTraits<float, float, inst_set_t::avx2> {
static constexpr int MR{3}; ///< Register block for M dimension
static constexpr int NR{32}; ///< Register block for N dimension
static constexpr int ROW_INTERLEAVE{1}; ///< No Row interleave.
static constexpr int MCB{
24}; ///< Cache block for M dimension (multiple of MR)
static constexpr int NCB{
64}; ///< Cache block for N dimension (multiple of NR)
static constexpr int KCB{256}; ///< Cache block for K dimension
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for fp16 input and float
* accumulation.
*
* This is picked when template parameter T is of float16 type and instruction
* set is avx2
*/
template <>
struct PackingTraits<float16, float, inst_set_t::avx2> {
static constexpr int BCOL{8};
static constexpr int ROW_INTERLEAVE{1};
};
/**
* @brief Packing parameter specialization for accumulation into 32-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512.
*/
template <typename T>
struct PackingTraits<
T,
std::int32_t,
inst_set_t::avx512,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{14}; ///< Register block for M dimension.
static constexpr int NR_MIN{
16}; ///< Minimum register block for N dimension.
///< 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector.
static constexpr int NR{
32}; ///< Register block for N dimension.
///< Must be a multiple of 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector. Total registers used for
///< N dimension: NR*ROW_INTERLEAVE*8/VLEN. We use MR x
///< NR*ROW_INTERLEAVE*8/VLEN zmm registers
///< for C accumulations.
static constexpr int ROW_INTERLEAVE{
4}; ///< 4 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
56}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
32}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{256}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for accumulation into 32-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512_ymm.
*/
template <typename T>
struct PackingTraits<
T,
std::int32_t,
inst_set_t::avx512_ymm,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{7}; ///< Register block for M dimension.
static constexpr int NR_MIN{16}; ///< Minimum register block for N dimension.
///< 8 because 8*ROW_INTERLEAVE int8 elements
///< completely fill a 256-bit wide vector.
static constexpr int NR{
32}; ///< Register block for N dimension.
///< NR = VLEN/8/ROW_INTERLEAVE = 256 / 8 / 4 = 8.
///< Total registers used for N dimension: NCB/NR.
///< Here we use 12 x 1 ymm register blocking for
///< the registers used for accumulation C.
static constexpr int ROW_INTERLEAVE{
4}; ///< 4 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
56}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
32}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{256}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for accumulation into 16-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512.
*/
template <typename T>
struct PackingTraits<
T,
std::int16_t,
inst_set_t::avx512,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{6}; ///< Register block for M dimension
static constexpr int NR_MIN{
32}; ///< Minimum register block for N dimension;
///< 32 because 32*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector.
static constexpr int NR{
128}; ///< Register block for N dimension;
///< Must be a multiple of 32 because 32*ROW_INTERLEAVE int8
///< elements completely fill a 512-bit wide vector. Total registers
///< used for N dimension: NR*ROW_INTERLEAVE*8/VLEN. We use MR x
///< NR*ROW_INTERLEAVE*8/VLEN zmm registers
///< for C accumulations.
static constexpr int ROW_INTERLEAVE{
2}; ///< 2 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
60}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
128}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{256}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for accumulation into 16-bit
* integers.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512_ymm.
*/
template <typename T>
struct PackingTraits<
T,
std::int16_t,
inst_set_t::avx512_ymm,
std::enable_if_t<is_8bit<T>::value>> {
static constexpr int MR{6}; ///< Register block for M dimension.
static constexpr int NR_MIN{
16}; ///< Minimum register block for N dimension.
///< 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 256-bit wide vector.
static constexpr int NR{
16}; ///< Register block for N dimension;
///< NR = VLEN/8/ROW_INTERLEAVE = 256 / 8 / 2 = 16.
///< Total registers used for N dimension: NCB/NR.
///< Here we use 3 x 4 ymm register blocking for the
///< registers used for accumulation C.
static constexpr int ROW_INTERLEAVE{
2}; ///< 2 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
60}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
64}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{256}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Helper struct to type specialize for int16_t and int32_t together.
*/
template <typename T>
struct is_16or32bit {
static constexpr bool value =
std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>;
};
/**
* @brief Packing parameter specialization for accumulation into 32-bit/16-bit
* integers.
*
* Since there is no int16_t accumulation for AVX512 VNNI, we redirect int16_t
* to int32_t accumulation and use the same blocking parameters as int32_t.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512_vnni.
*/
template <typename T, typename accT>
struct PackingTraits<
T,
accT,
inst_set_t::avx512_vnni,
std::enable_if_t<is_8bit<T>::value && is_16or32bit<accT>::value>> {
static constexpr int MR{8}; ///< Register block for M dimension.
static constexpr int NR_MIN{
16}; ///< Minimum register block for N dimension.
///< 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector.
static constexpr int NR{
48}; ///< Register block for N dimension.
///< Must be a multiple of 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector. Total registers used for
///< N dimension: NR*ROW_INTERLEAVE*8/VLEN. We use MR x
///< NR*ROW_INTERLEAVE*8/VLEN zmm registers
///< for C accumulations.
static constexpr int ROW_INTERLEAVE{
4}; ///< 4 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
384}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
48}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{512}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for accumulation into 32-bit/16-bit
* integers.
*
* Since there is no int16_t accumulation for AVX512 VNNI, we redirect int16_t
* to int32_t accumulation and use the same blocking parameters as int32_t.
*
* This is picked when T is of int8 type (signed or unsigned) and instruction
* set is avx512_vnni_ymm.
*/
template <typename T, typename accT>
struct PackingTraits<
T,
accT,
inst_set_t::avx512_vnni_ymm,
std::enable_if_t<is_8bit<T>::value && is_16or32bit<accT>::value>> {
static constexpr int MR{4}; ///< Register block for M dimension.
static constexpr int NR_MIN{
16}; ///< Minimum register block for N dimension.
///< 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector.
static constexpr int NR{
48}; ///< Register block for N dimension.
///< Must be a multiple of 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector. Total registers used for
///< N dimension: NR*ROW_INTERLEAVE*8/VLEN. We use MR x
///< NR*ROW_INTERLEAVE*8/VLEN zmm registers
///< for C accumulations.
static constexpr int ROW_INTERLEAVE{
4}; ///< 4 rows are interleaved to use vpmaddubsw instruction for packing
///< B matrix.
static constexpr int MCB{
384}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
48}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{512}; ///< Cache block for K dimension.
static std::tuple<int, int, int> getCacheBlockParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(MR));
}
static std::tuple<int, int, int, int> getKernelParams() {
return std::tuple<int, int, int, int>(
int(MCB), int(NCB), int(NR_MIN), int(NR));
}
static std::tuple<int, int, int> getMatrixPackAParams() {
return std::tuple<int, int, int>(int(MCB), int(KCB), int(ROW_INTERLEAVE));
}
static std::tuple<int, int, int> getMatrixPackBParams() {
return std::tuple<int, int, int>(int(KCB), int(NCB), int(ROW_INTERLEAVE));
}
};
/**
* @brief Packing parameter specialization for I64 GEMM
* integers.
*
* This is picked when T is of int64 type and instruction
* set is avx512.
*/
template <>
struct PackingTraits<int64_t, int64_t, inst_set_t::avx512> {
static constexpr int MR{2}; ///< Register block for M dimension.
static constexpr int NR_MIN{8}; ///< Minimum register block for N dimension.
///< 8 because 8 int64 elements
///< completely fill a 512-bit wide vector.
static constexpr int NR{
32}; ///< Register block for N dimension.
///< Must be a multiple of 16 because 16*ROW_INTERLEAVE int8 elements
///< completely fill a 512-bit wide vector. Total registers used for
///< N dimension: NR*8/VLEN. We use MR x
///< NR*8/VLEN zmm registers
///< for C accumulations.
static constexpr int MCB{
16}; ///< Cache block for M dimension (multiple of MR).
static constexpr int NCB{
64}; ///< Cache block for N dimension (multiple of NR).
static constexpr int KCB{8}; ///< Cache block for K dimension.
};
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/QuantUtils.h
+397
View File
@@ -0,0 +1,397 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include "./FbgemmBuild.h" // @manual
#include "./QuantUtilsAvx2.h" // @manual
#include "./QuantUtilsAvx512.h" // @manual
#include "./QuantUtilsNeon.h" // @manual
#include "./Types.h" // @manual
#include "./Utils.h" // @manual
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <limits>
/// @defgroup fbgemm-quant-utils-generic Quantization Utilities (Generic)
///
namespace fbgemm {
FBGEMM_API TensorQuantizationParams ChooseQuantizationParams(
float min,
float max,
std::int32_t qmin,
std::int32_t qmax,
bool preserve_sparsity = false,
bool force_scale_power_of_two = false);
FBGEMM_API void ChooseRequantizationMultiplier(
float real_multiplier,
std::int32_t* quantized_multiplier,
int* right_shift,
int requantization_multiplier_precision = 32);
////////////////////////////////////////////////////////////////////////////////
// Utility functions
// Clamp src in T1 to the desired precision and convert it to T2
// TODO: T26263653 fix signed-integer-overflow undefined behavior
template <typename T1, typename T2 = std::uint8_t>
NO_SANITIZE("signed-integer-overflow")
T2 clamp(T1 src, int precision, bool is_signed = false) {
std::int32_t min = is_signed ? -(1LL << (precision - 1)) : 0;
std::int32_t max =
is_signed ? ((1LL << (precision - 1)) - 1) : (1LL << precision) - 1;
// Make sure T1 and T2 can represent the precision
assert(min >= std::numeric_limits<T1>::lowest());
assert(min >= std::numeric_limits<T2>::lowest());
assert(max <= std::numeric_limits<T1>::max());
assert(max <= std::numeric_limits<T2>::max());
return std::min<T1>(std::max<T1>(src, min), max);
}
/// Quantize src using zero_point and scale, clamp to the specified precision,
/// and convert it to type T
template <typename T, bool LEGACY = true>
T Quantize(
float src,
std::int32_t zero_point,
float scale,
int result_precision,
bool result_is_signed = std::is_signed_v<T>) {
// Note: We want to multiply with src with inv_scale instead of
// dividing src by scale. The same is done in vector code and
// at other places.
//
// Example:
// With scale = 0.00214854861f, zero_point = 0 and src = 0.273939937f
// transformed_val is 127.5 for src * inv_scale while
// transformed_val is 127.499992 for src / scale.
// Eventually 127.5 gets rounded to 128 while 127.499992 gets rounded to 127.
float inv_scale = 1.0f / scale;
float transformed_val = src * inv_scale;
// nearbyint here performs round-to-nearest-ties-to-even with
// default rounding mode.
// For example, nearbyint(1.4) is 1.0, nearbyint(1.5) is 2.0
// and nearbyint(2.5) is 2.0
// Adding zero_point before or after rounding can make a difference
// in exactly halfway cases.
if constexpr (LEGACY) {
transformed_val = std::nearbyint(zero_point + transformed_val);
} else {
transformed_val = zero_point + std::nearbyint(transformed_val);
}
// Please note the use of double. Unlike float, a double can represent
// all int32 values exactly. Using a float results in a float value >
// INT32_MAX conversion to int32 in clamp function and hence an UBSAN error.
return clamp<double, T>(transformed_val, result_precision, result_is_signed);
}
template <typename T, bool LEGACY = true>
T Quantize(float src, const TensorQuantizationParams& qparams) {
return Quantize<T, LEGACY>(
src, qparams.zero_point, qparams.scale, qparams.precision);
}
template <typename T, bool LEGACY = true>
FBGEMM_API void Quantize(
const float* src,
T* dst,
std::int64_t len,
const TensorQuantizationParams& qparams,
int thread_id = 0,
int num_threads = 1);
/// @ingroup fbgemm-quant-utils-generic
///
/// Quantize floating point data in `src` to type `T`.
///
/// @tparam T output quantized data type (`int8_t`, `uint8_t`, and `int32_t` are
/// supported)
///
/// @tparam LAYOUT layout of input tensor in `src`. (`KCX` and `KXC` are
/// supported)
/// `KCX` corresponds to `KCRS` or `KCTRS` (for weight tensors with time
/// dimension)
/// `KXC` corresponds to `KRSC` or `KTRSC` (for weight tensors with time
/// dimension)
///
/// @param K Output channels for weight tensors
/// @param C Number of channels
/// @param X `R*S` or `T*R*S`
/// @param G Groups (if `G == C` the function performs channelwise
/// quantization;
/// if `1 < G < C` the function performs groupwise
/// quantization; if `G == 1` the function performs per tensor
/// quantization;)
/// @param scales floating point scales. Size should be equal `G`
/// @param zero_points zero points (should be reprsentable in type `T`).
/// Size should be equal `G`
template <typename T, layout_t LAYOUT = layout_t::KCX>
FBGEMM_API void QuantizeGroupwise(
const float* src,
int K,
int C,
int X,
int G,
const float* scales,
const std::int32_t* zero_points,
T* dst);
template <typename T>
float Dequantize(T src, const TensorQuantizationParams& qparams) {
return qparams.scale * (src - qparams.zero_point);
}
template <typename T>
void Dequantize(
const T* src,
float* dst,
std::int64_t len,
const TensorQuantizationParams& qparams,
int thread_id = 0,
int num_threads = 1) {
int64_t i_begin = 0, i_end = 0;
fbgemmPartition1D(thread_id, num_threads, len, i_begin, i_end);
for (int64_t i = i_begin; i < i_end; i++) {
dst[i] = Dequantize(src[i], qparams);
}
}
template <typename T>
float FusedQuantizeDequantize(
float src,
const TensorQuantizationParams& qparams) {
T q = Quantize<T, false>(
src, qparams.zero_point, qparams.scale, qparams.precision);
return Dequantize<T>(q, qparams);
}
/// @ingroup fbgemm-quant-utils-generic
///
/// Fused integer quantization dequantization kernel to accelerate
/// quantization-aware training. Quantize `fp32` values in src to `(u)int8`
/// using the provided qparams, and dequantize quantized integer values back
/// into `fp32`.
template <typename T>
FBGEMM_API void FusedQuantizeDequantize(
const float* src,
float* dst,
std::int64_t len,
const TensorQuantizationParams& qparams,
int thread_id = 0,
int num_threads = 1,
float noise_ratio = 0.0f);
////////////////////////////////////////////////////////////////////////////////
// Requantization (pure fixed-point)
FBGEMM_API std::int64_t
SaturatingRoundingMulWithShift(std::int32_t a, std::int32_t b, int right_shift);
template <typename T>
T Requantize(
std::int32_t src, // int32 input before requantization
std::int32_t zero_point,
std::int32_t multiplier,
int right_shift,
int result_precision,
bool result_is_signed = false) {
std::int64_t quantized_down =
zero_point + SaturatingRoundingMulWithShift(src, multiplier, right_shift);
return clamp<std::int64_t, T>(
quantized_down, result_precision, result_is_signed);
}
template <typename T>
T RequantizeFixedPoint(
std::int32_t src, // int32 input before requantization
const RequantizationParams& params) {
return Requantize<T>(
src,
params.target_qparams.zero_point,
params.multiplier,
params.right_shift,
params.target_qparams.precision);
}
template <typename T>
FBGEMM_API void RequantizeFixedPoint(
const std::int32_t* src,
T* dst,
std::int64_t len,
const RequantizationParams& params,
int thread_id = 0,
int num_threads = 1);
////////////////////////////////////////////////////////////////////////////////
// Requantization (with floats)
template <typename T>
T Requantize(
std::int32_t src, // int32 input before requantization
std::int32_t zero_point,
float multiplier,
int result_precision,
bool result_is_signed = false) {
long quantized_down = zero_point + std::lrintf(src * multiplier);
return clamp<long, T>(quantized_down, result_precision, result_is_signed);
}
template <typename T>
T Requantize(
std::int32_t src, // int32 input before requantization
const RequantizationParams& params) {
return Requantize<T>(
src,
params.target_qparams.zero_point,
params.real_multiplier,
params.target_qparams.precision);
}
template <typename T>
FBGEMM_API void Requantize(
const std::int32_t* src,
T* dst,
std::int64_t len,
const RequantizationParams& params,
int thread_id = 0,
int num_threads = 1);
/**
* @ingroup fbgemm-quant-utils-generic
*
* Convert float (fp32 or fp16) inputs to rowwise quantized outputs.
* bitrate specifies the number of bits in quantized output.
* Scale and Bias are in fp16. Each row's Scale and Bias are stored in
* the row itself (fused) at the end.
*
* @param bit_rate can be 2, 4, or 8
*/
template <typename InputType>
FBGEMM_API void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
int bit_rate,
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output,
const InputType* rowwise_min_max = nullptr);
/**
* Convert fused rowwise quantized inputs to float (fp32 or fp16).
* bitrate specifies the number of bits in quantized input.
* Scale and Bias are in fp16. Each row's Scale and Bias are stored in
* the row itself (fused) at the end.
*
* @param bit_rate can be 2, 4, or 8
*/
template <typename OutputType>
FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalf(
int bit_rate,
const uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output,
bool scale_bias_last = true);
/**
* Convert float or half inputs to rowwise quantized (8-bit) outputs.
* Scale and Bias are in float. Each row's Scale and Bias are stored in
* the row itself (fused) at the end.
*
* This version intentionally supports only 8-bit because we want to discourage
* the usage of float scale and bias with 2 and 4 bit cases as that diminishes
* the overall memory savings.
*/
template <typename InputType>
FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloat(
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output,
const InputType* rowwise_min_max = nullptr);
/**
* Convert fused rowwise quantized (8-bit) inputs to float or half outputs.
* Scale and Bias are in float. Each row's Scale and Bias are stored in
* the row itself (fused) at the end.
*
* This version intentionally supports only 8-bit because
* the corresponding quantize version only supports 8-bit.
*/
template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf(
const uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output,
const bool scale_bias_last = true,
const bool quant_padding_float_type = true);
/**
* Same as ToFusedNBitRowwiseQuantizedSBHalf but unoptimized.
* This should not be called directly except in testing.
*/
template <typename InputType>
FBGEMM_API void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfRef(
int bit_rate,
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output);
/**
* Same as FloatOrHalfToFused8BitRowwiseQuantizedSBFloat but unoptimized.
* This should not be called directly except in testing.
*/
template <typename InputType>
FBGEMM_API void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatRef(
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output);
/**
* Same as FusedNBitRowwiseQuantizedSBHalfToFloat but unoptimized.
* This should not be called directly except in testing.
*/
template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
FBGEMM_API void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfRef(
int bit_rate,
const uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output,
bool scale_bias_last = true);
/**
* Same as Fused8BitRowwiseQuantizedSBFloatToFloatOrHalf but unoptimized.
* This should not be called directly except in testing.
*/
template <typename OutputType, bool is_uint16_t_of_type_bf16 = false>
FBGEMM_API void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfRef(
const uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output,
const bool scale_bias_last = true,
const bool quant_padding_float_type = true);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/QuantUtilsAvx2.h
+192
View File
@@ -0,0 +1,192 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include "./FbgemmBuild.h" // @manual
#include "./UtilsAvx2.h" // @manual
/// @defgroup fbgemm-quant-utils-avx2 Quantization Utilities (AVX2)
///
namespace fbgemm {
/// Number of columns in the rowwise min/max buffer passed to the quantization
/// function(s)
constexpr int kRowwiseMinMaxNumCols = 2;
/// Struct from <a href="https://github.com/google/gemmlowp">`gemmlowp`</a>
///
/// A structure to hold quantization parameters `scale` and `zero_point`.
/// The meaning of these values is as the constants in the quantization equation
///
/// `real_value = scale * (quantized_value - zero_point)`
///
/// In other words, 'zero_point' is the quantized value that corresponds
/// to the real value 0, and 'scale' is the difference of real values
/// corresponding to consecutive quantized values.
struct FBGEMM_API TensorQuantizationParams {
float scale;
std::int32_t zero_point;
int precision;
float Min() const;
float Max() const;
};
/// Parameters when we scale from int32 intermediate matrix multiplication
/// results to 8-bit integers
struct FBGEMM_API RequantizationParams {
/// For floating-point requantization
float real_multiplier;
/// For fixed-point requantization
std::int32_t multiplier;
int right_shift;
TensorQuantizationParams target_qparams;
};
/// @ingroup fbgemm-quant-utils-avx2
///
/// @brief Find the min and max value in a float matrix.
void FBGEMM_API FindMinMax(const float* m, float* min, float* max, int64_t len);
#if !defined(__aarch64__)
////////////////////////////////////////////////////////////////////////////////
// Utility functions
////////////////////////////////////////////////////////////////////////////////
template <typename T = std::uint8_t, bool LEGACY = true>
void QuantizeAvx2(
const float* src,
T* dst,
int64_t len,
const TensorQuantizationParams& qparams);
template <typename T = std::uint8_t>
void FusedQuantizeDequantizeAvx2(
const float* src,
float* dst,
int len,
const TensorQuantizationParams& qparams,
float noise_ratio = 0.0f);
/// @ingroup fbgemm-quant-utils-avx2
///
/// Random number generator in [0, 9] based on
/// <a href="https://www.jstatsoft.org/v08/i14/paper">this paper</a>.
uint32_t FBGEMM_API Xor128();
void RequantizeFixedPointAvx2(
const std::int32_t* src,
std::uint8_t* dst,
int len,
const RequantizationParams& params);
void RequantizeAvx2(
const std::int32_t* src,
std::uint8_t* dst,
int len,
const RequantizationParams& params);
#endif // !defined(__aarch64__)
/// @ingroup fbgemm-quant-utils-avx2
///
/// Requantize with avx2 and bias is fused.
template <
bool A_SYMMETRIC,
bool B_SYMMETRIC,
QuantizationGranularity Q_GRAN,
bool HAS_BIAS,
bool FUSE_RELU,
typename BIAS_TYPE = std::int32_t,
bool DIRECT = false>
FBGEMM_API void requantizeOutputProcessingAvx2(
std::uint8_t* out,
const std::int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const requantizationParams_t<BIAS_TYPE>& r);
template <
bool A_SYMMETRIC,
bool B_SYMMETRIC,
QuantizationGranularity Q_GRAN,
bool HAS_BIAS,
bool FUSE_RELU,
int C_PER_G,
typename BIAS_TYPE = std::int32_t>
FBGEMM_API void requantizeOutputProcessingGConvAvx2(
std::uint8_t* out,
const std::int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const requantizationParams_t<BIAS_TYPE>& r);
template <
bool A_SYMMETRIC,
bool B_SYMMETRIC,
QuantizationGranularity Q_GRAN,
bool HAS_BIAS,
bool FUSE_RELU>
FBGEMM_API void requantizeForFloatAvx2(
float* out,
const std::int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const requantizationForFloatParams_t& r);
#if !defined(__aarch64__)
template <typename InputType, int BIT_RATE>
void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2(
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output,
const InputType* rowwise_min_max = nullptr);
template <typename InputType>
void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatAvx2(
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output,
const InputType* rowwise_min_max = nullptr);
template <typename OutputType, int BIT_RATE>
void FusedNBitRowwiseQuantizedSBHalfToFloatOrHalfAvx2(
const std::uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output);
template <
typename OutputType,
bool scale_bias_last = true,
bool quant_padding_float_type = true>
void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfAvx2(
const std::uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output);
#endif // !defined(__aarch64__)
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/QuantUtilsAvx512.h
+55
View File
@@ -0,0 +1,55 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include "Types.h"
#if !defined(__aarch64__)
#include <cstdint>
#include "./FbgemmBuild.h" // @manual
#include "./UtilsAvx2.h" // @manual
/// @defgroup fbgemm-quant-utils-avx512 Quantization Utilities (AVX512)
///
namespace fbgemm {
/// @ingroup fbgemm-quant-utils-avx512
///
/// Requantize with AVX512.
template <
bool A_SYMMETRIC,
bool B_SYMMETRIC,
QuantizationGranularity Q_GRAN,
bool HAS_BIAS,
bool FUSE_RELU,
int C_PER_G,
typename BIAS_TYPE = std::int32_t>
FBGEMM_API void requantizeOutputProcessingGConvAvx512(
std::uint8_t* out,
const std::int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const requantizationParams_t<BIAS_TYPE>& r);
template <bool scale_bias_last = true, bool quant_padding_float_type = true>
void Fused8BitRowwiseQuantizedSBFloatToBfloat16Avx512(
const std::uint8_t* input,
size_t input_rows,
int input_columns,
bfloat16* output);
} // namespace fbgemm
#endif
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/QuantUtilsNeon.h
+53
View File
@@ -0,0 +1,53 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#ifdef __aarch64__
#include <cstdint>
#include "./FbgemmBuild.h" // @manual
/// @defgroup fbgemm-quant-utils-avx2 Quantization Utilities (AVX2)
///
namespace fbgemm {
////////////////////////////////////////////////////////////////////////////////
// Utility functions
////////////////////////////////////////////////////////////////////////////////
template <typename InputType>
void FloatOrHalfToFused8BitRowwiseQuantizedSBFloatNeon(
const InputType* input,
size_t input_rows,
int input_columns,
uint8_t* output);
template <typename OutputType>
void Fused8BitRowwiseQuantizedSBFloatToFloatOrHalfNeon(
const std::uint8_t* input,
size_t input_rows,
int input_columns,
OutputType* output);
template <typename InputType, int BIT_RATE>
void FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfNeon(
const InputType* input,
size_t input_rows,
int input_columns,
std::uint8_t* output);
} // namespace fbgemm
#endif // __aarch64__
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/SimdUtils.h
+118
View File
@@ -0,0 +1,118 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include "./Utils.h" // @manual
#include <asmjit/core.h> // @manual
#include <asmjit/x86.h> // @manual
namespace fbgemm {
#if ASMJIT_LIBRARY_VERSION >= ASMJIT_LIBRARY_MAKE_VERSION(1, 17, 0)
//! 128-bit XMM register (SSE+).
class Xmm : public asmjit::x86::Vec {
public:
using Vec::Vec;
using Vec::operator=;
Xmm(uint32_t regId) : Vec(asmjit::x86::Vec::make_xmm(regId)) {}
//! Casts this register to a register that has half the size (XMM).
ASMJIT_INLINE_NODEBUG Xmm half() const noexcept {
return Xmm(id());
}
};
//! 256-bit YMM register (AVX+).
class Ymm : public asmjit::x86::Vec {
public:
using Vec::Vec;
using Vec::operator=;
Ymm(uint32_t regId) : Vec(asmjit::x86::Vec::make_ymm(regId)) {}
//! Casts this register to a register that has half the size (XMM).
ASMJIT_INLINE_NODEBUG Xmm half() const noexcept {
return Xmm(id());
}
};
//! 512-bit ZMM register (AVX512+).
class Zmm : public asmjit::x86::Vec {
public:
using Vec::Vec;
using Vec::operator=;
Zmm(uint32_t regId) : Vec(asmjit::x86::Vec::make_zmm(regId)) {}
//! Casts this register to a register that has half the size (YMM).
ASMJIT_INLINE_NODEBUG Ymm half() const noexcept {
return Ymm(id());
}
};
#else
using Xmm = asmjit::x86::Xmm;
using Ymm = asmjit::x86::Ymm;
using Zmm = asmjit::x86::Zmm;
#endif
/**
* @brief Some commonly used variables for different instruction sets
*/
template <inst_set_t inst_set>
struct simd_info;
template <>
struct simd_info<inst_set_t::avx2> {
static constexpr int WIDTH_BITS = 256;
static constexpr int WIDTH_BYTES = 32;
static constexpr int WIDTH_32BIT_ELEMS = 8;
static constexpr int NUM_VEC_REGS = 16;
using vec_reg_t = Ymm;
};
template <>
struct simd_info<inst_set_t::sve> {
// Implementation is unrolled to match params used on avx2
static constexpr int WIDTH_BITS = 256;
static constexpr int WIDTH_BYTES = 32;
static constexpr int WIDTH_32BIT_ELEMS = 8;
static constexpr int NUM_VEC_REGS = 32;
};
template <>
struct simd_info<inst_set_t::avx512> {
static constexpr int WIDTH_BITS = 512;
static constexpr int WIDTH_BYTES = 64;
static constexpr int WIDTH_32BIT_ELEMS = 16;
static constexpr int NUM_VEC_REGS = 32;
using vec_reg_t = Zmm;
};
template <>
struct simd_info<inst_set_t::avx512_vnni>
: public simd_info<inst_set_t::avx512> {};
template <>
struct simd_info<inst_set_t::avx512_ymm> {
static constexpr int WIDTH_BITS = 256;
static constexpr int WIDTH_BYTES = 32;
static constexpr int WIDTH_32BIT_ELEMS = 8;
static constexpr int NUM_VEC_REGS = 32;
using vec_reg_t = Ymm;
};
template <>
struct simd_info<inst_set_t::avx512_vnni_ymm>
: public simd_info<inst_set_t::avx512_ymm> {};
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/Types.h
+31
View File
@@ -0,0 +1,31 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
namespace fbgemm {
using float16 = std::uint16_t;
using bfloat16 = std::uint16_t;
inline int64_t round_up(int64_t val, int64_t unit) {
return (val + unit - 1) / unit * unit;
}
inline int64_t div_up(int64_t val, int64_t unit) {
return (val + unit - 1) / unit;
}
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/Utils.h
+505
View File
@@ -0,0 +1,505 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include "./FbgemmBuild.h" // @manual
#include "./UtilsAvx2.h" // @manual
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <string>
#include <type_traits>
#ifndef HAVE_SVE
#if defined(__aarch64__) && __ARM_FEATURE_SVE && \
__has_include(<arm_neon_sve_bridge.h>)
#define HAVE_SVE 1
#else
#define HAVE_SVE 0
#endif
#endif
namespace fbgemm {
/**
* @brief Helper struct to type specialize for uint8 and int8 together.
*/
template <typename T>
struct is_8bit {
static constexpr bool value =
std::is_same_v<T, int8_t> || std::is_same_v<T, uint8_t>;
};
/**
* @brief Typed enum to specify matrix operations.
*/
enum class matrix_op_t { NoTranspose, Transpose };
/**
* @brief Typed enum for supported instruction sets.
*/
enum class inst_set_t {
anyarch,
avx2,
avx512,
avx512_ymm,
avx512_vnni,
avx512_vnni_ymm,
sve
};
/**
* @brief Typed enum for optimized paths for convolutions
*/
enum class optimized_conv_t {
depthwise,
groupwise,
pointwise,
fastpath1d,
im2col,
directconv
};
/**
* @brief Typed enum for implementation type.
*
* ref is reference and opt is optimized.
*/
enum class impl_type_t { ref, opt };
/**
* @brief Typed enum to specify data layout.
* KCX can be KCRS format or KCTRS format (e.g., for 3-D convolutions)
* KXC can be KRSC format or KTRSC format (e.g., for 3-D convolutions)
*/
enum class FBGEMM_ENUM_CLASS_API layout_t { KCX, KXC };
/**
* @brief A function to compare data in two buffers for closeness/equality.
*/
template <typename T>
FBGEMM_API int compare_buffers(
const T* ref,
const T* test,
int m,
int n,
int ld,
size_t max_mismatches_to_report,
float atol = 1e-3);
/**
* @brief Print the matrix.
* @param op Transpose type of the matrix.
* @param R The height of the matrix.
* @param C The width of the matrix.
* @param ld The leading dimension of the matrix.
* @param name The prefix string before printing the matrix.
*/
template <typename T>
void printMatrix(
matrix_op_t op,
const T* inp,
size_t R,
size_t C,
size_t ld,
const std::string& name) {
// R: number of rows in op(inp)
// C: number of cols in op(inp)
// ld: leading dimension in inp
std::cout << name << ":" << "[" << R << ", " << C << "]" << '\n';
bool tr = (op == matrix_op_t::Transpose);
for (size_t r = 0; r < R; ++r) {
for (size_t c = 0; c < C; ++c) {
T res = tr ? inp[c * ld + r] : inp[r * ld + c];
if constexpr (std::is_integral_v<T>) {
std::cout << std::setw(5) << static_cast<int64_t>(res) << " ";
} else {
std::cout << std::setw(5) << res << " ";
}
}
std::cout << '\n';
}
}
/**
* @brief Transpose a matrix.
*
* @param M the number of rows of input matrix
* @param N the number of columns of input matrix
*/
template <typename T>
FBGEMM_API void transpose_simd(
int64_t M,
int64_t N,
const T* src,
int64_t ld_src,
T* dst,
int64_t ld_dst);
/**
* @brief Explicitly set instruction set to be used
*/
FBGEMM_API void fbgemmForceIsa(inst_set_t /*isa*/);
/**
* @brief Enable AVX512-256 path for Intel(r) Xeon(r) D servers
*/
FBGEMM_API void fbgemmEnableAvx512Ymm(bool /*flag*/);
/**
* @brief Are we running on a Xeon-D cpu?
*/
FBGEMM_API bool fbgemmIsIntelXeonD();
/**
* @brief Are we running on a AVX512 supported cpu?
*/
FBGEMM_API bool fbgemmHasAvx512Support();
/**
* @brief Are we running on a AVX2 supported cpu?
*/
FBGEMM_API bool fbgemmHasAvx2Support();
/**
* @brief Are we running on a AVX512_VNNI supported cpu?
*/
FBGEMM_API bool fbgemmHasAvx512VnniSupport();
/**
* @brief Are we running on a AVX512_BF16 supported cpu?
*/
FBGEMM_API bool fbgemmHasAvx512Bf16Support();
/**
* @brief Are we running on a ARM Neon supported cpu?
*/
FBGEMM_API bool fbgemmHasArmNeonSupport();
/**
* @brief Are we running on a ARM SVE supported cpu?
*/
FBGEMM_API bool fbgemmHasArmSveSupport();
/**
* @brief Are we running on a ARM SVE2 supported cpu?
*/
FBGEMM_API bool fbgemmHasArmSve2Support();
/**
* @brief Retrieve current CPU instruction set
*/
FBGEMM_API inst_set_t fbgemmInstructionSet();
/**
* @brief Is ISA is wide vector ZMM
*/
FBGEMM_API bool isZmm(inst_set_t /*isa*/);
/**
* @brief Is ISA is wide vector ZMM
*/
FBGEMM_API bool isYmm(inst_set_t /*isa*/);
/**
* @brief Helper struct to enable autotuning of FBGEMM packing and kernels.
*
* This structure is optional. If not used, the default values for these
* parameters are picked up from PackingTraits-inl.h. Please see this
* file for details on these parameters.
*/
struct FBGEMM_API BlockingFactors {
int MR;
int NR;
int NR_MIN;
int ROW_INTERLEAVE;
int MCB;
int KCB;
int NCB;
};
/**
* @brief A struct to represent the partition information for the threads on the
* m and n dimensions.
*/
struct FBGEMM_API thread_type_t {
int g_num_threads;
int m_num_threads;
int n_num_threads;
int g_thread_id;
int m_thread_id;
int n_thread_id;
std::string toString() const {
std::string out;
out += "g num threads: " + std::to_string(g_num_threads) + ", ";
out += "m num threads: " + std::to_string(m_num_threads) + ", ";
out += "n num threads: " + std::to_string(n_num_threads) + ", ";
out += "g thread id: " + std::to_string(g_thread_id) + ", ";
out += "m thread id: " + std::to_string(m_thread_id) + ", ";
out += "n thread id: " + std::to_string(n_thread_id);
return out;
}
};
/**
* @brief A heuristic algorithm to partition the threads across m and n
* dimensions for parallelization, ensuring the ratio between the number of rows
* allocated to each thread in the m dimension and the number of columns
* allocated to each thread in the n dimension is approximately aspect_ratio.
*
* The less aspect_ratio is, the more favorable it is to parallelize the m
* dimension over the n dimension.
*/
FBGEMM_API int fbgemmGet2DPartition(
int m,
int n,
int nthreads,
int n_align,
double aspect_ratio);
/**
* @brief A heuristic way to partition the threads across g, m and n dimensions
* for parallelization.
*/
FBGEMM_API thread_type_t fbgemmGetThreadPartition(
int g,
int m,
int n,
int thread_id,
int num_threads,
int n_align = 64);
template <int SIZE, typename T = std::int32_t>
std::string arrayToString(const std::array<T, SIZE>& inp) {
std::string out = "[";
for (int i = 0; i < SIZE; ++i) {
out += std::to_string(inp[i]);
out += (i != SIZE - 1) ? std::string(", ") : std::string("]");
}
return out;
}
template <typename accT = std::int32_t>
bool isValidBlockingFactor(const BlockingFactors* const param) {
constexpr bool is_32bit = std::is_same_v<accT, int32_t>;
constexpr bool is_16bit = std::is_same_v<accT, int16_t>;
static const auto iset = fbgemmInstructionSet();
if constexpr (is_32bit) {
if (param->ROW_INTERLEAVE != 4)
return false;
if (isZmm(iset)) {
if (param->NR_MIN != 16 || param->NR % param->NR_MIN)
return false;
} else if (isYmm(iset)) {
if (param->NR_MIN != 8 || param->NR % param->NR_MIN)
return false;
}
} else if constexpr (is_16bit) {
if (param->ROW_INTERLEAVE != 2)
return false;
if (isZmm(iset)) {
if (param->NR_MIN != 32 || param->NR % param->NR_MIN)
return false;
} else if (isYmm(iset)) {
if (param->NR_MIN != 16 || param->NR % param->NR_MIN)
return false;
}
}
if (param->MCB % param->MR)
return false;
if (param->NCB % param->NR)
return false;
if (isZmm(iset)) {
if constexpr (is_32bit) {
// Zmm register usage for C
if (param->MR * (param->NR / param->NR_MIN) > 28)
return false;
} else if constexpr (is_16bit) {
// Zmm register usage for C + one row for loading B
if ((param->MR * (param->NR / param->NR_MIN) +
(param->NR / param->NR_MIN)) > 28)
return false;
}
} else if (isYmm(iset)) {
if (param->MR * (param->NR / param->NR_MIN) > 12)
return false;
}
return true;
}
/**
* @brief Partition work across given number of threads
*
* @param start Given thread_id should execute starting from the index
* start
* @param stop Given thread_id should stop executing at the index stop
*
* i.e., the loop should be equivalent to for(int i = start; i < end; ++i)
*/
FBGEMM_API void fbgemmPartition1D(
int thread_id,
int num_threads,
std::int64_t total_work,
std::int64_t& start,
std::int64_t& end);
/**
* @brief Partition work across given number of threads in blocks
* of size block_size. Each thread gets a multiple of block_size
* work or nothing, except the last one. The last one might
* receive the fringe case.
*
* @param start Given thread_id should execute starting from the index
* start
* @param stop Given thread_id should stop executing at the index stop
*
* The loop can be equivalent to for(int i = start; i < end; i+=block_size)
* except for the last thread. (i.e., thread_id = num_threads - 1)
*
* Example 1: block_size = 2, num_threads = 2
* total_work start(th 0) end(th 0) start(th 1) end(th 1)
* 4 0 2 2 4
* 5 0 2 2 5
*
* Example 2: block_size = 2, num_threads = 3
* total_work start(th 0) end(th 0) start(th 1) end(th 1)
* 4 0 2 2 4
* 5 0 2 2 4
*
* total_work start(th 2) end(th 2)
* 4 4 4
* 5 4 5
*
* Example 3: block_size = 2, num_threads = 4
* total_work start(th 0) end(th 0) start(th 1) end(th 1)
* 4 0 2 2 4
* 5 0 2 2 4
*
* total_work start(th 2) end(th 2) start(th 3) end(th 3)
* 4 4 4 4 4
* 5 4 4 4 5
*/
FBGEMM_API void fbgemmPartition1DBlocked(
int thread_id,
int num_threads,
std::int64_t total_work,
int block_size,
std::int64_t& start,
std::int64_t& end);
/**
* @brief A stable sorting algorithm. It sorts 8 bits at a time, hence in a
* worst-case performing sizeof(K) / 8 passes. Providing meaningful max_value
* may help reduce the number of passes performed by radix_sort. If
* maybe_with_neg_vals is set to true, we are performing all possible passes,
* up to a sign bit. If OpenMP is available in a build system, radix_sort works
* in parallel.
*/
template <typename K, typename V>
FBGEMM_API std::pair<K*, V*> radix_sort_parallel(
K* const inp_key_buf,
V* const inp_value_buf,
K* const tmp_key_buf,
V* const tmp_value_buf,
const int64_t elements_count,
const int64_t max_value,
const bool maybe_with_neg_vals = false);
/**
* @brief Helper function that allows us to check whether radix_sort is
* accelerated with OpenMP or not.
*/
FBGEMM_API bool is_radix_sort_accelerated_with_openmp();
/**
* Choosing which kernel (autovec/asmjit/ref) to use for nbit-CPU-TBE
* Available kernels:
* * ref: non-optimized, reference implementation that focuses on
* correctness, not performance
* * asmjit: hand-optimized kernel by having asmjit emit SIMD
* instructions during runtime. Only supports x86_64 CPUs with
* AVX2/AVX512 instruction sets
* * autovec: the kernel written in regular C++ code but in a
* way that makes compilers easier to generate vectorized SIMD
* instructions out of it. Supports both x86_64 and aarch64 CPUs.
* Currently only available on Linux.
* How to set environment variables:
* * No environment variables: on x86_64 we will default to asmjit
* kernel, and on aarch64 and linux we will default to autovec.
* On non-linux aarch64 we will fall back to ref.
* * Set FBGEMM_NO_AUTOVEC: on aarch64 linux we will use ref. On other
* platforms this will have no effect.
* * Set FBGEMM_NO_ASMJIT: on x86_64 we will use ref. On other
* platforms this will have no effect.
* * Set FBGEMM_NO_ASMJIT AND FBGEMM_FORCE_AUTOVEC: on x86_64 we will
* use autovec if these two variables are set at the same time.
* No effect on other platforms.
* * FBGEMM_FORCE_AUTOVEC will override FBGEMM_NO_AUTOVEC if they
* are set at the same time.
* * These variables are considered set as long as they exist regardless
* of content. That means assigning values like "1", "true", "y", "0",
* "false" or "no" has the same effect. The easiest way of setting a
* variable is to prepend `<VARIABLE>=1` before the benchmarking command.
*/
FBGEMM_API bool is_autovec_disabled();
FBGEMM_API bool is_autovec_forced();
FBGEMM_API bool is_asmjit_disabled();
FBGEMM_API bool is_stats_enabled();
/**
* @brief A function to check if the input parameter in the nbit CPU TBE kernel
* is valid.
*/
template <typename OutType>
void nbit_embedding_sanity_check(
// assertions are ignored in release mode, in which case these parameters
// will be unused
[[maybe_unused]] const int input_bit_rate,
[[maybe_unused]] const int output_bit_rate,
[[maybe_unused]] const bool no_bag) {
assert(
(input_bit_rate == 2 || input_bit_rate == 4) &&
"input_bit_rate must be 2 or 4");
// NOLINTNEXTLINE(bugprone-branch-clone)
if constexpr (std::is_same_v<OutType, uint8_t>) {
assert(
(no_bag && input_bit_rate == 4 && output_bit_rate == 4) &&
"we currently only support int4 to int4 for sequential TBE");
} else {
assert(
(output_bit_rate == 8 * sizeof(OutType)) &&
"output_bit_rate should be equal to 8 * sizeof(OutType)");
}
}
#define WARN_ONCE(...) \
do { \
static bool _warned = false; \
if (!_warned) { \
_warned = true; \
fprintf(stderr, __VA_ARGS__); \
} \
} while (0)
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/UtilsAvx2.h
+97
View File
@@ -0,0 +1,97 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
// This file defines common utilities used in code compiled with avx2/avx512
// flags.
#include <cstdint>
#include <string>
namespace fbgemm {
enum class FBGEMM_ENUM_CLASS_API QuantizationGranularity {
TENSOR,
GROUP,
OUT_CHANNEL,
};
/**
* @brief A struct to represent a block of a matrix.
*/
struct FBGEMM_API block_type_t {
int row_start;
int row_size;
int col_start;
int col_size;
std::string toString() const {
std::string out;
out += "row start:" + std::to_string(row_start) + ", ";
out += "row size:" + std::to_string(row_size) + ", ";
out += "col start:" + std::to_string(col_start) + ", ";
out += "col size:" + std::to_string(col_size);
return out;
}
};
/**
* @brief A struct to represent all the requantization parameters.
*
* Please note that this is different from RequantizationParams in
* QuantUtilsAvx2.h as it combines all the parameters needed for various
* quantization granularities
*/
template <typename BIAS_TYPE = std::int32_t>
struct requantizationParams_t {
using BIAS_T = BIAS_TYPE;
std::int32_t A_zero_point;
const std::int32_t* B_zero_point;
std::int32_t C_zero_point;
const float* C_multiplier;
const std::int32_t* row_offsets;
const std::int32_t* col_offsets;
const BIAS_T* bias;
std::uint32_t ncols;
int groups;
const float* act_times_w_scale;
};
/**
* @brief A struct to represent all the parameters for requantizing for floats.
*/
struct requantizationForFloatParams_t {
std::int32_t A_zero_point;
const std::int32_t* B_zero_point;
float A_scale;
const float* B_scale;
const std::int32_t* row_offsets;
const std::int32_t* col_offsets;
const float* bias;
std::uint32_t ncols;
int groups;
};
/**
* @brief Allocate size bytes of uninitialized storage whose alignment is
* specified by align.
*/
FBGEMM_API void*
fbgemmAlignedAlloc(size_t align, size_t size, bool raiseException = false);
/**
* @brief Free memory allocated by fbgemmAlignedAlloc
*/
FBGEMM_API void fbgemmAlignedFree(void* p);
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/spmmUtils.h
+62
View File
@@ -0,0 +1,62 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <vector>
#include "fbgemm/FbgemmBuild.h"
#include "fbgemm/FbgemmSparse.h"
#include "fbgemm/UtilsAvx2.h"
#include "fbgemm/spmmUtilsAvx2.h"
namespace fbgemm {
FBGEMM_API void sparseDenseMMRef(
int M,
int N,
const int* row_ptr,
const int* col_idx,
const float* values,
const float* B,
int ldb,
float* C,
int ldc,
bool accum = false);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
FBGEMM_API void sparseDenseInt8MMRef(
int N,
const std::unique_ptr<BCSRMatrix<>>& bcsr,
const uint8_t* B,
int ldb,
int32_t* C_i32,
uint8_t* C_u8,
int ldc,
trRequantizationParams_t& rParams,
bool accum = false,
int thread_id = 0,
int num_threads = 1);
template <bool FUSE_RELU, QuantizationGranularity Q_GRAN>
FBGEMM_API void trRequantizeRef(
uint8_t* out,
const int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const trRequantizationParams_t& r);
// Get matrix shapes of interest
FBGEMM_API std::vector<std::vector<int>> getSparseMatrixShapes();
} // namespace fbgemm
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/fbgemm/spmmUtilsAvx2.h
+44
View File
@@ -0,0 +1,44 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <cstdint>
#include "./FbgemmBuild.h" // @manual
#include "fbgemm/UtilsAvx2.h"
namespace fbgemm {
struct FBGEMM_API trRequantizationParams_t {
std::int32_t act_zero_point; // activation zero point
const std::int32_t* weight_zero_points; // weight zero point(s)
std::int32_t C_zero_point;
const float C_scale;
const std::int32_t* weight_row_offsets;
const std::int32_t* act_col_offsets;
const float* bias;
const float* act_times_w_scale;
};
template <
bool FUSE_RELU,
bool ACT_SYMMETRIC, // whether activation matrix is symmetric
bool WEIGHT_SYMMETRIC, // whether weight matrix is symmetric
bool HAS_BIAS,
QuantizationGranularity Q_GRAN>
FBGEMM_API void trRequantizeOpt(
uint8_t* out,
const int32_t* inp,
const block_type_t& block,
int ld_out,
int ld_in,
const trRequantizationParams_t& rParams);
} // namespace fbgemm
#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)