86 lines
2.9 KiB
C++
86 lines
2.9 KiB
C++
#include "int8_linear_avx2.h"
|
|
|
|
#include <algorithm>
|
|
#include <immintrin.h>
|
|
|
|
namespace neural
|
|
{
|
|
namespace detail
|
|
{
|
|
namespace
|
|
{
|
|
|
|
template <std::size_t BlockRows>
|
|
void linear_block(const float *x, const std::int8_t *weights, const float *scales, const float *bias, std::size_t rows,
|
|
std::size_t inputs, std::size_t outputs, float *out)
|
|
{
|
|
for (std::size_t output = 0; output < outputs; ++output)
|
|
{
|
|
const std::int8_t *weight_row = weights + output * inputs;
|
|
const float scale = scales[output];
|
|
const float offset = bias == nullptr ? 0.0f : bias[output];
|
|
for (std::size_t base = 0; base < rows; base += BlockRows)
|
|
{
|
|
const std::size_t count = std::min(BlockRows, rows - base);
|
|
__m256 accumulators[BlockRows];
|
|
for (std::size_t row = 0; row < BlockRows; ++row)
|
|
{
|
|
accumulators[row] = _mm256_setzero_ps();
|
|
}
|
|
std::size_t input = 0;
|
|
for (; input + 8 <= inputs; input += 8)
|
|
{
|
|
// Each block widens the weight chunk once. Its rows share the result in a register.
|
|
const __m128i packed = _mm_loadl_epi64(reinterpret_cast<const __m128i *>(weight_row + input));
|
|
const __m256 widened = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(packed));
|
|
for (std::size_t row = 0; row < count; ++row)
|
|
{
|
|
const __m256 activation = _mm256_loadu_ps(x + (base + row) * inputs + input);
|
|
accumulators[row] = _mm256_fmadd_ps(activation, widened, accumulators[row]);
|
|
}
|
|
}
|
|
for (std::size_t row = 0; row < count; ++row)
|
|
{
|
|
alignas(32) float lanes[8];
|
|
_mm256_store_ps(lanes, accumulators[row]);
|
|
float total = 0.0f;
|
|
for (float lane : lanes)
|
|
{
|
|
total += lane;
|
|
}
|
|
for (std::size_t tail = input; tail < inputs; ++tail)
|
|
{
|
|
total += x[(base + row) * inputs + tail] * static_cast<float>(weight_row[tail]);
|
|
}
|
|
out[(base + row) * outputs + output] = total * scale + offset;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
void linear_int8_avx2(const float *x, const std::int8_t *weights, const float *scales, const float *bias,
|
|
std::size_t rows, std::size_t inputs, std::size_t outputs, float *out)
|
|
{
|
|
if (rows == 1)
|
|
{
|
|
linear_block<1>(x, weights, scales, bias, rows, inputs, outputs, out);
|
|
}
|
|
else if (rows == 2)
|
|
{
|
|
linear_block<2>(x, weights, scales, bias, rows, inputs, outputs, out);
|
|
}
|
|
else if (rows <= 4)
|
|
{
|
|
linear_block<4>(x, weights, scales, bias, rows, inputs, outputs, out);
|
|
}
|
|
else
|
|
{
|
|
linear_block<8>(x, weights, scales, bias, rows, inputs, outputs, out);
|
|
}
|
|
}
|
|
|
|
} // namespace detail
|
|
} // namespace neural
|