SIMD + Karatsuba polynomial multiplication experiments

polymul.c

#define _POSIX_C_SOURCE 199309L
#include <immintrin.h>
#include "profiler.h"
#include "nmod_poly.h"
#include "ulong_extras.h"
#include "nmod.h"


void n_poly_mul_n_classical(ulong * c, const ulong * a, const ulong * b, slong n)
{
    slong i, j;
    for (i = 0; i < 2 * n - 1; i++)
        c[i] = 0;
    for (i = 0; i < n; i++)
        for (j = 0; j < n; j++)
            c[i + j] += a[i] * b[j];
}

void d_poly_mul_n_classical(double * c, const double * a, const double * b, slong n)
{
    slong i, j;
    for (i = 0; i < 2 * n - 1; i++)
        c[i] = 0.0;
    for (i = 0; i < n; i++)
        for (j = 0; j < n; j++)
            c[i + j] += a[i] * b[j];
}

#define AVX2_ALIGNED __attribute__((aligned(32)))

FLINT_FORCE_INLINE void
n_poly_mul_4(ulong * res, const ulong * a, const ulong * b)
{
    ulong a0, a1, a2, a3;
    ulong b0, b1, b2, b3;

    a0 = a[0];
    b0 = b[0];
    res[0] = a0 * b0;
    a1 = a[1];
    b1 = b[1];
    res[1] = a0 * b1 + a1 * b0;
    a2 = a[2];
    b2 = b[2];
    res[2] = a0 * b2 + a1 * b1 + a2 * b0;
    a3 = a[3];
    b3 = b[3];
    res[3] = a0 * b3 + a1 * b2 + a2 * b1 + a3 * b0;
    res[4] = a1 * b3 + a2 * b2 + a3 * b1;
    res[5] = a2 * b3 + a3 * b2;
    res[6] = a3 * b3;
    /* Required by Karatsuba */
    res[7] = 0;
}

#define DEF_N_POLY_MUL_N_KARATSUBA(N, M) \
__attribute__((noinline)) void n_poly_mul_ ## N(ulong * c, \
            const ulong * a, \
            const ulong * b) \
{ \
    ulong a_sum[M] AVX2_ALIGNED, b_sum[M] AVX2_ALIGNED; \
    ulong z1[N] AVX2_ALIGNED; \
    slong i; \
    for (i = 0; i < M; i++) \
    { \
        a_sum[i] = a[i] + a[i + M]; \
        b_sum[i] = b[i] + b[i + M]; \
    } \
    n_poly_mul_ ## M(c, a, b); \
    n_poly_mul_ ## M(c + N, a + M, b + M); \
    n_poly_mul_ ## M(z1, a_sum, b_sum); \
    for (i = 0; i < N; i++) \
        z1[i] -= c[i] + c[i + N]; \
    for (i = 0; i < N; i++) \
        c[i + M] += z1[i]; \
}

DEF_N_POLY_MUL_N_KARATSUBA(8, 4)
DEF_N_POLY_MUL_N_KARATSUBA(16, 8)
DEF_N_POLY_MUL_N_KARATSUBA(32, 16)
DEF_N_POLY_MUL_N_KARATSUBA(64, 32)
DEF_N_POLY_MUL_N_KARATSUBA(128, 64)
DEF_N_POLY_MUL_N_KARATSUBA(256, 128)

void
n_poly_mul_n(ulong * c, const ulong * a, const ulong * b, slong n)
{
    switch (n)
    {
        case 16: n_poly_mul_16(c, a, b); break;
        case 32: n_poly_mul_32(c, a, b); break;
        case 64: n_poly_mul_64(c, a, b); break;
        case 128: n_poly_mul_128(c, a, b); break;
        case 256: n_poly_mul_256(c, a, b); break;
        default:
            flint_abort();
    }
}


#define NAILS 8
#define NAIL_MASK ((UWORD(1) << (64 - NAILS)) - 1)

FLINT_FORCE_INLINE
void
n_nailmul(ulong * hi, ulong * lo, ulong a, ulong b)
{
    ulong uhi, ulo;
    umul_ppmm(uhi, ulo, a, b);
    *lo = ulo & NAIL_MASK;
    *hi = (ulo >> (64 - NAILS)) | (uhi << NAILS);
}

