File size: 2,972 Bytes
97c3d9a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | #include <stdint.h>
#include <stddef.h>
#include <immintrin.h>
#ifdef _OPENMP
#include <omp.h>
#endif
/**
* Native C inference kernel for the PRIME discrete neural network architecture.
*
* Instead of taking an [out_features, in_features] array of 32-bit floats (which bottleneck memory bandwidth),
* this takes an array of 16-bit unsigned integers representing LUT indices.
* It expands the continuous float values on-the-fly directly inside the CPU's L1 cache during the FMA loop.
*
* @param x Input activations matrix [batch_size * in_features]
* @param weight_indices PRIME compressed weights [out_features * in_features]
* @param lut The 65536-element Prime Harmonic Look-Up Table
* @param y Output activations matrix [batch_size * out_features]
* @param batch_size Number of tokens/sequences
* @param in_features Input dimension
* @param out_features Output dimension
*/
void prime_linear_forward(
const float* restrict x,
const uint16_t* restrict weight_indices,
const float* restrict lut,
float* restrict y,
int batch_size,
int in_features,
int out_features
) {
#pragma omp parallel for collapse(2)
for (int b = 0; b < batch_size; ++b) {
for (int o = 0; o < out_features; ++o) {
const uint16_t* w_row = &weight_indices[o * in_features];
const float* x_row = &x[b * in_features];
int i = 0;
float sum = 0.0f;
#if defined(__AVX512F__) && defined(__AVX512BW__)
__m512 sum_vec = _mm512_setzero_ps();
// Unroll by 16 (AVX-512 processes 16 floats at once)
for (; i <= in_features - 16; i += 16) {
// Load 16 uint16_t weights (256 bits)
__m256i w_16bit = _mm256_loadu_si256((const __m256i*)&w_row[i]);
// Zero-extend 16 uint16_t -> 16 uint32_t (512 bits)
__m512i w_32bit = _mm512_cvtepu16_epi32(w_16bit);
// Hardware GATHER: fetch 16 floats from LUT simultaneously
__m512 lut_vals = _mm512_i32gather_ps(w_32bit, lut, 4); // scale by 4 bytes (sizeof float)
// Load 16 input floats
__m512 x_vals = _mm512_loadu_ps(&x_row[i]);
// Fused Multiply-Add
sum_vec = _mm512_fmadd_ps(lut_vals, x_vals, sum_vec);
}
// Horizontal sum of the 16 elements in the 512-bit register
sum += _mm512_reduce_add_ps(sum_vec);
#endif
// Handle tail elements (or fallback if AVX-512 is not compiled)
for (; i < in_features; ++i) {
uint16_t idx = w_row[i];
float w_val = lut[idx];
sum += x_row[i] * w_val;
}
y[b * out_features + o] = sum;
}
}
}
|