/*
References:
- duck.ac user saffah_cc_v41_260924, https://duck.ac/submission/114230:
copied its exact centered mixed-radix NTT and tail-recovery convolution.
The cached public source contains no separate license notice.
Idea:
- After both tiled corner transforms, complete A forward, B forward,
pointwise product, and A inverse consecutively for each independent row.
This loop fusion reduces the reuse distance of transform values from the
whole array to one row; modular arithmetic and reconstruction stay exact.
- Replace the reference's three inline assembly wrappers by corresponding
standard compiler SIMD intrinsics. No assembly or system calls are used.
Purpose:
- Test whether row-local scheduling improves official runtime. Local tests
check complete output equivalence, input preservation, and output guards.
*/
#ifndef LGR_
#define LGR_ 12
#endif
#ifndef M5L_
#define M5L_ 13
#endif
#ifndef LGR2_
#define LGR2_ 13
#endif
#ifndef LGC2_
#define LGC2_ 13
#endif
#define EXACT 1
#define PADC_ 16
#define PADC2_ 16
#define N5_ 5
#define M5_ (1u << M5L_)
#define NR_ (1u << LGR_)
#define NC_ (N5_ * M5_)
#define N1_ ((u32)(NR_ * NC_))
#define CSX_ (NC_ + PADC_)
#define AOFF(k) (((long)(k) / (long)NC_) * CSX_ + ((long)(k) % (long)NC_))
#define NR2_ (1u << LGR2_)
#define NC2_ (1u << LGC2_)
#define N2_ (1u << (LGR2_ + LGC2_))
#define CSX2_ (NC2_ + PADC2_)
#include <cstdint>
#define NO_AVX2_UNUSED 0
#include <cstring>
#ifndef NO_AVX2
#include <immintrin.h>
#endif
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;
typedef int64_t i64;
typedef int32_t i32;
typedef unsigned __int128 u128;
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __attribute__((always_inline)) __m256i ldu256(const void *p) {
return _mm256_loadu_si256((const __m256i *)p);
}
static inline __attribute__((always_inline)) void stu256(void *p, __m256i v) {
_mm256_storeu_si256((__m256i *)p, v);
}
static inline __attribute__((always_inline)) __m256i vaddm256(__m256i a, __m256i b, __m256i Pv) {
__m256i s = _mm256_add_epi32(a, b);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, Pv));
}
static inline __attribute__((always_inline)) __m256i vsubm256(__m256i a, __m256i b, __m256i Pv) {
__m256i d = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
}
#undef CMBSTEP
#pragma GCC pop_options
#pragma GCC optimize("O3","rename-registers","live-range-shrinkage","ira-loop-pressure","no-stack-protector")
#define TS 64
#define PADB_ 16
#define NRMAX_ ((NR_ > NR2_) ? NR_ : NR2_)
static const u32 P = 2013265921u;
static const u32 G = 31u;
static u64 IM;
static u32 TBUF[TS * TS];
static u32 TW[TS * TS];
static inline u32 mm(u32 a, u32 b) {
u64 z = (u64)a * b;
u64 q = (u64)(((u128)z * IM) >> 64);
u64 r = z - q * (u64)P;
return (u32)(r >= P ? r - P : r);
}
static inline u32 shoup_c(u32 w) { return (u32)(((u64)w << 32) / (u64)P); }
static u32 pw(u32 a, u32 e) {
u32 r = 1;
while (e) { if (e & 1) r = (u32)((u64)r * a % P); a = (u32)((u64)a * a % P); e >>= 1; }
return r;
}
static void build_tab(u32 *tw, int L, u32 w) {
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
u32 wl = pw(w, (u32)(L / len));
u32 c = 1;
int h = len >> 1;
for (int j = 0; j < h; j++) { tw[off + j] = c; c = mm(c, wl); }
off += h;
}
}
static void build_br(u16 *br, int L) {
int bits = 0; while ((1 << bits) < L) bits++;
for (int i = 0; i < L; i++) {
int r = 0;
for (int b = 0; b < bits; b++) if (i & (1 << b)) r |= 1 << (bits - 1 - b);
br[i] = (u16)r;
}
}
#ifdef NO_AVX2
#define HAVE_AVX2 0
#else
#define HAVE_AVX2 1
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __m256i vmulhi8(__m256i a, __m256i b) {
__m256i e = _mm256_mul_epu32(a, b);
__m256i o = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32));
e = _mm256_srli_epi64(e, 32);
return _mm256_blend_epi32(e, o, 0xAA);
}
static inline __attribute__((always_inline)) __m256i mullo256(__m256i a, __m256i b) {
return _mm256_mullo_epi32(a, b);
}
static inline __m256i shoup8(__m256i v, __m256i w, __m256i ws) {
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i t = mullo256(v, w);
__m256i q = vmulhi8(v, ws);
__m256i r = _mm256_sub_epi32(t, mullo256(q, Pv));
return _mm256_min_epu32(r, _mm256_sub_epi32(r, Pv));
}
#define SP(w, ws) _mm256_and_si256(w,M32), _mm256_srli_epi64(w,32), _mm256_and_si256(ws,M32), _mm256_srli_epi64(ws,32)
static void row_dit_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
int off = 0;
if (L >= 8) {
const __m256i V4 = _mm256_setr_epi32((int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2], 1, 1);
const __m256i V4s = _mm256_setr_epi32((int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2], 2, 2);
const __m256i V8 = _mm256_setr_epi32((int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6], 1, 1, 1, 1);
const __m256i V8s = _mm256_setr_epi32((int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6], 2, 2, 2, 2);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
const __m256i IDXQ4 = _mm256_setr_epi32(0, 1, 0, 1, 4, 5, 4, 5);
const __m256i IDXQ8 = _mm256_setr_epi32(0, 1, 2, 3, 0, 1, 2, 3);
int i4 = 0;
for (; i4 + 64 <= L; i4 += 64) {
__m256i vr0 = ldu256((const void *)((x + i4 + 0)));
__m256i vr1 = ldu256((const void *)((x + i4 + 8)));
__m256i vr2 = ldu256((const void *)((x + i4 + 16)));
__m256i vr3 = ldu256((const void *)((x + i4 + 24)));
__m256i vr4 = ldu256((const void *)((x + i4 + 32)));
__m256i vr5 = ldu256((const void *)((x + i4 + 40)));
__m256i vr6 = ldu256((const void *)((x + i4 + 48)));
__m256i vr7 = ldu256((const void *)((x + i4 + 56)));
{
{
__m256i y = _mm256_shuffle_epi32(vr0, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr0);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr0 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr1, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr1);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr1 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr2, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr2);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr2 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr3, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr3);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr3 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr4, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr4);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr4 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr5, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr5);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr5 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr6, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr6);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr6 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr7, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr7);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr7 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr0, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr0, y), _mm256_add_epi32(_mm256_sub_epi32(vr0, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr1, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr1, y), _mm256_add_epi32(_mm256_sub_epi32(vr1, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr2, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr2, y), _mm256_add_epi32(_mm256_sub_epi32(vr2, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr3, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr3, y), _mm256_add_epi32(_mm256_sub_epi32(vr3, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr4, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr4, y), _mm256_add_epi32(_mm256_sub_epi32(vr4, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr5, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr5, y), _mm256_add_epi32(_mm256_sub_epi32(vr5, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr6, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr6, y), _mm256_add_epi32(_mm256_sub_epi32(vr6, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr7, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr7, y), _mm256_add_epi32(_mm256_sub_epi32(vr7, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr0, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr0, y), _mm256_add_epi32(_mm256_sub_epi32(vr0, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr1, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr1, y), _mm256_add_epi32(_mm256_sub_epi32(vr1, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr2, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr2, y), _mm256_add_epi32(_mm256_sub_epi32(vr2, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr3, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr3, y), _mm256_add_epi32(_mm256_sub_epi32(vr3, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr4, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr4, y), _mm256_add_epi32(_mm256_sub_epi32(vr4, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr5, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr5, y), _mm256_add_epi32(_mm256_sub_epi32(vr5, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr6, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr6, y), _mm256_add_epi32(_mm256_sub_epi32(vr6, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr7, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr7, y), _mm256_add_epi32(_mm256_sub_epi32(vr7, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
vr7 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
stu256((void *)((x + i4 + 0)), vr0);
stu256((void *)((x + i4 + 8)), vr1);
stu256((void *)((x + i4 + 16)), vr2);
stu256((void *)((x + i4 + 24)), vr3);
stu256((void *)((x + i4 + 32)), vr4);
stu256((void *)((x + i4 + 40)), vr5);
stu256((void *)((x + i4 + 48)), vr6);
stu256((void *)((x + i4 + 56)), vr7);
}
for (; i4 < L; i4 += 8) {
__m256i v = ldu256((const void *)((x + i4)));
{
__m256i y = _mm256_shuffle_epi32(v, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, v);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
{
__m256i y = _mm256_shuffle_epi32(v, 0x4E);
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8);
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
stu256((void *)((x + i4)), v);
}
off = 7;
int len = 16;
for (; (len << 1) <= L; len <<= 2) {
int h = len >> 1;
const u32 *t1 = tw + off, *ts1 = tws + off;
const u32 *t2 = t1 + h, *ts2 = ts1 + h;
const u32 *t3 = t2 + h, *ts3 = ts2 + h;
off += 3 * h;
for (int i = 0; i < L; i += (len << 1)) {
u32 *p = x + i;
for (int j = 0; j < h; j += 8) {
__m256i WV1 = ldu256((const void *)((t1 + j))), WS1 = ldu256((const void *)((ts1 + j)));
__m256i WV2 = ldu256((const void *)((t2 + j))), WS2 = ldu256((const void *)((ts2 + j)));
__m256i WV3 = ldu256((const void *)((t3 + j))), WS3 = ldu256((const void *)((ts3 + j)));
__m256i A = ldu256((const void *)(p + j));
__m256i B = ldu256((const void *)(p + j + h));
__m256i C = ldu256((const void *)(p + j + 2 * h));
__m256i D = ldu256((const void *)(p + j + 3 * h));
__m256i Bm = shoup8(B, WV1, WS1), Dm = shoup8(D, WV1, WS1);
__m256i Ap = vaddm256(A, Bm, Pv), Bp = vsubm256(A, Bm, Pv);
__m256i Cp = vaddm256(C, Dm, Pv), Dp = vsubm256(C, Dm, Pv);
__m256i Cm = shoup8(Cp, WV2, WS2), Dn = shoup8(Dp, WV3, WS3);
stu256((void *)(p + j), vaddm256(Ap, Cm, Pv));
stu256((void *)(p + j + 2 * h), vsubm256(Ap, Cm, Pv));
stu256((void *)(p + j + h), vaddm256(Bp, Dn, Pv));
stu256((void *)(p + j + 3 * h), vsubm256(Bp, Dn, Pv));
}
}
}
if (len <= L) {
int h = len >> 1;
const u32 *t = tw + off, *ts = tws + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)(p + j));
__m256i v = shoup8(ldu256((const void *)(q + j)),
ldu256((const void *)(t + j)), ldu256((const void *)(ts + j)));
stu256((void *)(p + j), vaddm256(u, v, Pv));
stu256((void *)(q + j), vsubm256(u, v, Pv));
}
}
}
} else {
for (int len = 2; len <= L; len <<= 1) {
int h = len >> 1;
const u32 *t = tw + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = mm(q[j], t[j]);
u32 sa = u + v; if (sa >= P) sa -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = sa; q[j] = d;
}
}
}
}
}
static void row_dif_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
if (L >= 8) {
int len = L;
for (; len >= 32; len >>= 2) {
int h = len >> 1, h2 = len >> 2;
const u32 *t1 = tw + h - 1, *ts1 = tws + h - 1;
const u32 *t2 = t1 + h2, *ts2 = ts1 + h2;
const u32 *t3 = tw + h2 - 1, *ts3 = tws + h2 - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i;
for (int j = 0; j < h2; j += 8) {
__m256i WV1 = ldu256((const void *)((t1 + j))), WS1 = ldu256((const void *)((ts1 + j)));
__m256i WV2 = ldu256((const void *)((t2 + j))), WS2 = ldu256((const void *)((ts2 + j)));
__m256i WV3 = ldu256((const void *)((t3 + j))), WS3 = ldu256((const void *)((ts3 + j)));
__m256i v0 = ldu256((const void *)((p + j)));
__m256i v1 = ldu256((const void *)((p + j + h2)));
__m256i v2 = ldu256((const void *)((p + j + h)));
__m256i v3 = ldu256((const void *)((p + j + h + h2)));
__m256i s02 = _mm256_min_epu32(_mm256_add_epi32(v0, v2), _mm256_sub_epi32(_mm256_add_epi32(v0, v2), Pv));
__m256i d02 = _mm256_sub_epi32(v0, v2);
__m256i s13 = _mm256_min_epu32(_mm256_add_epi32(v1, v3), _mm256_sub_epi32(_mm256_add_epi32(v1, v3), Pv));
__m256i d13 = _mm256_sub_epi32(v1, v3);
d02 = _mm256_min_epu32(d02, _mm256_add_epi32(d02, Pv));
d13 = _mm256_min_epu32(d13, _mm256_add_epi32(d13, Pv));
d02 = shoup8(d02, WV1, WS1);
d13 = shoup8(d13, WV2, WS2);
__m256i ra = _mm256_min_epu32(_mm256_add_epi32(s02, s13), _mm256_sub_epi32(_mm256_add_epi32(s02, s13), Pv));
__m256i rb = _mm256_min_epu32(_mm256_sub_epi32(s02, s13), _mm256_add_epi32(_mm256_sub_epi32(s02, s13), Pv));
rb = shoup8(rb, WV3, WS3);
__m256i rc = _mm256_min_epu32(_mm256_add_epi32(d02, d13), _mm256_sub_epi32(_mm256_add_epi32(d02, d13), Pv));
__m256i rd = _mm256_min_epu32(_mm256_sub_epi32(d02, d13), _mm256_add_epi32(_mm256_sub_epi32(d02, d13), Pv));
rd = shoup8(rd, WV3, WS3);
stu256((void *)((p + j)), ra);
stu256((void *)((p + j + h2)), rb);
stu256((void *)((p + j + h)), rc);
stu256((void *)((p + j + h + h2)), rd);
}
}
}
if (len >= 16) {
int h = len >> 1;
const u32 *t = tw + h - 1, *ts = tws + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)((p + j)));
__m256i v = ldu256((const void *)((q + j)));
__m256i w = ldu256((const void *)((t + j)));
__m256i ws = ldu256((const void *)((ts + j)));
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(u, v), _mm256_sub_epi32(_mm256_add_epi32(u, v), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(u, v), _mm256_add_epi32(_mm256_sub_epi32(u, v), Pv));
d = shoup8(d, w, ws);
stu256((void *)((p + j)), sa);
stu256((void *)((q + j)), d);
}
}
}
const __m256i V8 = _mm256_setr_epi32(1, 1, 1, 1, (int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6]);
const __m256i V8s = _mm256_setr_epi32(2, 2, 2, 2, (int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6]);
const __m256i V4 = _mm256_setr_epi32(1, 1, (int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2]);
const __m256i V4s = _mm256_setr_epi32(2, 2, (int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2]);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
int i4 = 0;
for (; i4 + 64 <= L; i4 += 64) {
__m256i vr0 = ldu256((const void *)((x + i4 + 0)));
__m256i vr1 = ldu256((const void *)((x + i4 + 8)));
__m256i vr2 = ldu256((const void *)((x + i4 + 16)));
__m256i vr3 = ldu256((const void *)((x + i4 + 24)));
__m256i vr4 = ldu256((const void *)((x + i4 + 32)));
__m256i vr5 = ldu256((const void *)((x + i4 + 40)));
__m256i vr6 = ldu256((const void *)((x + i4 + 48)));
__m256i vr7 = ldu256((const void *)((x + i4 + 56)));
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr0, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr0);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr1, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr1);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr2, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr2);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr3, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr3);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr4, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr4);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr5, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr5);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr6, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr6);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr7, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr7);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr7 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr0, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr0);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr1, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr1);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr2, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr2);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr3, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr3);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr4, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr4);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr5, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr5);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr6, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr6);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr7, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr7);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr0, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr0), _mm256_add_epi32(_mm256_sub_epi32(y, vr0), Pv));
vr0 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr1, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr1), _mm256_add_epi32(_mm256_sub_epi32(y, vr1), Pv));
vr1 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr2, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr2), _mm256_add_epi32(_mm256_sub_epi32(y, vr2), Pv));
vr2 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr3, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr3), _mm256_add_epi32(_mm256_sub_epi32(y, vr3), Pv));
vr3 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr4, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr4), _mm256_add_epi32(_mm256_sub_epi32(y, vr4), Pv));
vr4 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr5, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr5), _mm256_add_epi32(_mm256_sub_epi32(y, vr5), Pv));
vr5 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr6, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr6), _mm256_add_epi32(_mm256_sub_epi32(y, vr6), Pv));
vr6 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr7, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr7), _mm256_add_epi32(_mm256_sub_epi32(y, vr7), Pv));
vr7 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
stu256((void *)((x + i4 + 0)), vr0);
stu256((void *)((x + i4 + 8)), vr1);
stu256((void *)((x + i4 + 16)), vr2);
stu256((void *)((x + i4 + 24)), vr3);
stu256((void *)((x + i4 + 32)), vr4);
stu256((void *)((x + i4 + 40)), vr5);
stu256((void *)((x + i4 + 48)), vr6);
stu256((void *)((x + i4 + 56)), vr7);
}
for (; i4 < L; i4 += 8) {
__m256i v = ldu256((const void *)((x + i4)));
{
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, v);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
{
__m256i y = _mm256_shuffle_epi32(v, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_sub_epi32(y, v);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{
__m256i y = _mm256_shuffle_epi32(v, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, v), _mm256_add_epi32(_mm256_sub_epi32(y, v), Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
stu256((void *)((x + i4)), v);
}
} else {
for (int len = L; len >= 2; len >>= 1) {
int h = len >> 1;
const u32 *t = tw + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = q[j];
u32 sa = u + v; if (sa >= P) sa -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = sa; q[j] = mm(d, t[j]);
}
}
}
}
}
static inline __m256i shoup_of(__m256i w) {
const __m256i ML = _mm256_set1_epi32(572662301);
__m256i lo = _mm256_mul_epu32(w, ML);
__m256i hi = _mm256_mul_epu32(_mm256_srli_epi64(w, 32), ML);
__m256i m = _mm256_blend_epi32(_mm256_srli_epi64(lo, 32), _mm256_slli_epi64(_mm256_srli_epi64(hi, 32), 32), 0xAA);
return _mm256_add_epi32(_mm256_add_epi32(w, w), m);
}
static inline void t8x8(__m256i *r) {
__m256i t0 = _mm256_unpacklo_epi32(r[0], r[1]), t1 = _mm256_unpackhi_epi32(r[0], r[1]);
__m256i t2 = _mm256_unpacklo_epi32(r[2], r[3]), t3 = _mm256_unpackhi_epi32(r[2], r[3]);
__m256i t4 = _mm256_unpacklo_epi32(r[4], r[5]), t5 = _mm256_unpackhi_epi32(r[4], r[5]);
__m256i t6 = _mm256_unpacklo_epi32(r[6], r[7]), t7 = _mm256_unpackhi_epi32(r[6], r[7]);
__m256i u0 = _mm256_unpacklo_epi64(t0, t2), u1 = _mm256_unpackhi_epi64(t0, t2);
__m256i u2 = _mm256_unpacklo_epi64(t1, t3), u3 = _mm256_unpackhi_epi64(t1, t3);
__m256i u4 = _mm256_unpacklo_epi64(t4, t6), u5 = _mm256_unpackhi_epi64(t4, t6);
__m256i u6 = _mm256_unpacklo_epi64(t5, t7), u7 = _mm256_unpackhi_epi64(t5, t7);
r[0] = _mm256_permute2x128_si256(u0, u4, 0x20); r[1] = _mm256_permute2x128_si256(u1, u5, 0x20);
r[2] = _mm256_permute2x128_si256(u2, u6, 0x20); r[3] = _mm256_permute2x128_si256(u3, u7, 0x20);
r[4] = _mm256_permute2x128_si256(u0, u4, 0x31); r[5] = _mm256_permute2x128_si256(u1, u5, 0x31);
r[6] = _mm256_permute2x128_si256(u2, u6, 0x31); r[7] = _mm256_permute2x128_si256(u3, u7, 0x31);
}
static void pointwise_mul_avx(u32 *A, const u32 *B, u32 n, u32 sc, u32 sx) {
(void)sc;
long k = 0;
for (int r = 0; k < (long)n; r++) {
long lim = (long)sx - PADC_;
if (lim > (long)n - k) lim = (long)n - k;
u32 *a = A + (long)r * sx;
const u32 *b = B + (long)r * sx;
long i = 0;
for (; i + 8 <= lim; i += 8) {
__m256i x = ldu256((const void *)(a + i));
__m256i y = ldu256((const void *)(b + i));
stu256((void *)(a + i), shoup8(x, y, shoup_of(y)));
}
for (; i < lim; i++) a[i] = mm(a[i], b[i]);
k += lim;
}
}
static u32 NPI = 0;
static const u32 R32 = (u32)((1ull << 32) % P);
static inline u32 mont1(u32 a, u32 b) {
u64 z = (u64)a * (u64)b;
u32 m = (u32)((u64)(u32)z * (u64)NPI);
u64 t = (z + (u64)m * (u64)P) >> 32;
return (u32)(t >= P ? t - P : t);
}
static inline __m256i mont8(__m256i a, __m256i b) {
__m256i P64 = _mm256_set1_epi64x((long long)P), NP = _mm256_set1_epi64x((long long)NPI);
__m256i Pv = _mm256_set1_epi32((int)P);
__m256i ahi = _mm256_srli_epi64(a, 32), bhi = _mm256_srli_epi64(b, 32);
__m256i tlo = _mm256_mul_epu32(a, b), thi = _mm256_mul_epu32(ahi, bhi);
__m256i e2 = _mm256_mul_epu32(tlo, NP), o2 = _mm256_mul_epu32(thi, NP);
__m256i e4 = _mm256_mul_epu32(e2, P64), o4 = _mm256_mul_epu32(o2, P64);
__m256i e5 = _mm256_add_epi64(tlo, e4), o5 = _mm256_add_epi64(thi, o4);
__m256i ue = _mm256_srli_epi64(e5, 32), uo = _mm256_srli_epi64(o5, 32);
__m256i u = _mm256_blend_epi32(ue, _mm256_slli_epi64(uo, 32), 0xAA);
return _mm256_min_epu32(u, _mm256_sub_epi32(u, Pv));
}
static void scale_row_mont(u32 *dst, const u32 *src, u32 kM) {
const __m256i kv = _mm256_set1_epi32((int)kM);
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), mont8(ldu256((const void *)((src + v))), kv));
}
static void twiddle_rows_mont(u32 *TW, const u32 *RVM) {
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8)
stu256((void *)((cur + v)), mont8(ldu256((const void *)((prv + v))), ldu256((const void *)((RVM + v)))));
}
}
static void scale_row_avx(u32 *dst, const u32 *src, u32 k) {
const __m256i kv = _mm256_set1_epi32((int)k), ks = _mm256_set1_epi32((int)shoup_c(k));
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), shoup8(ldu256((const void *)((src + v))), kv, ks));
}
static void twiddle_rows_avx(u32 *TW, const u32 *RV) {
__m256i RVS[TS / 8];
for (int v = 0; v < TS; v += 8) RVS[v >> 3] = shoup_of(ldu256((const void *)((RV + v))));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8) {
__m256i r = ldu256((const void *)((RV + v)));
stu256((void *)((cur + v)), shoup8(ldu256((const void *)((prv + v))), r, RVS[v >> 3]));
}
}
}
static void corner_tile_avx(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++) {
const u32 *s = in + (u64)(A0 + u0 + k) * strIn + B0 + v0;
const u32 *tw = TW + (u0 + k) * TS + v0;
__m256i x = ldu256((const void *)(s));
__m256i w = ldu256((const void *)(tw));
r[k] = mont8(x, w);
}
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8) stu256((void *)(d + u), ldu256((const void *)(s + u)));
}
}
static void corner_tile_avx_blk(const u32 *in, u32 *out, int B0, int strIn, int strOut) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++) {
const u32 *s = in + (u64)(u0 + k) * strIn + B0 + v0;
const u32 *tw = TW + (u0 + k) * TS + v0;
r[k] = mont8(ldu256((const void *)(s)), ldu256((const void *)(tw)));
}
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8) stu256((void *)(d + u), ldu256((const void *)(s + u)));
}
}
static void load_tile_avx(const u32 *src, u32 *S, int r0, int c0, int R, int C, int ctr, int allzero) {
for (int rb0 = 0; rb0 < TS; rb0 += 8) {
for (int cb0 = 0; cb0 < TS; cb0 += 16) {
__m256i r[8];
if (allzero) {
for (int k = 0; k < 8; k++) r[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r[k] = x;
}
}
t8x8(r);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0 + j) * TS + rb0), r[j]);
{
int cb0b = cb0 + 8;
__m256i r2[8];
if (allzero) {
for (int k = 0; k < 8; k++) r2[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0b;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r2[k] = x;
}
}
t8x8(r2);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0b + j) * TS + rb0), r2[j]);
}
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb += 8) stu256((void *)(d + rb), ldu256((const void *)(s + rb)));
}
}
static void load_tile_avx_d(const u32 *src, u32 *S, int r0, int c0, int C, int dstr, int ctr, int allzero) {
for (int rb0 = 0; rb0 < TS; rb0 += 8) {
for (int cb0 = 0; cb0 < TS; cb0 += 8) {
__m256i r[8];
if (allzero) {
for (int k = 0; k < 8; k++) r[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r[k] = x;
}
}
t8x8(r);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0 + j) * TS + rb0), r[j]);
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)cb * dstr + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb += 8) stu256((void *)(d + rb), ldu256((const void *)(s + rb)));
}
}
static void load_block(const u32 *src, u32 nsrc, u32 *S, int c0, int R, int C, int dstr, int ctr) {
for (int r0 = 0; r0 < R; r0 += TS) {
u64 firstidx = (u64)r0 * C + c0;
u64 lastidx = (u64)(r0 + TS - 1) * C + c0 + TS - 1;
if (lastidx < nsrc) { load_tile_avx_d(src, S, r0, c0, C, dstr, ctr, 0); continue; }
if (firstidx >= nsrc) { load_tile_avx_d(src, S, r0, c0, C, dstr, ctr, 1); continue; }
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)cb * dstr + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
static void copy_u32_avx(u32 *dst, const u32 *src, long n) {
long i = 0;
for (; i + 8 <= n; i += 8) stu256((void *)(dst + i), ldu256((const void *)(src + i)));
for (; i < n; i++) dst[i] = src[i];
}
static void copy_u32_nt(u32 *dst, const u32 *src, long n) {
long i = 0;
long ka = ((-(long)(((uintptr_t)dst) >> 2)) & 7);
for (; i < ka && i < n; i++) dst[i] = src[i];
for (; i + 8 <= n; i += 8) {
__m256i v = ldu256((const void *)(src + i));
_mm256_stream_si256((__m256i *)(dst + i), v);
}
for (; i < n; i++) dst[i] = src[i];
_mm_sfence();
}
static void untr_tile_avx(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 16) {
__m256i r[8];
for (int k = 0; k < 8; k++)
r[k] = ldu256((const void *)((in + (u64)(A0 + u0 + k) * strIn + B0 + v0)));
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
{
int v0b = v0 + 8;
__m256i r2[8];
for (int k = 0; k < 8; k++)
r2[k] = ldu256((const void *)((in + (u64)(A0 + u0 + k) * strIn + B0 + v0b)));
t8x8(r2);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0b + k) * TS + u0), r2[k]);
}
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8)
_mm256_stream_si256((__m256i *)(d + u), ldu256((const void *)(s + u)));
}
}
static void build_tab2(u32 *tw, u32 *tws, int L, u32 w) {
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
u32 wl = pw(w, (u32)(L / len));
u32 c = 1;
int h = len >> 1;
for (int j = 0; j < h; j++) { tw[off + j] = c; tws[off + j] = shoup_c(c); c = mm(c, wl); }
off += h;
}
}
#pragma GCC pop_options
#endif
static u32 R5F[6], R5FS[6];
static u32 R5I[6], R5IS[6];
static u32 T5F[4 * M5_], T5FS[4 * M5_];
static u32 T5I[4 * M5_], T5IS[4 * M5_];
static u32 tsubFw[M5_], tsubFsw[M5_], tsubIw[M5_], tsubIsw[M5_];
static void build_radix5(u32 root) {
const u32 H2 = 1006632961u;
const u32 H4 = mm(H2, H2);
u32 w5 = pw(root, (u32)(NC_ / 5));
u32 w2 = mm(w5, w5), w3 = mm(w2, w5), w4 = mm(w3, w5);
u32 uu = (u32)(((u64)w5 + w4) % P), vv = (u32)(((u64)w2 + w3) % P);
u32 pp = (u32)(((u64)w5 + P - w4) % P), qq = (u32)(((u64)w2 + P - w3) % P);
u32 c1 = mm((u32)(((u64)uu + vv) % P), H4);
u32 c2 = mm((u32)(((u64)uu + P - vv) % P), H4);
u32 c3 = mm((u32)(((u64)pp + qq) % P), H4);
u32 c4 = mm((u32)(((u64)pp + P - qq) % P), H4);
u32 c5 = mm((u32)(((u64)qq + P - pp) % P), H4);
u32 c6 = mm((u32)(((u64)qq + pp) % P), H4);
R5F[0] = c1; R5F[1] = c2; R5F[2] = c3; R5F[3] = c4; R5F[4] = c5; R5F[5] = c6;
R5I[0] = c1; R5I[1] = c2;
R5I[2] = (c3 ? P - c3 : 0); R5I[3] = (c4 ? P - c4 : 0);
R5I[4] = (c5 ? P - c5 : 0); R5I[5] = (c6 ? P - c6 : 0);
for (int i = 0; i < 6; i++) { R5FS[i] = shoup_c(R5F[i]); R5IS[i] = shoup_c(R5I[i]); }
for (int s = 1; s < 5; s++) {
u32 g = pw(root, (u32)s), gi = pw(g, P - 2u), c = 1, ci = 1;
u32 *pf = T5F + (s - 1) * M5_, *pf2 = T5FS + (s - 1) * M5_;
u32 *pi = T5I + (s - 1) * M5_, *pi2 = T5IS + (s - 1) * M5_;
for (int j = 0; j < (int)M5_; j++) { pf[j] = c; pf2[j] = shoup_c(c); c = mm(c, g);
pi[j] = ci; pi2[j] = shoup_c(ci); ci = mm(ci, gi); }
}
build_tab2(tsubFw, tsubFsw, (int)M5_, pw(root, 5));
build_tab2(tsubIw, tsubIsw, (int)M5_, pw(pw(root, 5), P - 2u));
}
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __attribute__((always_inline)) __m256i vaddm(__m256i a, __m256i b, __m256i Pv) {
__m256i s = _mm256_add_epi32(a, b);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, Pv));
}
static inline __attribute__((always_inline)) __m256i vsubm(__m256i a, __m256i b, __m256i Pv) {
__m256i d = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
}
static inline __attribute__((always_inline)) void dft5(__m256i *v, const __m256i *c,
const __m256i *cs, __m256i Pv) {
__m256i A = vaddm(v[1], v[4], Pv), B = vaddm(v[2], v[3], Pv);
__m256i Cc = vsubm(v[1], v[4], Pv), Dd = vsubm(v[2], v[3], Pv);
__m256i S = vaddm(A, B, Pv), D2 = vsubm(A, B, Pv);
__m256i S2 = vaddm(Cc, Dd, Pv), D3 = vsubm(Cc, Dd, Pv);
__m256i t1 = shoup8(S, c[0], cs[0]), t2 = shoup8(D2, c[1], cs[1]);
__m256i T = vaddm(t1, t2, Pv), V = vsubm(t1, t2, Pv);
__m256i t3 = shoup8(S2, c[2], cs[2]), t4 = shoup8(D3, c[3], cs[3]);
__m256i t5 = shoup8(S2, c[4], cs[4]), t6 = shoup8(D3, c[5], cs[5]);
__m256i U = vaddm(t3, t4, Pv), W = vaddm(t5, t6, Pv);
__m256i Z = vaddm(v[0], T, Pv), Y = vaddm(v[0], V, Pv);
__m256i X0 = vaddm(vaddm(v[0], A, Pv), B, Pv);
__m256i X1 = vaddm(Z, U, Pv), X4 = vsubm(Z, U, Pv);
__m256i X2 = vaddm(Y, W, Pv), X3 = vsubm(Y, W, Pv);
v[0] = X0; v[1] = X1; v[2] = X2; v[3] = X3; v[4] = X4;
}
static void col_dif5(u32 *x) {
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i c[6], cs[6];
for (int i = 0; i < 6; i++) { c[i] = _mm256_set1_epi32((int)R5F[i]); cs[i] = _mm256_set1_epi32((int)R5FS[i]); }
for (int j = 0; j < (int)M5_; j += 8) {
__m256i v[5];
for (int i = 0; i < 5; i++) v[i] = ldu256((const void *)(x + (u64)i * M5_ + j));
dft5(v, c, cs, Pv);
stu256((void *)(x + j), v[0]);
for (int s = 1; s < 5; s++) {
__m256i w = ldu256((const void *)(T5F + (s - 1) * M5_ + j));
__m256i ws = ldu256((const void *)(T5FS + (s - 1) * M5_ + j));
stu256((void *)(x + (u64)s * M5_ + j), shoup8(v[s], w, ws));
}
}
for (int s = 0; s < 5; s++) row_dif_avx(x + (u64)s * M5_, (int)M5_, tsubFw, tsubFsw);
}
static void col_dit5(u32 *x) {
for (int s = 0; s < 5; s++) row_dit_avx(x + (u64)s * M5_, (int)M5_, tsubIw, tsubIsw);
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i c[6], cs[6];
for (int i = 0; i < 6; i++) { c[i] = _mm256_set1_epi32((int)R5I[i]); cs[i] = _mm256_set1_epi32((int)R5IS[i]); }
for (int j = 0; j < (int)M5_; j += 8) {
__m256i v[5];
v[0] = ldu256((const void *)(x + j));
for (int s = 1; s < 5; s++) {
__m256i w = ldu256((const void *)(T5I + (s - 1) * M5_ + j));
__m256i ws = ldu256((const void *)(T5IS + (s - 1) * M5_ + j));
v[s] = shoup8(ldu256((const void *)(x + (u64)s * M5_ + j)), w, ws);
}
dft5(v, c, cs, Pv);
for (int i = 0; i < 5; i++) stu256((void *)(x + (u64)i * M5_ + j), v[i]);
}
}
#pragma GCC pop_options
static void row_dit(u32 *x, int L, const u32 *tw, const u32 *tws) {
#if HAVE_AVX2
row_dit_avx(x, L, tw, tws);
return;
#else
(void)tws;
#endif
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
int h = len >> 1;
const u32 *t = tw + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = mm(q[j], t[j]);
u32 s = u + v; if (s >= P) s -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = s; q[j] = d;
}
}
}
}
static void pointwise(u32 *A, const u32 *B, u32 n, u32 sc, u32 sx) {
#if HAVE_AVX2
pointwise_mul_avx(A, B, n, sc, sx);
return;
#else
for (long i = 0; i < (long)n; i++) { long o = (i / (long)(sx - PADC_)) * sx + (i % (long)(sx - PADC_)); A[o] = mm(mm(A[o], B[o]), sc); }
#endif
}
static void row_dif(u32 *x, int L, const u32 *tw, const u32 *tws) {
#if HAVE_AVX2
row_dif_avx(x, L, tw, tws);
return;
#else
(void)tws;
#endif
for (int len = L; len >= 2; len >>= 1) {
int h = len >> 1;
const u32 *t = tw + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = q[j];
u32 s = u + v; if (s >= P) s -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = s; q[j] = mm(d, t[j]);
}
}
}
}
static inline void bitrev_pass(u32 *x, int L, const u16 *br) {
for (int i = 0; i < L; i++) {
int k = br[i];
if (i < k) { u32 t = x[i]; x[i] = x[k]; x[k] = t; }
}
}
struct Eng {
int R, C, N, sx;
int blkr;
int r5;
u32 w;
u32 *tfR; u32 *tiR; u32 *tfC; u32 *tiC;
u32 *wsR; u32 *wsC;
u16 *brR; u16 *brC;
};
static void colfwd(Eng &e, u32 *row) {
#if HAVE_AVX2
if (e.r5) { col_dif5(row); return; }
#endif
row_dif(row, e.C, e.tfC, e.wsC);
}
static void colinv(Eng &e, u32 *row) {
#if HAVE_AVX2
if (e.r5) { col_dit5(row); return; }
#endif
row_dit(row, e.C, e.tiC, e.wsC + e.C);
}
static void load_T(Eng &e, const u32 *src, u32 nsrc, u32 *S, int ctr) {
const int R = e.R, C = e.C;
#if HAVE_AVX2
if (nsrc <= (u32)R * (u32)C) {
for (int r0 = 0; r0 < R; r0 += TS) {
u64 tilerow = (u64)r0 * C;
for (int c0 = 0; c0 < C; c0 += TS) {
u64 firstidx = tilerow + c0;
u64 lastidx = (u64)(r0 + TS - 1) * C + c0 + TS - 1;
if (lastidx < nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 0); continue; }
if (firstidx >= nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 1); continue; }
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
return;
}
#endif
for (int r0 = 0; r0 < R; r0 += TS) {
for (int c0 = 0; c0 < C; c0 += TS) {
u64 lastidx = (u64)(r0 + TS - 1) * C + (c0 + TS - 1);
if (lastidx < nsrc) {
for (int rb = 0; rb < TS; rb++) {
const u32 *s = src + (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u32 v = s[cb];
if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u));
TBUF[cb * TS + rb] = v;
}
}
} else {
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
}
static u32 REV6[TS];
static u32 g_inv_scale = 1u;
static u32 POWT[1024], POWT64[1024];
static void build_rev6(void) {
for (int i = 0; i < TS; i++) {
int r = 0;
for (int b = 0; b < 6; b++) if (i & (1 << b)) r |= 1 << (5 - b);
REV6[i] = (u32)r;
}
}
static inline u32 revbits(u32 x, int lg) {
u32 r = 0;
for (int b = 0; b < lg; b++) if (x & (1u << b)) r |= 1u << (lg - 1 - b);
return r;
}
static void corner_generic(const u32 *in, u32 *out, int nA, int nB, int strIn, int strOut, u32 root, int perm) {
u32 wv[TS];
{ u32 c = 1; for (int v = 0; v < TS; v++) { wv[v] = c; c = mm(c, root); } }
const u32 w64 = pw(root, TS);
int lgB = 0, lgA = 0;
while ((1 << lgB) < nB) lgB++;
while ((1 << lgA) < nA) lgA++;
const int S1 = nB / TS, S2 = nA / TS;
const int HB = lgB - 6, HA = lgA - 6;
u32 RV2[TS], RV2b[TS], PW2[TS], cu[TS], cu64[TS], base[TS];
if (perm == 1) {
u32 g = pw(root, (u32)S1);
{ u32 gM = mm(g, R32), c = R32; for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; c = mont1(c, gM); } }
{ u32 rootM = mm(root, R32), p = R32; for (int t = 0; t < S1; t++) { POWT[t] = p; p = mont1(p, rootM); } }
} else if (perm == 2) {
u32 g = pw(root, (u32)S2), g64 = pw(g, 64);
const u32 rootM2 = mm(root, R32), gM = mm(g, R32), g64M = mm(g64, R32);
{ u32 c = R32, c64 = R32;
for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; RV2b[REV6[t]] = c64; c = mont1(c, gM); c64 = mont1(c64, g64M); } }
{ u32 p = R32; for (int t = 0; t < S2; t++) { POWT[t] = p; p = mont1(p, rootM2); } }
{ u32 q = pw(root, 64), qM = mm(q, R32), p = R32; for (int t = 0; t < S2; t++) { POWT64[t] = p; p = mont1(p, qM); } }
}
u32 wr = 1;
const u32 INVSCALE = g_inv_scale;
for (int A0 = 0; A0 < nA; A0 += TS) {
u32 wr64 = pw(wr, TS);
if (perm == 1) {
u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[REV6[t]] = c; c = mont1(c, gM2); }
} else if (perm == 2) {
int hi = A0 / TS;
u32 rH = POWT[revbits((u32)hi, HA)], r64H = POWT64[revbits((u32)hi, HA)];
const u32 INVSCALE_M = mm(INVSCALE, R32);
for (int u = 0; u < TS; u++) { cu[u] = mont1(RV2[u], rH); cu64[u] = mont1(RV2b[u], r64H); base[u] = INVSCALE_M; }
} else {
u32 c = 1; for (int v = 0; v < TS; v++) { PW2[v] = c; c = mm(c, wr); }
}
u32 Vloc = 1;
u32 Vr = 1;
for (int B0 = 0; B0 < nB; B0 += TS) {
if (perm == 1) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 wrH = pw(wr, H), rootH = POWT[H];
u32 RV[TS];
#if HAVE_AVX2
scale_row_mont(RV, RV2, rootH);
scale_row_mont(TW, PW2, mm(wrH, R32));
twiddle_rows_mont(TW, RV);
#else
for (int v = 0; v < TS; v++) RV[v] = mont1(RV2[v], rootH);
for (int v = 0; v < TS; v++) TW[v] = mont1(PW2[v], mm(wrH, R32));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mont1(prv[v], RV[v]);
}
#endif
} else if (perm == 2) {
if (B0 == 0) {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
cur[0] = base[u];
for (int v = 1; v < TS; v++) cur[v] = mont1(cur[v - 1], cu[u]);
}
} else {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
#if HAVE_AVX2
scale_row_mont(cur, cur, cu64[u]);
#else
u32 d = cu64[u];
for (int v = 0; v < TS; v++) cur[v] = mont1(cur[v], d);
#endif
}
}
for (int u = 0; u < TS; u++) base[u] = mont1(base[u], cu64[u]);
} else {
u32 RV[TS];
for (int v = 0; v < TS; v++) RV[v] = mm(Vloc, wv[v]);
for (int v = 0; v < TS; v++) TW[v] = mm(Vr, PW2[v]);
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mm(prv[v], RV[v]);
}
}
#if HAVE_AVX2
corner_tile_avx(in, out, A0, B0, strIn, strOut);
#else
for (int u = 0; u < TS; u++) {
const u32 *s = in + (u64)(A0 + u) * strIn + B0;
const u32 *tw = TW + u * TS;
for (int v = 0; v < TS; v++) TBUF[u * TS + v] = mont1(s[v], tw[v]);
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
for (int u = 0; u < TS; u++) d[u] = TBUF[u * TS + v];
}
#endif
Vr = mm(Vr, wr64);
Vloc = mm(Vloc, w64);
}
wr = mm(wr, w64);
}
}
static inline void corner_fwd(Eng &e, const u32 *S, u32 *X, u32 root) {
corner_generic(S, X, e.C, e.R, e.R, e.sx, root, 1);
}
static inline void corner_inv(Eng &e, const u32 *X, u32 *S, u32 root) {
corner_generic(X, S, e.R, e.C, e.sx, e.R, root, 2);
}
static void untranspose(Eng &e, const u32 *S, u32 *dst) {
const int R = e.R, C = e.C;
#if HAVE_AVX2
for (int c0 = 0; c0 < C; c0 += TS)
for (int r0 = 0; r0 < R; r0 += TS)
untr_tile_avx(S, dst, c0, r0, R, e.sx);
return;
#endif
for (int c0 = 0; c0 < C; c0 += TS) {
for (int r0 = 0; r0 < R; r0 += TS) {
for (int u = 0; u < TS; u++) {
const u32 *s = S + (u64)(c0 + u) * R + r0;
u32 *d = TBUF + u * TS;
for (int v = 0; v < TS; v++) d[v] = s[v];
}
for (int v = 0; v < TS; v++) {
u32 *d = dst + (u64)(r0 + v) * e.sx + c0;
for (int u = 0; u < TS; u++) d[u] = TBUF[u * TS + v];
}
}
}
}
static u32 SBLK[TS * (NRMAX_ + PADB_)] __attribute__((aligned(64)));
static void corner_fwd_stream(Eng &e, const u32 *src, u32 nsrc, u32 *X, int ctr) {
const int nA = e.C, nB = e.R, strIn = e.blkr, strOut = e.sx;
const u32 root = e.w;
u32 *S = SBLK;
int lgB = 0; while ((1 << lgB) < nB) lgB++;
const int S1 = nB / TS, HB = lgB - 6;
u32 RV2[TS], PW2[TS];
{ u32 g = pw(root, (u32)S1), gM = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; c = mont1(c, gM); } }
{ u32 rootM = mm(root, R32), p = R32;
for (int t = 0; t < S1; t++) { POWT[t] = p; p = mont1(p, rootM); } }
const u32 w64 = pw(root, TS);
u32 wr = 1;
for (int A0 = 0; A0 < nA; A0 += TS) {
u32 wr64 = pw(wr, TS);
{ u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[REV6[t]] = c; c = mont1(c, gM2); } }
load_block(src, nsrc, S, A0, nB, nA, strIn, ctr);
for (int u = 0; u < TS; u++)
row_dif(S + (u64)u * strIn, nB, e.tfR, e.wsR);
for (int B0 = 0; B0 < nB; B0 += TS) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 wrH = pw(wr, H), rootH = POWT[H];
u32 RV[TS];
scale_row_mont(RV, RV2, rootH);
scale_row_mont(TW, PW2, mm(wrH, R32));
twiddle_rows_mont(TW, RV);
corner_tile_avx_blk(S, X + (u64)A0, B0, strIn, strOut);
}
wr = mm(wr, w64);
}
}
static void fwd(Eng &e, const u32 *src, u32 nsrc, u32 *S, u32 *X, int ctr) {
(void)S;
corner_fwd_stream(e, src, nsrc, X, ctr);
for (int k0 = 0; k0 < e.R; k0++) colfwd(e, X + (u64)k0 * e.sx);
}
static void inv(Eng &e, u32 *X, u32 *S) {
for (int k0 = 0; k0 < e.R; k0++) colinv(e, X + (u64)k0 * e.sx);
corner_inv(e, X, S, pw(e.w, P - 2u));
for (int j0 = 0; j0 < e.C; j0++) {
u32 *row = S + (u64)j0 * e.R;
row_dit(row, e.R, e.tiR, e.wsR + e.R);
}
}
// Complete each transform row while both input rows and its output are nearby.
static void multiply_rows(Eng &e, const u32 *a, u32 na,
const u32 *b, u32 nb, u32 *S,
u32 *A, u32 *B, int ctr, u32 scale) {
corner_fwd_stream(e, a, na, A, ctr);
corner_fwd_stream(e, b, nb, B, ctr);
for (int r = 0; r < e.R; ++r) {
u32 *ar = A + (u64)r * e.sx;
u32 *br = B + (u64)r * e.sx;
colfwd(e, ar);
colfwd(e, br);
pointwise(ar, br, e.C, scale, e.sx);
colinv(e, ar);
}
corner_inv(e, A, S, pw(e.w, P - 2u));
for (int j = 0; j < e.C; ++j)
row_dit(S + (u64)j * e.R, e.R, e.tiR, e.wsR + e.R);
}
static void eng_init(Eng &e, int lgR, int nC, u32 w, int sx,
u32 *tfR, u32 *tiR, u32 *tfC, u32 *tiC, u16 *brR, u16 *brC,
u32 *wsR, u32 *wsC) {
e.R = 1 << lgR; e.C = nC; e.N = e.R * e.C; e.w = w; e.sx = sx;
e.blkr = e.R + PADB_;
e.r5 = (nC == (int)NC_);
e.tfR = tfR; e.tiR = tiR; e.tfC = tfC; e.tiC = tiC; e.brR = brR; e.brC = brC;
e.wsR = wsR; e.wsC = wsC;
u32 wR = pw(w, (u32)e.C), wC = pw(w, (u32)e.R);
#if HAVE_AVX2
build_tab2(tfR, e.wsR, e.R, wR); build_tab2(tiR, e.wsR + e.R, e.R, pw(wR, P - 2u));
if (!e.r5) { build_tab2(tfC, e.wsC, e.C, wC); build_tab2(tiC, e.wsC + e.C, e.C, pw(wC, P - 2u)); }
#else
build_tab(tfR, e.R, wR); build_tab(tiR, e.R, pw(wR, P - 2u));
build_tab(tfC, e.C, wC); build_tab(tiC, e.C, pw(wC, P - 2u));
#endif
build_br(brR, e.R);
if (!e.r5) build_br(brC, e.C);
build_rev6();
}
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __attribute__((always_inline)) __m256i scan8(__m256i v, __m256i &carry) {
v = _mm256_add_epi32(v, _mm256_slli_si256(v, 4));
v = _mm256_add_epi32(v, _mm256_slli_si256(v, 8));
__m256i t = _mm256_permute2x128_si256(_mm256_setzero_si256(), v, 0x20);
t = _mm256_shuffle_epi32(t, 0xFF);
v = _mm256_add_epi32(v, t);
v = _mm256_add_epi32(v, carry);
carry = _mm256_permutevar8x32_epi32(v, _mm256_set1_epi32(7));
return v;
}
#define CMBSTEP() \
{ u32 d; \
if (PH == 0) d = 4u * (a[k] + b[k]) - 16u; \
else if (PH == 1) { if (nm_lt) d = 4u * (b[k] - b[k - 1 - n]); \
else d = 4u * (a[k] - a[k - 1 - m]); } \
else d = 16u - 4u * (a[k - 1 - m] + b[k - 1 - n]); \
SUMr += d; \
u32 tab; \
if (TB == 1) { u32 x = A[AOFF(k)], y = c[N1_ + k]; tab = (x >= y) ? (x - y) : (x + (P - y)); } \
else if (TB == 2) tab = A[AOFF(k)]; \
else tab = c[k]; \
i32 ts = (i32)tab; \
if (ts > (i32)(P >> 1)) ts -= (i32)P; \
c[k] = (u32)((u32)ts + SUMr); \
}
template<int PH, int TB, int NT>
static void cmbv(long k0, long k1, const u32 *a, const u32 *b, u32 *A, u32 *c,
long n, long m, long nmax, u32 &SUMr, int nm_lt)
{
if (k0 > k1) return;
long k = k0;
{
long ka = k + ((-(long)k) & 7);
long kend = (ka < k1 + 1) ? ka : k1 + 1;
if (ka > k) { const __m256i Pv0 = _mm256_set1_epi32((int)P); (void)Pv0; }
for (; k < kend; k++) { CMBSTEP(); }
}
const __m256i Pv = _mm256_set1_epi32((int)P);
const __m256i Phv = _mm256_set1_epi32((int)(P >> 1));
#if HAVE_AVX2
const __m256i c16 = _mm256_set1_epi32(16);
const __m256i idx7 = _mm256_set1_epi32(7);
__m256i carry = _mm256_set1_epi32((int)SUMr);
if (NT) {
carry = _mm256_set1_epi32((int)SUMr);
}
for (; k + 7 <= k1; k += 8) {
__m256i d;
if (PH == 0) {
__m256i va = ldu256(a + k), vb = ldu256(b + k);
d = _mm256_sub_epi32(_mm256_slli_epi32(_mm256_add_epi32(va, vb), 2), c16);
} else if (PH == 1) {
if (nm_lt) { __m256i v1 = ldu256(b + k), v0 = ldu256(b + k - 1 - n);
d = _mm256_slli_epi32(_mm256_sub_epi32(v1, v0), 2); }
else { __m256i v1 = ldu256(a + k), v0 = ldu256(a + k - 1 - m);
d = _mm256_slli_epi32(_mm256_sub_epi32(v1, v0), 2); }
} else {
__m256i va = ldu256(a + k - 1 - m), vb = ldu256(b + k - 1 - n);
d = _mm256_sub_epi32(c16, _mm256_slli_epi32(_mm256_add_epi32(va, vb), 2));
}
d = scan8(d, carry);
__m256i tab;
if (TB == 1) {
__m256i x = ldu256(A + AOFF(k)), y = ldu256(c + N1_ + k);
__m256i dd = _mm256_sub_epi32(x, y);
__m256i mk = _mm256_cmpgt_epi32(y, x);
tab = _mm256_add_epi32(dd, _mm256_and_si256(mk, Pv));
} else if (TB == 2) {
tab = ldu256(A + AOFF(k));
} else {
tab = ldu256(c + k);
}
__m256i ts = _mm256_sub_epi32(tab, _mm256_and_si256(_mm256_cmpgt_epi32(tab, Phv), Pv));
__m256i out = _mm256_add_epi32(ts, d);
if (NT && (((uintptr_t)(c + k)) & 31) == 0) _mm256_stream_si256((__m256i *)(c + k), out);
else stu256(c + k, out);
}
SUMr = (u32)_mm256_extract_epi32(carry, 0);
#endif
for (; k <= k1; k++) { CMBSTEP(); }
}
#pragma GCC pop_options
static u32 Aarr[NR_ * CSX_] __attribute__((aligned(64)));
static u32 Barr[NR_ * CSX_] __attribute__((aligned(64)));
static Eng E1, E2;
static u32 t1fR[NR_], t1iR[NR_], t1fC[1], t1iC[1];
static u16 b1R[NR_], b1C[1];
static u32 w1R[2 * NR_], w1C[2];
static u32 t2fR[NR2_], t2iR[NR2_], t2fC[NC2_], t2iC[NC2_];
static u16 b2R[NR2_], b2C[NC2_];
static u32 w2R[2 * NR2_], w2C[2 * NC2_];
static int inited = 0;
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c)
{
IM = ~0ull / (u64)P;
if (!inited) {
inited = 1;
{ u64 inv = 1; for (int i = 0; i < 5; i++) inv = inv * (2 - (u64)P * inv); NPI = (u32)(-(u64)inv); }
u32 w1 = pw(G, (u32)((P - 1) / N1_));
u32 w2 = pw(G, (u32)((P - 1) / N2_));
eng_init(E1, LGR_, (int)NC_, w1, (int)CSX_, t1fR, t1iR, t1fC, t1iC, b1R, b1C, w1R, w1C);
eng_init(E2, LGR2_, (1 << LGC2_), w2, (int)CSX2_, t2fR, t2iR, t2fC, t2iC, b2R, b2C, w2R, w2C);
E1.r5 = 1; E2.r5 = 0;
#if HAVE_AVX2
build_radix5(pw(w1, (u32)E1.R));
#endif
}
u32 *S = (u32 *)c;
u32 *A = Aarr, *B = Barr;
const int ctr = EXACT;
const u32 sc = pw((u32)(N1_ % P), P - 2u);
const u32 sc2 = pw((u32)(N2_ % P), P - 2u);
const long Sa = (long)N1_ - (long)m;
const long Sb = (long)N1_ - (long)n;
const long M = (long)n + (long)m - (long)N1_;
if (M >= 0) {
g_inv_scale = sc2;
multiply_rows(E2, a + Sa, (u32)(M + 1), b + Sb, (u32)(M + 1),
S, A, B, ctr, sc2);
untranspose(E2, S, A);
{
long k = M, e = M + M + 1;
u32 *d = c + N1_;
while (k < e) {
long col = k % (long)NC2_;
long run = (long)NC2_ - col;
if (run > e - k) run = e - k;
const u32 *s = A + (k / (long)NC2_) * CSX2_ + col;
if ((d - c) % 8 == 0 && col % 8 == 0 && run % 8 == 0 && run > 0) copy_u32_nt(d, s, run);
else for (long i = 0; i < run; i++) d[i] = s[i];
d += run; k += run;
}
}
}
g_inv_scale = sc;
multiply_rows(E1, a, (u32)n + 1, b, (u32)m + 1,
S, A, B, ctr, sc);
untranspose(E1, S, A);
long nmax = (long)n + (long)m;
#if EXACT
{
const long kA = ((long)n < (long)m) ? (long)n : (long)m;
const long kB = ((long)n < (long)m) ? (long)m : (long)n;
const int nm_lt = ((long)n < (long)m);
const long M1 = (M + 1 > 0) ? (M + 1) : 0;
u32 SUM = 0;
cmbv<0,1,1>(0, (M < kA) ? M : kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<0,2,1>(M1, (N1_ - 1 < kA) ? (long)(N1_ - 1) : kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<0,3,0>((long)N1_, kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<1,2,1>((kA + 1 > M1) ? kA + 1 : M1, (N1_ - 1 < kB) ? (long)(N1_ - 1) : kB,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<1,3,0>((kA + 1 > (long)N1_) ? kA + 1 : (long)N1_, kB,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<2,2,1>((kB + 1 > M1) ? kB + 1 : M1, (N1_ - 1 < nmax) ? (long)(N1_ - 1) : nmax,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<2,3,0>((kB + 1 > (long)N1_) ? kB + 1 : (long)N1_, nmax,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
}
#else
for (long k = 0; k <= nmax; k++) c[k] = c[k];
#endif
_mm_sfence();
(void)E2;
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 5.748 s | 2046 MB + 480 KB | Accepted | Score: 100 | 显示更多 |