FLINT_FORCE_INLINE void
n2_poly_mul_4(ulong * res, const ulong * a, const ulong * b)
{
    ulong a0, a1, a2, a3;
    ulong b0, b1, b2, b3;
    ulong s0, s1, t0, t1;

    // c0 = a0 b0
    a0 = a[0];
    b0 = b[0];
    n_nailmul(&s1, &s0, a0, b0);
    res[0] = s0;
    res[1] = s1;

    // c1 = a0 b1 + a1 b0
    a1 = a[1];
    b1 = b[1];
    n_nailmul(&s1, &s0, a0, b1);
    n_nailmul(&t1, &t0, a1, b0);
    s0 += t0;
    s1 += t1;
    res[2] = s0;
    res[3] = s1;

    // c2 = a0 b2 + a1 b1 + a2 b0
    a2 = a[2];
    b2 = b[2];
    n_nailmul(&s1, &s0, a0, b2);
    n_nailmul(&t1, &t0, a1, b1);
    s0 += t0;
    s1 += t1;
    n_nailmul(&t1, &t0, a2, b0);
    s0 += t0;
    s1 += t1;
    res[4] = s0;
    res[5] = s1;

    // c3 = a0 b3 + a1 b2 + a2 b1 + a3 b0
    a3 = a[3];
    b3 = b[3];
    n_nailmul(&s1, &s0, a0, b3);
    n_nailmul(&t1, &t0, a1, b2);
    s0 += t0;
    s1 += t1;
    n_nailmul(&t1, &t0, a2, b1);
    s0 += t0;
    s1 += t1;
    n_nailmul(&t1, &t0, a3, b0);
    s0 += t0;
    s1 += t1;
    res[6] = s0;
    res[7] = s1;

    // c4 = a1 b3 + a2 b2 + a3 b1
    n_nailmul(&s1, &s0, a1, b3);
    n_nailmul(&t1, &t0, a2, b2);
    s0 += t0;
    s1 += t1;
    n_nailmul(&t1, &t0, a3, b1);
    s0 += t0;
    s1 += t1;
    res[8] = s0;
    res[9] = s1;

    // c5 = a2 b3 + a3 b2
    n_nailmul(&s1, &s0, a2, b3);
    n_nailmul(&t1, &t0, a3, b2);
    s0 += t0;
    s1 += t1;
    res[10] = s0;
    res[11] = s1;

    // c6 = a3 b3
    n_nailmul(&s1, &s0, a3, b3);
    res[12] = s0;
    res[13] = s1;

    // c7 = 0
    res[14] = 0;
    res[15] = 0;
}

#define DEF_N2_POLY_MUL_N_KARATSUBA(N, M) \
__attribute__((noinline)) void n2_poly_mul_ ## N(ulong * c, \
            const ulong * a, \
            const ulong * b) \
{ \
    ulong a_sum[M] AVX2_ALIGNED, b_sum[M] AVX2_ALIGNED; \
    ulong z1[2 * N] AVX2_ALIGNED; \
    slong i; \
    for (i = 0; i < M; i++) \
    { \
        a_sum[i] = a[i] + a[i + M]; \
        b_sum[i] = b[i] + b[i + M]; \
    } \
    n2_poly_mul_ ## M(c, a, b); \
    n2_poly_mul_ ## M(c + 2 * N, a + M, b + M); \
    n2_poly_mul_ ## M(z1, a_sum, b_sum); \
    for (i = 0; i < 2 * N; i++) \
        z1[i] -= c[i] + c[i + 2 * N]; \
    for (i = 0; i < 2 * N; i++) \
        c[i + 2 * M] += z1[i]; \
}

DEF_N2_POLY_MUL_N_KARATSUBA(8, 4)
DEF_N2_POLY_MUL_N_KARATSUBA(16, 8)
DEF_N2_POLY_MUL_N_KARATSUBA(32, 16)
DEF_N2_POLY_MUL_N_KARATSUBA(64, 32)
DEF_N2_POLY_MUL_N_KARATSUBA(128, 64)
DEF_N2_POLY_MUL_N_KARATSUBA(256, 128)


void
n2_poly_mul_n(ulong * c, const ulong * a, const ulong * b, slong n)
{
    switch (n)
    {
        case 4: n2_poly_mul_4(c, a, b); break;
        case 8: n2_poly_mul_8(c, a, b); break;
        case 16: n2_poly_mul_16(c, a, b); break;
        case 32: n2_poly_mul_32(c, a, b); break;
        case 64: n2_poly_mul_64(c, a, b); break;
        case 128: n2_poly_mul_128(c, a, b); break;
        case 256: n2_poly_mul_256(c, a, b); break;
        default:
            flint_abort();
    }
}


