FPS の融合畳み込み

components/fps_convolution.hpp · GitHub で見る

共通ページ · 軽量実装なし

Linux x86_64 / GCC の fps::multiply と convolution(vector<mint>,vector<mint>) が使う内部部品です。公開 API から利用し、呼び出し側が AVX2 対応を実行時に確認します。

操作	内容・返り値	計算量
fps_detail::multiply_detail::multiply(a,b)	内部 Montgomery 係数の二つの非空配列の積を返す	\(O(n\log n)\)

通常は FPS の f*g を使います。ここで n は入出力の最大係数数です。

fps_detail::multiply_detail::multiply(a,b) は、非空の Montgomery 係数列から積を返します。出力長は \(2^{23}\) 以下とし、入力を変更しません。同一配列の二乗にも対応します。入力に alignment や padding は要求しません。

両入力が 4096 項以上で、それぞれの列の全係数が一定であることを正規化した値で厳密に確認できる場合は、定数係数畳み込み を使います。重なりの項数から積を直接作り、時間は \(O(n+m)\)、返す配列を除く追加作業領域は \(O(1)\) です。入力の先頭だけで判定を確定せず、末尾も確認します。

短辺が 32 より大きく、長辺が 32768 以上かつ短辺の 32 倍以上のときは、非対称畳み込み を使います。短辺の変換を再利用して長辺をブロックごとに処理し、短辺を \(m\)、長辺を \(n\) として時間は \(O((n+m)\log m)\)、追加作業領域は \(O(m)\) です。

次数 8 の環を葉とし、2 本の前進変換・葉の積・逆変換を部分木内で続けて処理します。条件を満たす長さでは最下段を 2 群ずつ処理し、半分が零の最上段と最後の倍率を融合します。入力の共通因子 \(R=2^{32}\) を保持し、結果を直接 Montgomery 表現で返します。

Subtrees・forward_radix8 は Taylor shift の middle product でも共有します。作業配列・root 表は RAII で管理し、assembly が読む末尾の余白を含めて初期化します。

専用 Product と汎用 ntt::Convolution のどちらでも、返す vector<Mint> 内に変換長と末尾 16 係数の領域を確保し、32 byte 境界から作業配列 A として使います。全係数の生存期間を開始して零で埋めてから入力をコピーし、変換後は先頭へ移動して公開サイズに縮めます。別の結果配列へのコピーを省き、返す配列の容量には変換用の余白が残ることがあります。

内部の Product(n,m,borrowed_a) と ntt::Convolution(n,m,borrowed_a) は A を借用し、B と root 表だけを所有します。borrowed_a は 32 byte alignment、length()+16 個以上の生存中の 4 byte 係数、零初期化済みの末尾を必要とします。呼び出し側が入力を正規化してコピーし、plan の利用中は領域を維持します。借用引数のない二つのコンストラクタは従来どおり A も所有し、借用のない raw 配列用の経路を維持します。

借用領域への入力コピーは Mint と memcpy を使います。変換中の SIMD アクセスは GCC の alias 対応 vector 型を使い、公開入力の const 領域には書き込みません。最終係数は正規化した Montgomery 表現です。

出力長を \(n\) として、時間は \(O(n\log n)\)、追加空間は \(O(n)\) です。

依存ライブラリ

