FPS の定数係数畳み込み

components/fps_convolution_constant.hpp · GitHub で見る

共通ページ

Linux x86_64 / GCC の AVX2 用内部部品です。公開の fps::multiply と convolution が、両入力の全係数がそれぞれ一定の場合に利用します。

操作	内容	計算量
fps_detail::multiply_detail::constant::fits(a,b)	両入力が 4096 項以上で、各列の全係数がそれぞれ同じ剰余か厳密に判定	最悪 \(O(n+m)\)
fps_detail::multiply_detail::constant::multiply(a,b)	確認済みの二つの定数係数列の積を返す	\(O(n+m)\)

fits は空入力や短い入力には false を返します。最初に両列の先頭 8 項を比較してから残りを調べるため、先頭で一致しない入力は \(O(1)\) で棄却します。遅延剰余を正規化して比較し、末尾まで確認します。

multiply の呼び出し側は、AVX2 対応、両 span が生存中の Mint 要素を指すこと、各列がそれぞれ定数であること、非空入力、内部 Montgomery 値が \([0,2p)\)、出力長 \(n+m-1\) が \(2^{23}\) 以下であることを保証します。公開 dispatch は fits が true のときだけこの関数を呼びます。

各列の係数を \(a_0,b_0\) とすると、出力の第 \(i\) 係数は \(a_0b_0\min(i+1,n,m,n+m-1-i)\) です。正規化した Montgomery 表現の積から等差列を生成し、反対側へ反転して書き、中央の一定区間を埋めます。零の場合も公開サイズを保ち、返すすべての係数を \([0,p)\) に正規化します。

入力を変更せず、alignment や padding を要求しません。同一入力や長短の交換にも対応します。結果は独立した vector<Mint> の生存中の要素として構築し、SIMD による入力・出力の受け渡しには memcpy を使います。返す配列を除く追加作業領域は \(O(1)\)、返す配列は \(O(n+m)\) です。

依存ライブラリ

• FPS のブロック変換
• FPS の結果配列への一括変換
• FPS の作業領域
• modint998244353 の共通剰余演算
• modint998244353

検証コード

• 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/ordered_convolution.test.cpp
• test/polynomial_inv_mod.test.cpp
• test/power_projection.test.cpp

コード

#ifndef CP_COMPONENTS_FPS_CONVOLUTION_CONSTANT_HPP
#define CP_COMPONENTS_FPS_CONVOLUTION_CONSTANT_HPP

#include "fps_blocked.hpp"
#include <algorithm>
#include <bit>
#include <cstddef>
#include <cstring>
#include <span>
#include <vector>

#if defined(__linux__) && defined(__x86_64__) && defined(__GNUC__) && !defined(__clang__)
#pragma GCC push_options
#pragma GCC optimize("O3")
_Pragma("GCC target(\"avx2\")")
namespace fps_detail::multiply_detail::constant {
inline unsigned canonical(mint x) { return std::min(x.v, x.v - mint::mod()); }
inline bool same_eight(const mint* p, unsigned value) {
	__m256i v;
	std::memcpy(&v, p, sizeof(v));
	v = _mm256_min_epu32(v, _mm256_sub_epi32(v, _mm256_set1_epi32(mint::mod())));
	return _mm256_movemask_epi8(_mm256_cmpeq_epi32(v, _mm256_set1_epi32(value))) == -1;
}
inline bool remaining(std::span<const mint> a, unsigned value) {
	size_t i = 8;
	for (; i + 8 <= a.size(); i += 8) if (!same_eight(a.data() + i, value)) return false;
	for (; i < a.size(); ++i) if (canonical(a[i]) != value) return false;
	return true;
}
[[gnu::noinline]] inline bool fits(const std::vector<mint>& a, const std::vector<mint>& b) {
	if (std::min(a.size(), b.size()) < 4096) return false;
	const unsigned av = canonical(a[0]), bv = canonical(b[0]);
	if (!same_eight(a.data(), av) || !same_eight(b.data(), bv)) return false;
	return remaining(a, av) && (&a == &b || remaining(b, bv));
}
[[gnu::noinline]] inline std::vector<mint> multiply(std::span<const mint> a, std::span<const mint> b) {
	const size_t n = std::max(a.size(), b.size()), m = std::min(a.size(), b.size()), size = n + m - 1;
	const unsigned c = canonical(mint::word(canonical(a[0])) * mint::word(canonical(b[0])));
	std::vector<mint> result(size);
	if (c == 0) return result;
	alignas(32) unsigned first[8];
	for (unsigned i = 0; i < 8; ++i) first[i] = std::uint64_t(c) * (i + 1) % mint::mod();
	__m256i value = _mm256_load_si256(reinterpret_cast<const __m256i*>(first));
	const auto step = _mm256_set1_epi32(std::uint64_t(c) * 8 % mint::mod()), p = _mm256_set1_epi32(mint::mod());
	const auto reverse = _mm256_setr_epi32(7, 6, 5, 4, 3, 2, 1, 0);
	size_t i = 0;
	for (; i + 8 <= m; i += 8) {
		std::memcpy(static_cast<void*>(result.data() + i), &value, sizeof(value));
		const auto reflected = _mm256_permutevar8x32_epi32(value, reverse);
		std::memcpy(static_cast<void*>(result.data() + size - i - 8), &reflected, sizeof(reflected));
		value = _mm256_add_epi32(value, step);
		value = _mm256_min_epu32(value, _mm256_sub_epi32(value, p));
	}
	unsigned current = _mm256_extract_epi32(value, 0);
	for (; i < m; ++i) {
		result[i] = result[size - i - 1] = mint::word(current);
		current += c;
		current = std::min(current, current - mint::mod());
	}
	if (n > m + 1) std::fill(result.begin() + m, result.begin() + n - 1, mint::word(std::uint64_t(c) * m % mint::mod()));
	return result;
}
}
#pragma GCC pop_options
#endif
#endif

Back to top page