// Claude-generated
void d_poly_mul_16_unaligned(double * c, const double * a, const double * b)
{
    // -------------------------------------------------------------------------
    // Load A (pass 0 — aligned, no shift)
    // -------------------------------------------------------------------------
    __m256d A0 = _mm256_load_pd(a +  0); // [a0,  a1,  a2,  a3]
    __m256d A1 = _mm256_load_pd(a +  4); // [a4,  a5,  a6,  a7]
    __m256d A2 = _mm256_load_pd(a +  8); // [a8,  a9,  a10, a11]
    __m256d A3 = _mm256_load_pd(a + 12); // [a12, a13, a14, a15]

    __m256d S0, S1, S2, S3, S4, S5, S6, S7;

#define B(k) _mm256_set1_pd(b[k])

    // =========================================================================
    // PASS 0: no shift, multiply by b[0], b[4], b[8], b[12]
    // =========================================================================
    S0 = _mm256_mul_pd(A0, B( 0));
    S1 = _mm256_mul_pd(A1, B( 0));
    S2 = _mm256_mul_pd(A2, B( 0));
    S3 = _mm256_mul_pd(A3, B( 0));

    S1 = _mm256_fmadd_pd(A0, B( 4), S1);
    S2 = _mm256_fmadd_pd(A1, B( 4), S2);
    S3 = _mm256_fmadd_pd(A2, B( 4), S3);
    S4 = _mm256_mul_pd  (A3, B( 4));

    S2 = _mm256_fmadd_pd(A0, B( 8), S2);
    S3 = _mm256_fmadd_pd(A1, B( 8), S3);
    S4 = _mm256_fmadd_pd(A2, B( 8), S4);
    S5 = _mm256_mul_pd  (A3, B( 8));

    S3 = _mm256_fmadd_pd(A0, B(12), S3);
    S4 = _mm256_fmadd_pd(A1, B(12), S4);
    S5 = _mm256_fmadd_pd(A2, B(12), S5);
    S6 = _mm256_mul_pd  (A3, B(12));

    // =========================================================================
    // PASS 1: shift = 1  →  loadu(a + offset - 1), mask lane 0 of first chunk
    //
    //   A0 = loadu(a-1) & mask(0x1) = [0,   a0,  a1,  a2]
    //   A1 = loadu(a+3)              = [a3,  a4,  a5,  a6]
    //   A2 = loadu(a+7)              = [a7,  a8,  a9,  a10]
    //   A3 = loadu(a+11)             = [a11, a12, a13, a14]
    //   A4 = loadu(a+15) & mask(0x1) = [a15, 0,   0,   0]   (only lane 0 valid)
    //
    // A4 here is a single element.  loadu(a+15) reads a[15..18]; a[16..18] are
    // garbage, but we zero lanes 1-3 via blend before use.
    // =========================================================================
    {
        __m256d zero = _mm256_setzero_pd();

        // Left edge: loadu reads [a[-1], a[0], a[1], a[2]]; lane 0 is garbage → zero it.
        A0 = _mm256_blend_pd(_mm256_loadu_pd(a - 1), zero, 0x1); // [0, a0, a1, a2]

        // Interior chunks: fully in-bounds, no masking needed.
        A1 = _mm256_loadu_pd(a + 3);   // [a3, a4, a5, a6]
        A2 = _mm256_loadu_pd(a + 7);   // [a7, a8, a9, a10]
        A3 = _mm256_loadu_pd(a + 11);  // [a11, a12, a13, a14]

        // Right edge: loadu reads [a[15], a[16], a[17], a[18]]; lanes 1-3 garbage → zero.
        __m256d A4 = _mm256_blend_pd(zero, _mm256_loadu_pd(a + 15), 0x1); // [a15, 0, 0, 0]

        // -- b[1] -------------------------------------------------------------
        S0 = _mm256_fmadd_pd(A0, B( 1), S0);
        S1 = _mm256_fmadd_pd(A1, B( 1), S1);
        S2 = _mm256_fmadd_pd(A2, B( 1), S2);
        S3 = _mm256_fmadd_pd(A3, B( 1), S3);
        S4 = _mm256_fmadd_pd(A4, B( 1), S4);

        // -- b[5] -------------------------------------------------------------
        S1 = _mm256_fmadd_pd(A0, B( 5), S1);
        S2 = _mm256_fmadd_pd(A1, B( 5), S2);
        S3 = _mm256_fmadd_pd(A2, B( 5), S3);
        S4 = _mm256_fmadd_pd(A3, B( 5), S4);
        S5 = _mm256_fmadd_pd(A4, B( 5), S5);

        // -- b[9] -------------------------------------------------------------
        S2 = _mm256_fmadd_pd(A0, B( 9), S2);
        S3 = _mm256_fmadd_pd(A1, B( 9), S3);
        S4 = _mm256_fmadd_pd(A2, B( 9), S4);
        S5 = _mm256_fmadd_pd(A3, B( 9), S5);
        S6 = _mm256_fmadd_pd(A4, B( 9), S6);

        // -- b[13] ------------------------------------------------------------
        S3 = _mm256_fmadd_pd(A0, B(13), S3);
        S4 = _mm256_fmadd_pd(A1, B(13), S4);
        S5 = _mm256_fmadd_pd(A2, B(13), S5);
        S6 = _mm256_fmadd_pd(A3, B(13), S6);
        S7 = _mm256_mul_pd  (A4, B(13));     // first write to S7
    }

    // =========================================================================
    // PASS 2: shift = 2  →  loadu(a + offset - 2), mask lanes 0-1 of first chunk
    //
    //   A0 = loadu(a-2) & mask(0x3) = [0,   0,   a0,  a1]
    //   A1 = loadu(a+2)              = [a2,  a3,  a4,  a5]
    //   A2 = loadu(a+6)              = [a6,  a7,  a8,  a9]
    //   A3 = loadu(a+10)             = [a10, a11, a12, a13]
    //   A4 = loadu(a+14) & mask(0x3) = [a14, a15, 0,   0]
    // =========================================================================
    {
        __m256d zero = _mm256_setzero_pd();

        // Left edge: lanes 0-1 are garbage (a[-2], a[-1]) → zero both.
        A0 = _mm256_blend_pd(_mm256_loadu_pd(a - 2), zero, 0x3); // [0, 0, a0, a1]

        A1 = _mm256_loadu_pd(a + 2);   // [a2, a3, a4, a5]
        A2 = _mm256_loadu_pd(a + 6);   // [a6, a7, a8, a9]
        A3 = _mm256_loadu_pd(a + 10);  // [a10, a11, a12, a13]

        // Right edge: lanes 2-3 are garbage (a[16], a[17]) → zero them.
        __m256d A4 = _mm256_blend_pd(zero, _mm256_loadu_pd(a + 14), 0x3); // [a14, a15, 0, 0]

        // -- b[2] -------------------------------------------------------------
        S0 = _mm256_fmadd_pd(A0, B( 2), S0);
        S1 = _mm256_fmadd_pd(A1, B( 2), S1);
        S2 = _mm256_fmadd_pd(A2, B( 2), S2);
        S3 = _mm256_fmadd_pd(A3, B( 2), S3);
        S4 = _mm256_fmadd_pd(A4, B( 2), S4);

        // -- b[6] -------------------------------------------------------------
        S1 = _mm256_fmadd_pd(A0, B( 6), S1);
        S2 = _mm256_fmadd_pd(A1, B( 6), S2);
        S3 = _mm256_fmadd_pd(A2, B( 6), S3);
        S4 = _mm256_fmadd_pd(A3, B( 6), S4);
        S5 = _mm256_fmadd_pd(A4, B( 6), S5);

        // -- b[10] ------------------------------------------------------------
        S2 = _mm256_fmadd_pd(A0, B(10), S2);
        S3 = _mm256_fmadd_pd(A1, B(10), S3);
        S4 = _mm256_fmadd_pd(A2, B(10), S4);
        S5 = _mm256_fmadd_pd(A3, B(10), S5);
        S6 = _mm256_fmadd_pd(A4, B(10), S6);

        // -- b[14] ------------------------------------------------------------
        S3 = _mm256_fmadd_pd(A0, B(14), S3);
        S4 = _mm256_fmadd_pd(A1, B(14), S4);
        S5 = _mm256_fmadd_pd(A2, B(14), S5);
        S6 = _mm256_fmadd_pd(A3, B(14), S6);
        S7 = _mm256_fmadd_pd(A4, B(14), S7);
    }

    // =========================================================================
    // PASS 3: shift = 3  →  loadu(a + offset - 3), mask lanes 0-2 of first chunk
    //
    //   A0 = loadu(a-3) & mask(0x7) = [0,   0,   0,   a0]
    //   A1 = loadu(a+1)              = [a1,  a2,  a3,  a4]
    //   A2 = loadu(a+5)              = [a5,  a6,  a7,  a8]
    //   A3 = loadu(a+9)              = [a9,  a10, a11, a12]
    //   A4 = loadu(a+13) & mask(0x7) = [a13, a14, a15, 0]
    // =========================================================================
    {
        __m256d zero = _mm256_setzero_pd();

        // Left edge: lanes 0-2 are garbage (a[-3], a[-2], a[-1]) → zero all three.
        A0 = _mm256_blend_pd(_mm256_loadu_pd(a - 3), zero, 0x7); // [0, 0, 0, a0]

        A1 = _mm256_loadu_pd(a + 1);   // [a1, a2, a3, a4]
        A2 = _mm256_loadu_pd(a + 5);   // [a5, a6, a7, a8]
        A3 = _mm256_loadu_pd(a + 9);   // [a9, a10, a11, a12]

        // Right edge: lane 3 is garbage (a[16]) → zero it.
        __m256d A4 = _mm256_blend_pd(zero, _mm256_loadu_pd(a + 13), 0x7); // [a13, a14, a15, 0]

        // -- b[3] -------------------------------------------------------------
        S0 = _mm256_fmadd_pd(A0, B( 3), S0);
        S1 = _mm256_fmadd_pd(A1, B( 3), S1);
        S2 = _mm256_fmadd_pd(A2, B( 3), S2);
        S3 = _mm256_fmadd_pd(A3, B( 3), S3);
        S4 = _mm256_fmadd_pd(A4, B( 3), S4);

        // -- b[7] -------------------------------------------------------------
        S1 = _mm256_fmadd_pd(A0, B( 7), S1);
        S2 = _mm256_fmadd_pd(A1, B( 7), S2);
        S3 = _mm256_fmadd_pd(A2, B( 7), S3);
        S4 = _mm256_fmadd_pd(A3, B( 7), S4);
        S5 = _mm256_fmadd_pd(A4, B( 7), S5);

        // -- b[11] ------------------------------------------------------------
        S2 = _mm256_fmadd_pd(A0, B(11), S2);
        S3 = _mm256_fmadd_pd(A1, B(11), S3);
        S4 = _mm256_fmadd_pd(A2, B(11), S4);
        S5 = _mm256_fmadd_pd(A3, B(11), S5);
        S6 = _mm256_fmadd_pd(A4, B(11), S6);

        // -- b[15] ------------------------------------------------------------
        S3 = _mm256_fmadd_pd(A0, B(15), S3);
        S4 = _mm256_fmadd_pd(A1, B(15), S4);
        S5 = _mm256_fmadd_pd(A2, B(15), S5);
        S6 = _mm256_fmadd_pd(A3, B(15), S6);
        S7 = _mm256_fmadd_pd(A4, B(15), S7);
    }

#undef B

    _mm256_store_pd(c +  0, S0);
    _mm256_store_pd(c +  4, S1);
    _mm256_store_pd(c +  8, S2);
    _mm256_store_pd(c + 12, S3);
    _mm256_store_pd(c + 16, S4);
    _mm256_store_pd(c + 20, S5);
    _mm256_store_pd(c + 24, S6);
    _mm256_store_pd(c + 28, S7);
}



