#pragma GCC target("avx,avx2")
#include <bits/extc++.h>
#include <immintrin.h>
static const unsigned sModulus = 998244353U;
static const unsigned sInverse = 998244351U;
static const unsigned sRsquare = 932051910U;
static const __m256i vModulus = _mm256_set1_epi32(sModulus);
static const __m256i vInverse = _mm256_set1_epi32(sInverse);
static const __m256i vRsquare = _mm256_set1_epi32(sRsquare);
static const unsigned sSize = sizeof(uint32_t);
static const unsigned vSize = sizeof(__m256i) / sSize;
static auto shrink = [](auto V) -> __m256i {
return _mm256_min_epu32(V, _mm256_sub_epi32(V, vModulus));
};
static auto dilate = [](auto V) -> __m256i {
return _mm256_min_epu32(V, _mm256_add_epi32(V, vModulus));
};
static auto place = [](auto& self, auto u, auto v, auto n, auto N, auto o, auto O) -> void {
if (n == vSize) {
for (auto i = 0U; i != vSize; ++i) {
auto value = _mm256_load_si256(reinterpret_cast<__m256i*>(o + (v + i) * N + u));
_mm256_store_si256(reinterpret_cast<__m256i*>(O + i * vSize), value);
}
} else {
auto m = n / 2;
auto M = m * m;
self(self, u + 0 * m, v + 0 * m, m, N, o, O + 0 * M);
self(self, u + 1 * m, v + 0 * m, m, N, o, O + 1 * M);
self(self, u + 0 * m, v + 1 * m, m, N, o, O + 2 * M);
self(self, u + 1 * m, v + 1 * m, m, N, o, O + 3 * M);
}
};
static auto antiplace = [](auto& self, auto u, auto v, auto n, auto N, auto o, auto O) -> void {
if (n == vSize) {
for (auto i = 0U; i != vSize; ++i) {
auto value = _mm256_load_si256(reinterpret_cast<__m256i*>(O + i * vSize));
auto P = _mm256_mul_epu32(value, vRsquare);
auto Q = _mm256_mul_epu32(_mm256_srli_epi64(value, 32), vRsquare);
auto U = _mm256_mul_epu32(_mm256_mul_epu32(P, vInverse), vModulus);
auto V = _mm256_mul_epu32(_mm256_mul_epu32(Q, vInverse), vModulus);
auto X = _mm256_srli_epi64(_mm256_add_epi64(U, P), 32);
auto Y = _mm256_add_epi64(V, Q);
auto T = shrink(_mm256_blend_epi32(X, Y, 0xAA));
_mm256_store_si256(reinterpret_cast<__m256i*>(o + (v + i) * N + u), T);
}
} else {
auto m = n / 2;
auto M = m * m;
self(self, u + 0 * m, v + 0 * m, m, N, o, O + 0 * M);
self(self, u + 1 * m, v + 0 * m, m, N, o, O + 1 * M);
self(self, u + 0 * m, v + 1 * m, m, N, o, O + 2 * M);
self(self, u + 1 * m, v + 1 * m, m, N, o, O + 3 * M);
}
};
static auto multiplies = [](auto& self, auto A, auto B, auto C, auto n, auto P, auto Q, auto R) -> void {
if (n == vSize) {
__m256i s[vSize]{};
__m256i S[vSize]{};
for (auto w = 0U; w != vSize; ++w) {
auto o = _mm256_load_si256(reinterpret_cast<__m256i*>(B + w * vSize));
auto O = _mm256_srli_epi64(o, 32);
for (auto i = 0U; i != vSize; ++i) {
s[i] = _mm256_add_epi64(s[i], _mm256_mul_epu32(_mm256_set1_epi32(int(A[i * vSize + w])), o));
S[i] = _mm256_add_epi64(S[i], _mm256_mul_epu32(_mm256_set1_epi32(int(A[i * vSize + w])), O));
}
}
for (auto i = 0U; i != vSize; ++i) {
auto U = _mm256_mul_epu32(_mm256_mul_epu32(s[i], vInverse), vModulus);
auto V = _mm256_mul_epu32(_mm256_mul_epu32(S[i], vInverse), vModulus);
auto X = _mm256_srli_epi64(_mm256_add_epi64(U, s[i]), 32);
auto Y = _mm256_add_epi64(V, S[i]);
auto T = shrink(_mm256_blend_epi32(X, Y, 0xAA));
T = shrink(_mm256_min_epu32(T, _mm256_sub_epi32(T, _mm256_set1_epi32(sModulus * 2))));
_mm256_store_si256(reinterpret_cast<__m256i*>(C + i * vSize), T);
}
} else {
auto m = n / 2;
auto M = m * m;
auto plus = [M](auto U, auto V, auto W) -> void {
for (auto i = 0U; i != M; i += vSize) {
auto Ui = _mm256_load_si256(reinterpret_cast<__m256i*>(U + i));
auto Vi = _mm256_load_si256(reinterpret_cast<__m256i*>(V + i));
auto Wi = shrink(_mm256_add_epi32(Ui, Vi));
_mm256_store_si256(reinterpret_cast<__m256i*>(W + i), Wi);
}
};
auto minus = [M](auto U, auto V, auto W) -> void {
for (auto i = 0U; i != M; i += vSize) {
auto Ui = _mm256_load_si256(reinterpret_cast<__m256i*>(U + i));
auto Vi = _mm256_load_si256(reinterpret_cast<__m256i*>(V + i));
auto Wi = dilate(_mm256_sub_epi32(Ui, Vi));
_mm256_store_si256(reinterpret_cast<__m256i*>(W + i), Wi);
}
};
auto i00 = M * 0, i01 = M * 1;
auto i10 = M * 2, i11 = M * 3;
minus(A + i01, A + i11, P);
plus(B + i10, B + i11, Q);
self(self, P, Q, C + i00, m, P + M, Q + M, R);
plus(A + i00, A + i01, P);
self(self, P, B + i11, C + i01, m, P + M, Q + M, R);
minus(C + i00, C + i01, C + i00);
minus(B + i10, B + i00, Q);
self(self, A + i11, Q, C + i10, m, P + M, Q + M, R);
plus(C + i00, C + i10, C + i00);
minus(A + i10, A + i00, P);
plus(B + i00, B + i01, Q);
self(self, P, Q, C + i11, m, P + M, Q + M, R);
plus(A + i10, A + i11, P);
self(self, P, B + i00, R, m, P + M, Q + M, R + M);
plus(C + i10, R, C + i10);
minus(C + i11, R, C + i11);
plus(A + i00, A + i11, P);
plus(B + i00, B + i11, Q);
self(self, P, Q, R, m, P + M, Q + M, R + M);
plus(C + i00, R, C + i00);
plus(C + i11, R, C + i11);
minus(B + i01, B + i11, Q);
self(self, A + i00, Q, R, m, P + M, Q + M, R + M);
plus(C + i01, R, C + i01);
plus(C + i11, R, C + i11);
}
};
static unsigned bit_ceil_u(unsigned x) {
unsigned r = vSize;
while (r < x) r <<= 1;
return r;
}
void matrix_multiply(int n, const int *A, const int *B, int *C) {
static const unsigned sSize2 = sizeof(uint32_t);
auto N = bit_ceil_u((unsigned)n);
auto align = alignof(__m256i);
auto bytes = (size_t)N * N * sSize2;
auto alloc = [align](size_t sz) -> uint32_t* {
void *p = nullptr;
if (posix_memalign(&p, align, sz) != 0) return nullptr;
return static_cast<uint32_t*>(p);
};
auto dA = alloc(bytes);
auto dB = alloc(bytes);
auto dC = alloc(bytes);
std::memset(dA, 0, bytes);
std::memset(dB, 0, bytes);
std::memset(dC, 0, bytes);
for (auto i = 0U; i != (unsigned)n; ++i)
for (auto j = 0U; j != (unsigned)n; ++j)
dA[i * N + j] = (uint32_t)A[i * n + j];
for (auto i = 0U; i != (unsigned)n; ++i)
for (auto j = 0U; j != (unsigned)n; ++j)
dB[i * N + j] = (uint32_t)B[i * n + j];
auto oA = alloc(bytes);
auto oB = alloc(bytes);
auto oC = alloc(bytes);
auto P = alloc(bytes);
auto Q = alloc(bytes);
auto R = alloc(bytes);
place(place, 0U, 0U, N, N, dA, oA);
place(place, 0U, 0U, N, N, dB, oB);
multiplies(multiplies, oA, oB, oC, N, P, Q, R);
antiplace(antiplace, 0U, 0U, N, N, dC, oC);
for (auto i = 0U; i != (unsigned)n; ++i)
for (auto j = 0U; j != (unsigned)n; ++j)
C[i * n + j] = (int)dC[i * N + j];
free(dA); free(dB); free(dC);
free(oA); free(oB); free(oC);
free(P); free(Q); free(R);
}