• FPS のブロック変換
• FPS の定数係数畳み込み
• 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_HPP
#define CP_COMPONENTS_FPS_CONVOLUTION_HPP
#include "fps_blocked.hpp"
#include "fps_output.hpp"
#include "fps_convolution_unbalanced.hpp"
#include "fps_convolution_constant.hpp"
#include <algorithm>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#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 {
namespace ntt {
namespace kernels = blocked_ntt::kernels;
namespace detail = blocked_ntt::detail;
class Convolution {
	size_t n_, m_, len_;
	blocked_poly::Arena arena_;
	blocked_poly::Transform transform_;
	Word *a_, *b_;
	static int log_length(size_t n, size_t m) {
		if (n == 0 || m == 0 || n > (1u << 23) || m > (1u << 23) || n + m - 1 > (1u << 23)) std::abort();
		return std::max(6, int(std::bit_width(n + m - 2)));
	}
public:
	Convolution(size_t n, size_t m) : Convolution(n, m, nullptr) {}
	Convolution(size_t n, size_t m, Word* borrowed)
		: n_(n), m_(m), len_(size_t(1) << log_length(n, m)),
		  arena_(blocked_poly::Transform::words(log_length(n, m)) + (borrowed ? 1 : 2) * blocked_poly::Arena::footprint(len_)),
		  transform_(arena_, log_length(n, m)), a_(borrowed ? borrowed : arena_.take(len_).data()), b_(arena_.take(len_).data()) {
		if (n == 0 || m == 0 || log_length(n, m) > 23) std::abort();
	}
	static bool fits(size_t n, size_t m) {
		const int lg = log_length(n, m);
		return lg % 2 == 0 && lg >= 10 && 2 * std::max(n, m) <= size_t(1) << lg;
	}
	Word* a() { return a_; }
	Word* b() { return b_; }
	const Word* multiply(bool square = false) {
		transform_.forward(ConstWords(b_, m_), 0, Words(b_, len_));
		if (square) transform_.inverse_product(ConstWords(b_, len_), ConstWords(b_, len_), Words(a_, len_),
			blocked_poly::Half::kBoth, detail::power(detail::kR, kernels::kP - 2));
		else transform_.cyclic_product(ConstWords(a_, n_), 0,
			Words(a_, len_), ConstWords(b_, len_), blocked_poly::Half::kBoth, detail::power(detail::kR, kernels::kP - 2));
		return a_;
	}
};
}
using Product = ntt::Convolution;
using ntt::detail::Vec, ntt::detail::Factor, ntt::detail::add, ntt::detail::diff, ntt::detail::kP, ntt::detail::reduce;
inline Vec times(Vec x, const Factor& f) { return ntt::detail::multiply(x, f); }
inline Vec diff_canonical(Vec x, Vec y) { return _mm256_sub_epi32(_mm256_add_epi32(x, ntt::detail::broadcast(kP)), y); }
class Subtrees {
	const Word *roots_, *inverse_;
public:
	Subtrees(const Word* r, const Word* ir) : roots_(r), inverse_(ir) {}
	void visit(Vec* a, Vec* b, size_t nv, size_t k) const {
		using namespace blocked_poly::detail;
		Recursion<ForwardBottom>(roots_, inverse_, ForwardBottom{roots_}).visit(reinterpret_cast<Word*>(b), nv, k);
		ProductBottom bottom{roots_, inverse_, reinterpret_cast<Word*>(b), k * nv};
		Recursion<ProductBottom>(roots_, inverse_, bottom).visit(reinterpret_cast<Word*>(a), nv, k);
	}
};
inline void forward_radix8(Vec* f, size_t q, const Word* roots) {
	const Factor i(roots[1], roots[9]), y(roots[2], roots[10]), z(roots[3], roots[11]);
	for (size_t j = 0; j < q; ++j) {
		const Vec f0 = f[j], f1 = f[j + q], f2 = f[j + 2 * q], f3 = f[j + 3 * q];

		const Vec g0 = add(f0, f2), g1 = add(f1, f3);
		const Vec h0 = diff_canonical(f0, f2), ih1 = times(diff_canonical(f1, f3), i);
		f[j] = add(g0, g1), f[j + q] = diff(g0, g1);
		f[j + 2 * q] = add(h0, ih1), f[j + 3 * q] = diff(h0, ih1);

		const Vec if2 = times(f2, i), if3 = times(f3, i);
		const Vec u0 = reduce(add(f0, if2), 2 * kP), v0 = reduce(diff(f0, if2), 2 * kP);
		const Vec yu1 = times(add(f1, if3), y), zv1 = times(diff(f1, if3), z);
		f[j + 4 * q] = add(u0, yu1), f[j + 5 * q] = diff(u0, yu1);
		f[j + 6 * q] = add(v0, zv1), f[j + 7 * q] = diff(v0, zv1);
	}
}
inline void copy_input(const std::vector<mint> &in, Word *out, size_t len) {
	size_t i = 0;
	const auto p = _mm256_set1_epi32(mint::mod());
	for (; i + 8 <= in.size(); i += 8) {
		const auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(in.data() + i));
		_mm256_store_si256(reinterpret_cast<__m256i *>(out + i), _mm256_min_epu32(x, _mm256_sub_epi32(x, p)));
	}
	for (; i < in.size(); ++i) out[i] = std::min(in[i].v, in[i].v - mint::mod());
	std::fill(out + i, out + (i <= len / 2 ? len / 2 : len), 0);
	std::fill(out + len, out + len + 16, 0);
}