#define d_poly_mul_16 d_poly_mul_16_unaligned


#define DEF_D_POLY_MUL_N_KARATSUBA(N, M) \
__attribute__((noinline)) void d_poly_mul_ ## N(double * c, \
            const double * a, \
            const double * b) \
{ \
    double a_sum[M] AVX2_ALIGNED, b_sum[M] AVX2_ALIGNED; \
    double z1[N] AVX2_ALIGNED; \
    slong i; \
    for (i = 0; i < M; i++) \
    { \
        a_sum[i] = a[i] + a[i + M]; \
        b_sum[i] = b[i] + b[i + M]; \
    } \
    d_poly_mul_ ## M(c, a, b); \
    d_poly_mul_ ## M(c + N, a + M, b + M); \
    d_poly_mul_ ## M(z1, a_sum, b_sum); \
    for (i = 0; i < N; i++) \
        z1[i] -= c[i] + c[i + N]; \
    for (i = 0; i < N; i++) \
        c[i + M] += z1[i]; \
}

DEF_D_POLY_MUL_N_KARATSUBA(32, 16)
DEF_D_POLY_MUL_N_KARATSUBA(64, 32)
DEF_D_POLY_MUL_N_KARATSUBA(128, 64)
DEF_D_POLY_MUL_N_KARATSUBA(256, 128)
DEF_D_POLY_MUL_N_KARATSUBA(512, 256)

