/*
 * Fused scalar×slice upsampling kernel for ctypes.
 *
 * v1: basic auto-vectorised loop, strided scatter output
 * v2: explicit AVX2 FMA, shared input reads, vectorised interleave
 * v3: v2 + identity shortcut (sub-kernel 0 = delta → memcpy)
 *
 * gcc -O3 -march=native -ffast-math -mavx2 -mfma -shared -fPIC \
 *     -o upsample_c.so upsample_c.c
 */

#include <immintrin.h>
#include <stdlib.h>
#include <string.h>

/* ── v1: auto-vectorised baseline ──────────────────────────────────── */

void upsample_fast_c(float *restrict output,
                     const float *restrict input,
                     const float *restrict kernel,
                     int B, int T, int n_win, int klen, int factor) {
    float *acc = (float *)malloc(n_win * sizeof(float));

    for (int f = 0; f < factor; f++) {
        const float *kv = kernel + f * klen;
        for (int b = 0; b < B; b++) {
            const float *in_row = input + b * T;

            float w0 = kv[0];
            for (int t = 0; t < n_win; t++)
                acc[t] = w0 * in_row[t];
            for (int k = 1; k < klen; k++) {
                float wk = kv[k];
                for (int t = 0; t < n_win; t++)
                    acc[t] += wk * in_row[t + k];
            }

            float *out_base = output + b * (n_win * factor) + f;
            for (int t = 0; t < n_win; t++)
                out_base[t * factor] = acc[t];
        }
    }
    free(acc);
}

/* ── v2: AVX2 FMA + shared reads + vectorised interleave ───────────── */
/*    Specialised for factor=2 (the gstlal upsample-by-2 case).        */

void upsample_fast_c_v2(float *restrict output,
                         const float *restrict input,
                         const float *restrict kernel,
                         int B, int T, int n_win, int klen, int factor) {
    /*
     * Both sub-kernels are computed per batch so each input row (~4 KB)
     * is loaded once from L1 rather than once per sub-kernel.
     *
     * Inner loop uses explicit _mm256_fmadd_ps.
     * Output interleave uses AVX2 unpack+lane-permute → two contiguous
     * 256-bit stores, eliminating the scalar stride-2 scatter entirely.
     */
    int acc_len = (n_win + 7) & ~7;
    float *acc0 = (float *)aligned_alloc(32, acc_len * sizeof(float));
    float *acc1 = (float *)aligned_alloc(32, acc_len * sizeof(float));

    const float *kv0 = kernel;
    const float *kv1 = kernel + klen;

    for (int b = 0; b < B; b++) {
        const float *in_row = input + b * T;
        float *out_row = output + b * n_win * 2;

        /* ---- first tap: initialise both accumulators ---- */
        {
            __m256 w0 = _mm256_set1_ps(kv0[0]);
            __m256 w1 = _mm256_set1_ps(kv1[0]);
            for (int t = 0; t < n_win; t += 8) {
                __m256 v = _mm256_loadu_ps(in_row + t);
                _mm256_store_ps(acc0 + t, _mm256_mul_ps(w0, v));
                _mm256_store_ps(acc1 + t, _mm256_mul_ps(w1, v));
            }
        }

        /* ---- remaining taps: shared input load + FMA ---- */
        for (int k = 1; k < klen; k++) {
            __m256 w0 = _mm256_set1_ps(kv0[k]);
            __m256 w1 = _mm256_set1_ps(kv1[k]);
            const float *src = in_row + k;
            for (int t = 0; t < n_win; t += 8) {
                __m256 v  = _mm256_loadu_ps(src + t);
                __m256 a0 = _mm256_load_ps(acc0 + t);
                __m256 a1 = _mm256_load_ps(acc1 + t);
                _mm256_store_ps(acc0 + t, _mm256_fmadd_ps(w0, v, a0));
                _mm256_store_ps(acc1 + t, _mm256_fmadd_ps(w1, v, a1));
            }
        }

        /* ---- AVX2 interleave: acc0,acc1 → [a0,a1,a0,a1,...] ---- */
        for (int t = 0; t < n_win; t += 8) {
            __m256 a0 = _mm256_load_ps(acc0 + t);
            __m256 a1 = _mm256_load_ps(acc1 + t);

            __m256 lo = _mm256_unpacklo_ps(a0, a1);   /* 0 8 1 9 | 4 C 5 D */
            __m256 hi = _mm256_unpackhi_ps(a0, a1);   /* 2 A 3 B | 6 E 7 F */

            _mm256_storeu_ps(out_row + 2*t,
                             _mm256_permute2f128_ps(lo, hi, 0x20));
            _mm256_storeu_ps(out_row + 2*t + 8,
                             _mm256_permute2f128_ps(lo, hi, 0x31));
        }
    }

    free(acc0);
    free(acc1);
}

/* ── v3: v2 + identity shortcut for sub-kernel 0 ──────────────────── */
/*    Sub-kernel 0 is delta (centre=1, rest≈0): memcpy replaces FMA.   */
/*    Only sub-kernel 1 is convolved → half the FMA work of v2.        */

void upsample_fast_c_v3(float *restrict output,
                         const float *restrict input,
                         const float *restrict kernel,
                         int B, int T, int n_win, int klen, int factor) {
    int acc_len = (n_win + 7) & ~7;
    float *acc0 = (float *)aligned_alloc(32, acc_len * sizeof(float));
    float *acc1 = (float *)aligned_alloc(32, acc_len * sizeof(float));

    const int half_length = klen / 2;   /* centre offset for delta copy */
    const float *kv1 = kernel + klen;   /* sub-kernel 1 */

    for (int b = 0; b < B; b++) {
        const float *in_row = input + b * T;
        float *out_row = output + b * n_win * 2;

        /* ---- sub-kernel 0: identity (delta) → aligned copy ---- */
        {
            const float *src = in_row + half_length;
            for (int t = 0; t < n_win; t += 8)
                _mm256_store_ps(acc0 + t, _mm256_loadu_ps(src + t));
        }

        /* ---- sub-kernel 1: first tap → initialise acc1 ---- */
        {
            __m256 w1 = _mm256_set1_ps(kv1[0]);
            for (int t = 0; t < n_win; t += 8) {
                __m256 v = _mm256_loadu_ps(in_row + t);
                _mm256_store_ps(acc1 + t, _mm256_mul_ps(w1, v));
            }
        }

        /* ---- sub-kernel 1: remaining taps via FMA ---- */
        for (int k = 1; k < klen; k++) {
            __m256 w1 = _mm256_set1_ps(kv1[k]);
            const float *src = in_row + k;
            for (int t = 0; t < n_win; t += 8) {
                __m256 v  = _mm256_loadu_ps(src + t);
                __m256 a1 = _mm256_load_ps(acc1 + t);
                _mm256_store_ps(acc1 + t, _mm256_fmadd_ps(w1, v, a1));
            }
        }

        /* ---- AVX2 interleave: acc0,acc1 → [a0,a1,a0,a1,...] ---- */
        for (int t = 0; t < n_win; t += 8) {
            __m256 a0 = _mm256_load_ps(acc0 + t);
            __m256 a1 = _mm256_load_ps(acc1 + t);

            __m256 lo = _mm256_unpacklo_ps(a0, a1);
            __m256 hi = _mm256_unpackhi_ps(a0, a1);

            _mm256_storeu_ps(out_row + 2*t,
                             _mm256_permute2f128_ps(lo, hi, 0x20));
            _mm256_storeu_ps(out_row + 2*t + 8,
                             _mm256_permute2f128_ps(lo, hi, 0x31));
        }
    }

    free(acc0);
    free(acc1);
}