inline void copy_input_mint(const std::vector<mint>& in, mint* out) {
	static_assert(sizeof(mint) == sizeof(Word) && std::is_trivially_copyable_v<mint>);
	size_t i = 0;
	const auto p = _mm256_set1_epi32(mint::mod());
	for (; i + 8 <= in.size(); i += 8) {
		__m256i x;
		std::memcpy(&x, in.data() + i, sizeof(x));
		x = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
		std::memcpy(static_cast<void*>(out + i), &x, sizeof(x));
	}
	for (; i < in.size(); ++i) out[i] = mint::word(std::min(in[i].v, in[i].v - mint::mod()));
}

template <class Plan>
inline std::vector<mint> run_borrowed_as(const std::vector<mint>& a, const std::vector<mint>& b) {
	const size_t len = size_t(1) << std::max(6, int(std::bit_width(a.size() + b.size() - 2)));
	constexpr size_t alignment = 32;
	const size_t words = len + 16 + alignment / sizeof(mint) - 1, size = a.size() + b.size() - 1;
	std::vector<mint> result;
	prepare_fps_output(result, words);
	result.assign(words, mint::word(0));
	const auto address = reinterpret_cast<std::uintptr_t>(result.data());
	const size_t offset = ((alignment - (address & (alignment - 1))) & (alignment - 1)) / sizeof(mint);
	mint* const to = result.data() + offset;
	Plan plan(a.size(), b.size(), reinterpret_cast<Word*>(to));
	if (&a != &b) copy_input_mint(a, to);
	copy_input(b, plan.b(), len);
	plan.multiply(&a == &b);
	if (offset) std::memmove(static_cast<void*>(result.data()), to, size * sizeof(mint));
	result.resize(size);
	return result;
}

inline std::vector<mint> run_borrowed(const std::vector<mint>& a, const std::vector<mint>& b) {
	return run_borrowed_as<Product>(a, b);
}

template <class Plan>
inline std::vector<mint> run(const std::vector<mint> &a, const std::vector<mint> &b) {
	Plan plan(a.size(), b.size());
	const auto len = size_t(1) << std::max(6, int(std::bit_width(a.size() + b.size() - 2)));
	copy_input(a, plan.a(), len);
	copy_input(b, plan.b(), len);
	const auto *p = plan.multiply();
	std::vector<mint> result;
	const size_t size = a.size() + b.size() - 1;
	prepare_fps_output(result, size);
	append_fps_words<true>(result, p, size);
	return result;
}


inline std::vector<mint> multiply(const std::vector<mint>& a, const std::vector<mint>& b) {
	assert(!a.empty() && !b.empty() && a.size() + b.size() - 1 <= (1 << 23));
	if (constant::fits(a, b)) return constant::multiply(a, b);
	if (unbalanced::fits(a.size(), b.size())) return unbalanced::multiply(a, b);
	return run_borrowed_as<ntt::Convolution>(a, b);
}
}
#pragma GCC pop_options
#endif
#endif

Back to top page