void
d_poly_mul_n(double * c, const double * a, const double * b, slong n)
{
    switch (n)
    {
        case 16: d_poly_mul_16(c, a, b); break;
        case 32: d_poly_mul_32(c, a, b); break;
        case 64: d_poly_mul_64(c, a, b); break;
        case 128: d_poly_mul_128(c, a, b); break;
        case 256: d_poly_mul_256(c, a, b); break;
        case 512: d_poly_mul_512(c, a, b); break;
        default:
            flint_abort();
    }
}



#include "machine_vectors.h"


FLINT_FORCE_INLINE void vec8n_store_unaligned(ulong* z, vec8n a) {
    vec4n_store_unaligned(z+0, a.e1);
    vec4n_store_unaligned(z+4, a.e2);
}


FLINT_FORCE_INLINE vec8n vec8d_convert_limited_vec8n(vec8d a) {
    vec8n z = {vec4d_convert_limited_vec4n(a.e1), vec4d_convert_limited_vec4n(a.e2)};
    return z;
}

FLINT_FORCE_INLINE vec4n vec4n_bit_shift_left(vec4n a, ulong b) {
    return _mm256_slli_epi64(a, b);
}

FLINT_FORCE_INLINE vec8n vec8n_bit_shift_left(vec8n a, ulong b) {
    vec8n z = {vec4n_bit_shift_left(a.e1, b), vec4n_bit_shift_left(a.e2, b)};
    return z;
}


