modint998244353 の共通剰余演算

components/mint_arithmetic.hpp · GitHub で見る

共通ページ

mint 用の階乗生成・多項式評価・NTT が共有する内部演算です。通常は modint998244353 を利用します。NTT の変換処理や変換表には依存しません。

mint_detail::Ops は Montgomery 表現のワードを扱い、mul・add・sub、通常表現との変換 in・out、定数乗算 fixed_mul を提供します。one は Montgomery 表現の 1、mod は 998244353、twice はその 2 倍です。

x86 では vectorized() で AVX2 の利用可否を確認でき、load8・store8・mul8・reduce8・add8・sub8・fixed_mul8・fixed_even8 を提供します。AVX2 の関数は対応 CPU から呼び出します。

reduce8(even, odd) は、偶数・奇数位置ごとに 4 本の 64 bit レーンへ蓄積した積を Montgomery 還元し、8 本の 32 bit レーンへ戻す内部部品です。補正項を加えても 64 bit でオーバーフローしない範囲で使い、返されたワードの範囲の調整は呼び出し側で行います。

各演算は \(O(1)\) 時間・追加空間です。変換表の前計算や動的なメモリ確保は行いません。

依存ライブラリ

• modint998244353

検証コード

• test/binomial_prefix_sum.test.cpp
• test/binomial_prefix_sum_online.test.cpp
• test/bitwise_and_convolution.test.cpp
• test/bitwise_xor_convolution.test.cpp
• test/comb.test.cpp
• test/comb_prime_mod.test.cpp
• test/convolution.test.cpp
• test/cyclic_convolution.test.cpp
• test/fps_bostan_mori.test.cpp
• test/fps_compose.test.cpp
• test/fps_compose_large.test.cpp
• test/fps_compositional_inverse.test.cpp
• test/fps_compositional_inverse_large.test.cpp
• test/fps_divmod.test.cpp
• test/fps_exp.test.cpp
• test/fps_exp_sparse.test.cpp
• test/fps_interpolate.test.cpp
• test/fps_inv.test.cpp
• test/fps_inv_sparse.test.cpp
• test/fps_linear_recurrence.test.cpp
• test/fps_log.test.cpp
• test/fps_log_sparse.test.cpp
• test/fps_pow.test.cpp
• test/fps_pow_sparse.test.cpp
• test/fps_rational_coefficients.test.cpp
• test/fps_sqrt.test.cpp
• test/fps_sqrt_sparse.test.cpp
• test/fps_taylor_shift.test.cpp
• test/multiplicative_convolution.test.cpp
• test/multipoint_eval.test.cpp
• test/multivariate_convolution.test.cpp
• test/multivariate_convolution_light.test.cpp
• test/ntt.test.cpp
• test/online_convolution.test.cpp
• test/online_convolution_exp.test.cpp
• test/online_convolution_inv.test.cpp
• test/ordered_convolution.test.cpp
• test/polynomial_inv_mod.test.cpp
• test/power_projection.test.cpp

コード

#ifndef MINT_ARITHMETIC_HPP
#define MINT_ARITHMETIC_HPP

#include "../math/mint.hpp"
#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#endif

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC push_options
#pragma GCC optimize("O3")
#endif

namespace mint_detail {
struct Ops {
	using T = mint;
	typedef unsigned int U __attribute__((may_alias));
	static constexpr U mod = 998244353, twice = 2 * mod, one = 301989884;
	static U reduce(unsigned long long x) { return T::reduce(x); }
	static U mul(U a, U b) { return reduce((unsigned long long)a * b); }
	static U add(U a, U b) {
		U c = a + b;
		return c < twice ? c : c - twice;
	}
	static U sub(U a, U b) {
		U c = a - b;
		return c < twice ? c : c + twice;
	}
	static U in(U a) { return mul(a, 932051910); }
	static U out(U a) {
		a = reduce(a);
		return a < mod ? a : a - mod;
	}
	static U quotient(U r) { return ((unsigned long long)r << 32) / mod; }
	static U fixed_mul(U a, U r, U q) {
		return a * r - U(((unsigned long long)a * q) >> 32) * mod;
	}

#if defined(__x86_64__) || defined(__i386__)
	static bool vectorized() {
#if defined(__AVX2__)
		return true;
#else
		static const bool available = __builtin_cpu_supports("avx2");
		return available;
#endif
	}
	using I8 = __m256i;
	__attribute__((target("avx2"))) static I8 load8(const U *a) {
		return _mm256_loadu_si256(reinterpret_cast<const I8 *>(a));
	}
	__attribute__((target("avx2"))) static void store8(U *a, I8 x) {
		_mm256_storeu_si256(reinterpret_cast<I8 *>(a), x);
	}
	__attribute__((target("avx2"), always_inline)) static inline I8 reduce8(I8 x, I8 y) {
		I8 r = _mm256_set1_epi32(998244351), p = _mm256_set1_epi32(mod);
		x = _mm256_add_epi64(x, _mm256_mul_epu32(_mm256_mul_epu32(x, r), p));
		y = _mm256_add_epi64(y, _mm256_mul_epu32(_mm256_mul_epu32(y, r), p));
		return _mm256_blend_epi32(_mm256_srli_epi64(x, 32), y, 0xaa);
	}
	__attribute__((target("avx2"))) static I8 mul8(I8 a, I8 b) {
		return reduce8(_mm256_mul_epu32(a, b), _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32)));
	}
	__attribute__((target("avx2"))) static I8 fixed_mul8(I8 a, I8 r, I8 q) {
		I8 x = _mm256_srli_epi64(_mm256_mul_epu32(a, q), 32);
		I8 y = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), q);
		I8 t = _mm256_blend_epi32(x, y, 0xaa);
		return _mm256_sub_epi32(_mm256_mullo_epi32(a, r),
								_mm256_mullo_epi32(t, _mm256_set1_epi32(mod)));
	}
	__attribute__((target("avx2"))) static I8 fixed_even8(I8 a, I8 r, I8 q) {
		I8 t = _mm256_srli_epi64(_mm256_mul_epu32(a, q), 32);
		return _mm256_sub_epi64(_mm256_mul_epu32(a, r),
								_mm256_mul_epu32(t, _mm256_set1_epi32(mod)));
	}
	__attribute__((target("avx2"))) static I8 add8(I8 a, I8 b) {
		I8 x = _mm256_add_epi32(a, b);
		return _mm256_min_epu32(x, _mm256_sub_epi32(x, _mm256_set1_epi32(twice)));
	}
	__attribute__((target("avx2"))) static I8 sub8(I8 a, I8 b) {
		I8 x = _mm256_sub_epi32(a, b);
		return _mm256_min_epu32(x, _mm256_add_epi32(x, _mm256_set1_epi32(twice)));
	}

#endif
};
}

#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
#pragma GCC pop_options
#endif

#endif

Back to top page