serge-sans-paille
72abc9ee48
Also reflect gemmology changes to our codebase. Basically using the default execution engine now that we need to state one explicitly. Differential Revision: https://phabricator.services.mozilla.com/D296101
303 lines
9.5 KiB
C++
303 lines
9.5 KiB
C++
/***************************************************************
|
|
* _ *
|
|
* | | *
|
|
* __ _ ___ _ __ ___ _ __ ___ ___ | | ___ __ _ _ _ *
|
|
* / _` |/ _ \ '_ ` _ \| '_ ` _ \ / _ \| |/ _ \ / _` | | | | *
|
|
* | (_| | __/ | | | | | | | | | | (_) | | (_) | (_| | |_| | *
|
|
* \__, |\___|_| |_| |_|_| |_| |_|\___/|_|\___/ \__, |\__, | *
|
|
* __/ | __/ | __/ | *
|
|
* |___/ |___/ |___/ *
|
|
* *
|
|
* version 0.1 *
|
|
***************************************************************/
|
|
|
|
#ifndef GEMMOLOGY_FWD_H
|
|
#define GEMMOLOGY_FWD_H
|
|
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <tuple>
|
|
#include <xsimd/xsimd.hpp>
|
|
|
|
#ifdef GEMMOLOGY_WITH_STD_THREAD
|
|
#include <thread>
|
|
#include <vector>
|
|
#endif
|
|
|
|
namespace gemmology {
|
|
|
|
struct SequentialExecutionEngine {
|
|
|
|
template<class F>
|
|
inline void operator()(size_t Start, size_t End, size_t Stride, F&& f) {
|
|
for(size_t i = Start; i < End; i += Stride) {
|
|
f(i);
|
|
}
|
|
}
|
|
|
|
};
|
|
|
|
|
|
#ifdef GEMMOLOGY_WITH_STD_THREAD
|
|
struct StdThreadExecutionEngine {
|
|
|
|
StdThreadExecutionEngine(size_t PoolSize) : MaxPoolSize(PoolSize) {
|
|
Pool.reserve(PoolSize - 1);
|
|
}
|
|
|
|
template<class F>
|
|
inline void operator()(size_t Start, size_t End, size_t Stride, F&& f) {
|
|
const size_t NbIter = (End - Start) / Stride;
|
|
const size_t NbThread = std::min(NbIter, MaxPoolSize);
|
|
const size_t Chunk = (NbIter / NbThread) * Stride;
|
|
|
|
size_t Curr = Start, Next = Start;
|
|
|
|
for(size_t threadID = 0; threadID < NbThread - 1; ++threadID) {
|
|
Next += Chunk;
|
|
Pool.emplace_back([=]() {
|
|
for(size_t i = Curr; i < Next; i += Stride) {
|
|
f(i);
|
|
};
|
|
});
|
|
Curr = Next;
|
|
}
|
|
|
|
for(size_t i = Next; i < End; i += Stride) {
|
|
f(i);
|
|
};
|
|
for(size_t threadID = 0; threadID < Pool.size(); ++threadID) {
|
|
Pool[threadID].join();
|
|
}
|
|
Pool.clear();
|
|
}
|
|
|
|
private:
|
|
const size_t MaxPoolSize;
|
|
std::vector<std::thread> Pool;
|
|
|
|
};
|
|
|
|
#endif
|
|
|
|
#ifdef _OPENMP
|
|
struct OpenMPExecutionEngine {
|
|
|
|
template<class F>
|
|
inline void operator()(size_t Start, size_t End, size_t Stride, F&& f) {
|
|
#pragma omp parallel for
|
|
for(size_t i = Start; i < End; i += Stride) {
|
|
f(i);
|
|
}
|
|
}
|
|
|
|
};
|
|
#endif
|
|
|
|
namespace callbacks {
|
|
|
|
struct Unquantize {
|
|
float unquant_mult;
|
|
template <class Arch>
|
|
xsimd::batch<float, Arch> operator()(xsimd::batch<int32_t, Arch> total, size_t, size_t, size_t);
|
|
template <class Arch>
|
|
std::tuple<xsimd::batch<float, Arch>, xsimd::batch<float, Arch>> operator()(
|
|
std::tuple<xsimd::batch<int32_t, Arch>, xsimd::batch<int32_t, Arch>>
|
|
total,
|
|
size_t, size_t, size_t);
|
|
};
|
|
|
|
struct AddBias {
|
|
const float *bias_addr;
|
|
template <class Arch>
|
|
xsimd::batch<float, Arch> operator()(xsimd::batch<float, Arch> total, size_t, size_t col_idx,
|
|
size_t);
|
|
template <class Arch>
|
|
std::tuple<xsimd::batch<float, Arch>, xsimd::batch<float, Arch>>
|
|
operator()(
|
|
std::tuple<xsimd::batch<float, Arch>, xsimd::batch<float, Arch>> total,
|
|
size_t, size_t col_idx, size_t);
|
|
};
|
|
|
|
struct Write {
|
|
float *output_addr;
|
|
|
|
Write(float *o) : output_addr(o) {}
|
|
Write(int32_t *o) : output_addr(reinterpret_cast<float*>(o)) {}
|
|
|
|
template <class Arch>
|
|
void operator()(xsimd::batch<float, Arch> result, size_t row_idx,
|
|
size_t col_idx, size_t col_size);
|
|
template <class Arch>
|
|
void operator()(xsimd::batch<int32_t, Arch> result, size_t row_idx,
|
|
size_t col_idx, size_t col_size);
|
|
|
|
template <class Arch>
|
|
void operator()(
|
|
std::tuple<xsimd::batch<float, Arch>, xsimd::batch<float, Arch>> result,
|
|
size_t row_idx, size_t col_idx, size_t col_size);
|
|
|
|
template <class Arch>
|
|
void operator()(
|
|
std::tuple<xsimd::batch<int32_t, Arch>, xsimd::batch<int32_t, Arch>>
|
|
result,
|
|
size_t row_idx, size_t col_idx, size_t col_size);
|
|
};
|
|
|
|
struct UnquantizeAndWrite {
|
|
|
|
Unquantize unquantize;
|
|
Write write;
|
|
|
|
UnquantizeAndWrite(float factor, float *output)
|
|
: unquantize{factor}, write{output} {}
|
|
|
|
template <class T>
|
|
void operator()(T const &total, size_t row_idx, size_t col_idx,
|
|
size_t col_size);
|
|
};
|
|
|
|
struct UnquantizeAndAddBiasAndWrite {
|
|
|
|
Unquantize unquantize;
|
|
AddBias add_bias;
|
|
Write write;
|
|
|
|
UnquantizeAndAddBiasAndWrite(float factor, const float *bias, float *output)
|
|
: unquantize{factor}, add_bias{bias}, write{output} {}
|
|
|
|
template <class T>
|
|
void operator()(T const &total, size_t row_idx, size_t col_idx,
|
|
size_t col_size);
|
|
};
|
|
|
|
} // namespace callbacks
|
|
|
|
//
|
|
// Arch-specific implementation of each routine
|
|
//
|
|
template <class Arch> struct Engine {
|
|
|
|
static void QuantizeU(const float *input, uint8_t *output, float quant_mult,
|
|
size_t size);
|
|
|
|
static void Quantize(const float *const input, int8_t *const output,
|
|
float quant_mult, size_t size);
|
|
|
|
template <typename IntegerTy>
|
|
static void SelectColumnsB(const int8_t *input, int8_t *output, size_t rows,
|
|
const IntegerTy *cols_begin,
|
|
const IntegerTy *cols_end);
|
|
|
|
static void PrepareBTransposed(const float *input, int8_t *output,
|
|
float quant_mult, size_t cols, size_t rows);
|
|
|
|
static void PrepareBQuantizedTransposed(const int8_t *input, int8_t *output,
|
|
size_t cols, size_t rows);
|
|
|
|
static void PrepareBQuantized(const int8_t *input, int8_t *output,
|
|
size_t cols, size_t rows);
|
|
|
|
static void PrepareB(const float *input, int8_t *output_shadow,
|
|
float quant_mult, size_t rows, size_t cols);
|
|
|
|
static void PrepareA(const float *input, int8_t *output, float quant_mult,
|
|
size_t rows, size_t cols);
|
|
|
|
struct Shift {
|
|
|
|
static void PrepareA(const float *input, uint8_t *output, float quant_mult,
|
|
size_t rows, size_t cols);
|
|
|
|
template <class Callback, class ExecutionEngine>
|
|
static void Multiply(const uint8_t *A, const int8_t *B, size_t A_rows,
|
|
size_t width, size_t B_cols, Callback callback,
|
|
ExecutionEngine& engine);
|
|
|
|
template <class Callback>
|
|
static void PrepareBias(const int8_t *B, size_t width, size_t B_cols,
|
|
Callback C);
|
|
};
|
|
};
|
|
|
|
//
|
|
// Top-level wrappers that mostly match intgemm API
|
|
//
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void QuantizeU(const float *input, uint8_t *output, float quant_mult,
|
|
size_t size) {
|
|
return Engine<Arch>::QuantizeU(input, output, quant_mult, size);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void Quantize(const float *const input, int8_t *const output,
|
|
float quant_mult, size_t size) {
|
|
return Engine<Arch>::Quantize(input, output, quant_mult, size);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch, typename IntegerTy>
|
|
inline void SelectColumnsB(const int8_t *input, int8_t *output, size_t rows,
|
|
const IntegerTy *cols_begin,
|
|
const IntegerTy *cols_end) {
|
|
return Engine<Arch>::SelectColumnsB(input, output, rows, cols_begin,
|
|
cols_end);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareBTransposed(const float *input, int8_t *output,
|
|
float quant_mult, size_t cols, size_t rows) {
|
|
return Engine<Arch>::PrepareBTransposed(input, output, quant_mult, cols,
|
|
rows);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareBQuantized(const int8_t *input, int8_t *output,
|
|
size_t cols, size_t rows) {
|
|
return Engine<Arch>::PrepareBQuantized(input, output, cols, rows);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareBQuantizedTransposed(const int8_t *input, int8_t *output,
|
|
size_t cols, size_t rows) {
|
|
return Engine<Arch>::PrepareBQuantizedTransposed(input, output, cols, rows);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareB(const float *input, int8_t *output_shadow,
|
|
float quant_mult, size_t rows, size_t cols) {
|
|
return Engine<Arch>::PrepareB(input, output_shadow, quant_mult, rows, cols);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareA(const float *input, int8_t *output, float quant_mult,
|
|
size_t rows, size_t cols) {
|
|
return Engine<Arch>::PrepareA(input, output, quant_mult, rows, cols);
|
|
}
|
|
|
|
namespace Shift {
|
|
|
|
template <class Arch = xsimd::default_arch>
|
|
inline void PrepareA(const float *input, uint8_t *output, float quant_mult,
|
|
size_t rows, size_t cols) {
|
|
return Engine<Arch>::Shift::PrepareA(input, output, quant_mult, rows, cols);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch, class Callback, class ExecutionEngine=SequentialExecutionEngine>
|
|
inline void Multiply(const uint8_t *A, const int8_t *B, size_t A_rows,
|
|
size_t width, size_t B_cols, Callback C, ExecutionEngine&& engine={}) {
|
|
return Engine<Arch>::Shift::Multiply(A, B, A_rows, width, B_cols, C, engine);
|
|
}
|
|
|
|
template <class Arch = xsimd::default_arch, class Callback>
|
|
inline void PrepareBias(const int8_t *B, size_t width, size_t B_cols,
|
|
Callback C) {
|
|
return Engine<Arch>::Shift::PrepareBias(B, width, B_cols, C);
|
|
}
|
|
|
|
} // namespace Shift
|
|
|
|
} // namespace gemmology
|
|
|
|
#endif
|