void
_nmod_vec_reduce_simd_d(nn_ptr res, const double * vec, slong len, nmod_t mod)
{
    slong i;

    vec4d p = vec4d_set_d((double) mod.n);
    vec4d pinv = vec4d_set_d(1.0 / mod.n);

    for (i = 0; i < len; i += 4)
    {
        vec4d d = vec4d_load_unaligned(vec + i);
        vec4d r = vec4d_reduce_to_0n(d, p, pinv);
        vec4n u = vec4d_convert_limited_vec4n(r);
        vec4n_store_unaligned(res + i, u);
    }
}

void
_nmod_poly_mul_n_double(nn_ptr res, nn_srcptr a, nn_srcptr b, slong len, nmod_t mod)
{
    double R[512] AVX2_ALIGNED;
    /* 4 padding entries to allow kernel to read out of bounds */
    double A[4 + 256 + 4] AVX2_ALIGNED;
    double B[4 + 256 + 4] AVX2_ALIGNED;

    slong i;
    for (i = 0; i < len; i++)
        A[4 + i] = a[i];
    for (i = 0; i < len; i++)
        B[4 + i] = b[i];

    d_poly_mul_n(R, A + 4, B + 4, len);

    _nmod_vec_reduce_simd_d(res, R, 2 * len, mod);
}

void
_nmod_poly_mul_n_uint64(nn_ptr res, nn_srcptr a, nn_srcptr b, slong len, nmod_t mod)
{
    ulong R[512] AVX2_ALIGNED;

    slong i;
    n_poly_mul_n(R, a, b, len);

    ulong npre = n_barrett_precomp(mod.n);

    for (i = 0; i < 2 * len - 1; i++)
        res[i] = n_mod_barrett(R[i], mod.n, npre);
}

void
_nmod_poly_mul_n_2x_double(nn_ptr res, nn_srcptr a, nn_srcptr b, slong len, nmod_t mod)
{
    double Rlo[512] AVX2_ALIGNED;
    double Rhi[512] AVX2_ALIGNED;
    double A[4 + 256 + 4] AVX2_ALIGNED;
    double Blo[4 + 256 + 4] AVX2_ALIGNED;
    double Bhi[4 + 256 + 4] AVX2_ALIGNED;

    slong i;
    for (i = 0; i < len; i++)
        A[i + 4] = a[i];

    for (i = 0; i < len; i++)
    {
        Blo[i + 4] = b[i] & ((UWORD(1) << 14) - 1);
        Bhi[i + 4] = b[i] >> 14;
    }

    d_poly_mul_n(Rlo, A + 4, Blo + 4, len);
    d_poly_mul_n(Rhi, A + 4, Bhi + 4, len);

    ulong npre = n_barrett_precomp(mod.n);

    for (i = 0; i < len; i++)
    {
        res[i] = n_mod_barrett((ulong) Rlo[i] + ((ulong) Rhi[i] << 14), mod.n, npre);
    }
}

#include <stdint.h>
#include "nmod.h"
#include "nmod_vec.h"

// Claude-generated
FLINT_FORCE_INLINE ulong reconstruct_hi128(ulong x_lo, double x_approx)
{
    static const double TWO64 = 0x1p64;
    static const double INV_TWO64 = 0x1p-64;

    ulong hi = (int64_t)(x_approx * INV_TWO64);
    double rem_approx = x_approx - (double)hi * TWO64;
    int64_t delta = (int64_t)x_lo - (int64_t)(ulong)rem_approx;
    if (delta > (int64_t)0x7FFFFFFFFFFFFFFFll / 2) {
        hi += 1;
    } else if (0 && delta < -(int64_t)0x7FFFFFFFFFFFFFFFll / 2) {
        hi -= 1;
    }

    return hi;
}

void
_nmod_poly_mul_n_uint64_double(nn_ptr res, nn_srcptr a, nn_srcptr b, slong len, nmod_t mod)
{
    ulong Rlo[1024] AVX2_ALIGNED;

    n_poly_mul_n(Rlo, a, b, len);

    double Rhi[1024] AVX2_ALIGNED;
    double A[4 + 512 + 4] AVX2_ALIGNED;
    double B[4 + 512 + 4] AVX2_ALIGNED;

    slong i;
    for (i = 0; i < len; i++)
        A[i + 4] = a[i];

    for (i = 0; i < len; i++)
        B[i + 4] = b[i];

    d_poly_mul_n(Rhi, A + 4, B + 4, len);

    for (i = 0; i < 2 * len - 1; i++)
    {
        ulong lo = Rlo[i];
        ulong hi = reconstruct_hi128(lo, Rhi[i]);
        NMOD_RED2(res[i], hi, lo, mod);
    }
}

// Claude-generated
/**
 * Convert a 128-bit value from two-word signed base-2^56 redundant
 * representation to standard base-2^64 (two uint64_t limbs).
 *
 * Input:  V = lo + hi * 2^56  (lo, hi are int64_t; V is guaranteed >= 0)
 * Output: V = r0 + r1 * 2^64  (r0, r1 are uint64_t; standard representation)
 *
 * lo or hi may be negative; the overall value V must be non-negative
 * and fit in 128 bits.
 */
FLINT_FORCE_INLINE void
from_signed_base2_56_redundant(int64_t lo, int64_t hi,
                                     uint64_t *r0, uint64_t *r1)
{
    /*
     * Step 1: Extract the carry from lo with a SIGNED arithmetic shift.
     *
     *   lo = lo_56 + carry * 2^56
     *
     * We want carry = floor(lo / 2^56), the signed quotient, so that
     * lo_56 = lo - carry * 2^56 is in [0, 2^56).
     *
     * carry = lo >> 56  (arithmetic right shift, rounds toward -inf for negative lo)
     *
     * In C, arithmetic right shift of int64_t is NOT guaranteed by the standard,
     * but is universally supported on two's-complement targets. We confirm:
     *   - If lo >= 0: carry = lo >> 56 in [0, 255],  lo_56 = lo & MASK56 in [0, 2^56)
     *   - If lo <  0: carry = -((-lo + 2^56-1) >> 56), i.e. carry in [-128, -1],
     *                 lo_56 = lo - carry*2^56 in [0, 2^56)   ✓
     *
     * Alternatively, the portable version uses:
     *   carry = (lo >> 63) | ((uint64_t)lo >> 56)   -- see note below
     * but the arithmetic shift path is cleaner and universally used in practice.
     */

    const uint64_t MASK56 = (UINT64_C(1) << 56) - 1;

    int64_t  carry  = lo >> 56;                        /* signed floor-divide by 2^56 */
    uint64_t lo_56  = (uint64_t)lo & MASK56;           /* lo mod 2^56, always in [0, 2^56) */

    /*
     * Step 2: Absorb the carry into hi.
     *
     *   M = hi + carry
     *
     * Since V >= 0, M >= 0 is guaranteed:
     *
     * Proof: V = lo_56 + (hi + carry) * 2^56 >= 0, and lo_56 in [0, 2^56),
     *        so (hi + carry) >= 0 follows from V >= 0 and lo_56 < 2^56.
     * More precisely: M = (V - lo_56) / 2^56 >= -lo_56/2^56 > -1, so M >= 0.
     *
     * carry in [-128, 255] (8-bit signed range from int64_t >> 56)
     * hi in [-2^63, 2^63-1]
     * → M = hi + carry fits in int64_t (no overflow since |carry| <= 255)
     *   and M >= 0, so we can safely cast to uint64_t.
     */

    int64_t  M_signed = hi + carry;                    /* guaranteed >= 0             */
    uint64_t M        = (uint64_t)M_signed;            /* safe: M_signed >= 0         */

    /*
     * Step 3: Reconstruct base-2^64 limbs (identical to unsigned case).
     *
     *   V = lo_56 + M * 2^56
     *     = lo_56 + (M & 0xFF) * 2^56 + (M >> 8) * 2^64
     *
     *   r0 = lo_56 | (M << 56)
     *   r1 = M >> 8
     */

    *r0 = lo_56 | (M << 56);
    *r1 = M >> 8;
}

void
_nmod_poly_mul_n_uint64_ll(nn_ptr res, nn_srcptr a, nn_srcptr b, slong len, nmod_t mod)
{
    ulong R[1024] AVX2_ALIGNED;
    slong i;

    n2_poly_mul_n(R, a, b, len);

    for (i = 0; i < 2 * len - 1; i++)
    {
        ulong lo = R[2 * i];
        ulong hi = R[2 * i + 1];
        from_signed_base2_56_redundant(lo, hi, &lo, &hi);
        NMOD_RED2(res[i], hi, lo, mod);
    }
}


#define MAXN 256

#include "nmod.h"

int main(void)
{
    slong n, i;
    flint_rand_t state;
    flint_rand_init(state);

    double FLINT_SET_BUT_UNUSED(tcpu), t1, t2, t3;
    nn_ptr A, B, C, D;

    A = flint_malloc(sizeof(ulong) * MAXN);
    B = flint_malloc(sizeof(ulong) * MAXN);
    C = flint_malloc(sizeof(ulong) * 2 * MAXN);
    D = flint_malloc(sizeof(ulong) * 2 * MAXN);

    nmod_t mod;

    nmod_init(&mod, 1000001);

    flint_printf("MOD = %wu\n", mod.n);

    flint_printf("%4s   %14s %14s     (speedup)\n", "len", "_nmod_poly_mul", "double");

    for (n = 16; n <= MAXN; n *= 2)
    {
        for (i = 0; i < n; i++)
        {
            A[i] = n_randint(state, mod.n);
            B[i] = n_randint(state, mod.n);
            C[i] = n_randint(state, mod.n);
            D[i] = n_randint(state, mod.n);
        }

        TIMEIT_START;
        _nmod_poly_mul(C, A, n, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t1);

        TIMEIT_START;
        _nmod_poly_mul_n_double(D, A, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t2);

        if (!_nmod_vec_equal(C, D, 2 * n - 1)) flint_abort();

        flint_printf("%4wd   %14g %14g     (%.3f)\n", n, t1, t2, t1 / t2);
    }

    flint_printf("\n");



    nmod_init(&mod, 100000001);

    flint_printf("MOD = %wu\n", mod.n);

    flint_printf("%4s   %14s %14s   (%s)  %14s    (%s)\n",
        "len", "_nmod_poly_mul", "uint64", "speedup", "2x double", "speedup");

    for (n = 16; n <= MAXN; n *= 2)
    {
        for (i = 0; i < n; i++)
        {
            A[i] = n_randint(state, mod.n);
            B[i] = n_randint(state, mod.n);
            C[i] = n_randint(state, mod.n);
            D[i] = n_randint(state, mod.n);
        }

        TIMEIT_START;
        _nmod_poly_mul(C, A, n, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t1);

        TIMEIT_START;
        _nmod_poly_mul_n_uint64(D, A, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t2);

        if (!_nmod_vec_equal(C, D, 2 * n - 1)) flint_abort();

        TIMEIT_START;
        _nmod_poly_mul_n_2x_double(D, A, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t3);

        if (!_nmod_vec_equal(C, D, 2 * n - 1)) flint_abort();

        flint_printf("%4wd   %14g %14g     (%.3f)  %14g     (%.3f)\n",
            n, t1, t2, t1 / t2, t3, t1 / t3);
    }

    flint_printf("\n");


    nmod_init(&mod, n_pow(10, 15) + 1);

    flint_printf("MOD = %wu\n", mod.n);

    flint_printf("%4s   %14s %14s   (%s)  %14s    (%s)\n",
        "len", "_nmod_poly_mul", "uint64+double", "speedup", "uint64-ll", "speedup");

    for (n = 16; n <= MAXN; n *= 2)
    {
        for (i = 0; i < n; i++)
        {
            A[i] = n_randint(state, mod.n);
            B[i] = n_randint(state, mod.n);
            C[i] = n_randint(state, mod.n);
            D[i] = n_randint(state, mod.n);
        }

        TIMEIT_START;
        _nmod_poly_mul(C, A, n, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t1);

        TIMEIT_START;
        _nmod_poly_mul_n_uint64_double(D, A, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t2);

        if (!_nmod_vec_equal(C, D, 2 * n - 1)) flint_abort();

        TIMEIT_START;
        _nmod_poly_mul_n_uint64_ll(D, A, B, n, mod);
        TIMEIT_STOP_VALUES(tcpu, t3);

        if (!_nmod_vec_equal(C, D, 2 * n - 1)) flint_abort();

        flint_printf("%4wd   %14g %14g     (%.3f)  %14g     (%.3f)\n",
            n, t1, t2, t1 / t2, t3, t1 / t3);
    }

    flint_free(A);
    flint_free(B);
    flint_free(C);
    flint_free(D);
    return 0;
}