FPS の対数の内部変換

components/fps_log.hpp · GitHub で見る

共通ページ

fps_detail::LogTransform は fps/fps.hpp の対数と冪が使う x86 / AVX2 用の内部部品です。通常は fps::log() から利用します。AVX2 が使えない環境では FPS 側が汎用の級数除算へ戻ります。

#include "lib/components/fps_log.hpp"

#if defined(__x86_64__) || defined(__i386__)
if (fps_detail::NTT::vectorized()) {
	std::vector<fps_detail::NTT::U> data(64, 0);
	auto &plan = fps_detail::ntt_plan(64 / 16);
	fps_detail::LogTransform transform{plan};
	transform.forward(data.data(), 64);
}
#endif

操作	内容・返り値	計算量
LogTransform{plan}	前処理済みの NTT プランを参照して作成	\(O(1)\)
forward(a,n,norm = 1,half = false)	a を部分 NTT し、norm 倍を組み込む	\(O(n\log n)\)
inverse(a,n)	a を未正規化で逆変換	\(O(n\log n)\)
inverse_half<upper>(a,n)	指定した半分だけを逆変換	\(O(n\log n)\)
products(dst,a,b,terms,n,side = nullptr)	terms 組の変換済み配列の環積の和を dst に格納	\(O(n)\)
toeplitz(dst,side,w1,w2,q0,q1,n)	二つの環積の和を dst・side に格納	\(O(n)\)
integrate(dst,src,prefix,first,n)	src[i] / (first+i) を dst[i] に格納	\(O(n)\)

この型は x86 でのみ定義されます。各関数を呼ぶ前に AVX2 対応を確認します。保持する NTT プランの寿命は、この型より長くします。

forward(a,n,norm,half) は次数 8 の環を葉とする部分 NTT を行い、通常の整数 norm 倍を組み込みます。norm は \([0,p)\)、入力は \([0,2p)\) とします。結果は \([0,p)\) の内部 Montgomery ワードです。half=true は後半を 0 として扱い、読み取る入力は前半のみです。inverse(a,n) は未正規化の逆変換で、両半分を復元します。比較用の inverse_half<upper>(a,n) は、upper が真なら後半、偽なら前半だけを復元します。正規化係数は \(8/n\) です。入力は \([0,2p)\)、復元した半分も \([0,2p)\) です。もう半分の値は未規定です。

変換長 n は 64 以上 \(2^{23}\) 以下の 2 冪とし、配列には n 要素、プランには少なくとも n/16 個の根を用意します。

products(dst,a,b,terms,n,side) は同じ部分変換の 1 または 2 組の環積を加算します。side が非 null なら、既存の環積も加えます。a・b は terms 個の配列ポインタです。各入力は n 要素で、a・b は \([0,p)\)、side は \([0,2p)\) とします。結果は \([0,2p)\) です。dst は入力の 1 配列と同一でも構いません。それ以外の部分的な重なりは認めません。出力以外の入力は変更しません。

toeplitz(dst,side,w1,w2,q0,q1,n) は、環積 \(w_2q_0+w_1q_1\) を dst に、\(w_3q_0+w_2q_1\) を side に計算します。side は最初に \(w_3\) を保持します。全入力は \([0,p)\)、出力は \([0,2p)\) です。dst は q1 と同一でも構いません。他の配列同士は重ねません。両積を 3 回の環積で求めます。

products・toeplitz の n は 64 の倍数とし、根を n/16 個以上用意します。

integrate(dst,src,prefix,first,n) は dst[i] = src[i] / (first+i) を計算します。src は \([0,2p)\) の内部 Montgomery ワードで、出力も同範囲です。first >= 1、n >= 0、first+n <= p とします。src・dst はそれぞれちょうど n 要素で構いません。末尾はマスク付きの読み書きを使います。prefix は n を 32 の倍数へ切り上げた要素数の作業領域とし、3 配列は重ねません。ここで \(p=998244353\) です。

計算量: forward・inverse・inverse_half は \(O(n\log n)\)、products・toeplitz・integrate は \(O(n)\) 時間です。integrate の作業領域は \(O(n)\)、その他の追加空間は \(O(1)\) です。NTT プランは \(O(n)\) 空間です。

依存ライブラリ

• modint998244353 の共通剰余演算
• NTT（998244353）（高速版）
• modint998244353

コード

#ifndef CP_COMPONENTS_FPS_LOG_HPP
#define CP_COMPONENTS_FPS_LOG_HPP

#include "ntt.hpp"

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

namespace fps_detail {

#if defined(__x86_64__) || defined(__i386__)
struct LogTransform {
	using U = NTT::U;
	using I8 = NTT::I8;
	static constexpr U mod = NTT::mod, twice = NTT::twice, one = NTT::one;
	const NTT &plan;

	__attribute__((target("avx2"))) void forward(U *a, int n, U norm = 1, bool half = false) const {
		assert(n >= 64 && !(n & (n - 1)) && n <= (1 << 23) && static_cast<int>(plan.roots.size()) >= n / 16);
		if (half) std::copy_n(a, n / 2, a + n / 2);
		plan.template inverse_partial8<false>(a, n, half ? n / 2 : n);
		I8 p = _mm256_set1_epi32(mod), pp = _mm256_set1_epi32(twice);
		I8 r = _mm256_set1_epi32(norm), rq = _mm256_set1_epi32(NTT::quotient(norm));
		for (int i = 0; i < n; i += 8) {
			I8 x = NTT::load8(a + i);
			x = _mm256_min_epu32(x, _mm256_sub_epi32(x, pp));
			if (norm != 1) x = NTT::fixed_mul8(x, r, rq);
			NTT::store8(a + i, _mm256_min_epu32(x, _mm256_sub_epi32(x, p)));
		}
	}
	__attribute__((target("avx2"))) void inverse(U *a, int n) const {
		assert(n >= 64 && !(n & (n - 1)) && n <= (1 << 23) && static_cast<int>(plan.roots.size()) >= n / 16);
		plan.template inverse_partial8<true>(a, n);
	}
	template <bool upper> __attribute__((target("avx2"))) void inverse_half(U *a, int n) const {
		assert(n >= 64 && !(n & (n - 1)) && n <= (1 << 23) && static_cast<int>(plan.roots.size()) >= n / 16);
		int h = 8;
		for (; 4 * h < n; h <<= 2) plan.template radix4_8<true>(a, h, n);
		if (2 * h == n) {
			for (int i = 0; i < h; i += 8) {
				I8 x = upper ? NTT::sub8(NTT::load8(a + i), NTT::load8(a + h + i)) : NTT::add8(NTT::load8(a + i), NTT::load8(a + h + i));
				NTT::store8(a + (upper ? h : 0) + i, x);
			}
		} else {
			I8 im = _mm256_set1_epi32(plan.inverse_roots[1]), iq = _mm256_set1_epi32(plan.inverse_root_quotients[1]);
			for (int i = 0; i < h; i += 8) {
				I8 x0 = NTT::load8(a + i), x1 = NTT::load8(a + h + i), x2 = NTT::load8(a + 2 * h + i), x3 = NTT::load8(a + 3 * h + i);
				I8 a0 = NTT::add8(x0, x1), a1 = NTT::sub8(x0, x1), a2 = NTT::add8(x2, x3), a3 = NTT::fixed_mul8(NTT::sub8(x2, x3), im, iq);
				NTT::store8(a + (upper ? 2 * h : 0) + i, upper ? NTT::sub8(a0, a2) : NTT::add8(a0, a2));
				NTT::store8(a + (upper ? 3 * h : h) + i, upper ? NTT::sub8(a1, a3) : NTT::add8(a1, a3));
			}
		}
	}
	__attribute__((target("avx2"))) void products(U *dst, const U *const *a, const U *const *b, int terms, int n, const U *side = nullptr) const {
		assert(n >= 64 && n % 64 == 0 && terms >= 1 && terms <= 2 && static_cast<int>(plan.roots.size()) >= n / 16);
		for (int i = 0; i < n; i += 64) {
			U ws[8], qs[8];
			for (int j = 0; j < 8; ++j) {
				int k = i / 8 + j;
				ws[j] = k & 1 ? mod - plan.roots[k / 2] : plan.roots[k / 2];
				qs[j] = k & 1 ? ~plan.root_quotients[k / 2] : plan.root_quotients[k / 2];
			}
			alignas(32) U x[64], sum[64];
			for (int t = 0; t < terms; ++t) {
				std::copy_n(a[t] + i, 64, x);
				plan.template inverse_leaf_products8<8, true>(x, const_cast<U *>(b[t]) + i, ws, qs);
				for (int j = 0; j < 64; j += 8) NTT::store8(sum + j, t ? NTT::add8(NTT::load8(sum + j), NTT::load8(x + j)) : NTT::load8(x + j));
			}
			for (int j = 0; j < 64; j += 8)
				NTT::store8(dst + i + j, side ? NTT::add8(NTT::load8(sum + j), NTT::load8(side + i + j)) : NTT::load8(sum + j));
		}
	}

	__attribute__((target("avx2"))) static void integrate(U *dst, const U *src, U *prefix, int first, int n) {
		assert(n >= 0 && first >= 1 && static_cast<unsigned long long>(first) + n <= mod);
		if (!n) return;
		I8 x[4], prod[4], inv[4], step = _mm256_set1_epi32(NTT::in(32)), onev = _mm256_set1_epi32(one);
		alignas(32) U initial[32];
		for (int i = 0; i < 32; ++i) initial[i] = NTT::in(first + i);
		const I8 lanes = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7), limit = _mm256_set1_epi32(n);
		auto mask = [&](int at) __attribute__((target("avx2"))) { return _mm256_cmpgt_epi32(limit, _mm256_add_epi32(_mm256_set1_epi32(at), lanes)); };
		auto factor = [&](I8 a, int at) __attribute__((target("avx2"))) { return at + 8 <= n ? a : _mm256_blendv_epi8(onev, a, mask(at)); };
		for (int v = 0; v < 4; ++v) x[v] = NTT::load8(initial + 8 * v), prod[v] = factor(x[v], 8 * v);
		for (int i = 32; i < n; i += 32)
			for (int v = 0; v < 4; ++v) {
				NTT::store8(prefix + i - 32 + 8 * v, prod[v]);
				x[v] = NTT::add8(x[v], step);
				prod[v] = NTT::mul8(prod[v], factor(x[v], i + 8 * v));
			}
		for (int v = 0; v < 4; ++v) inv[v] = NTT::inverse8(prod[v]);
		auto put = [&](int at, I8 reciprocal) __attribute__((target("avx2"))) {
			if (at >= n) return;
			I8 a = at + 8 <= n ? NTT::load8(src + at) : _mm256_maskload_epi32(reinterpret_cast<const int *>(src + at), mask(at));
			I8 r = NTT::mul8(a, reciprocal);
			if (at + 8 <= n)
				NTT::store8(dst + at, r);
			else
				_mm256_maskstore_epi32(reinterpret_cast<int *>(dst + at), mask(at), r);
		};
		for (int i = (n - 1) / 32 * 32; i >= 32; i -= 32)
			for (int v = 0; v < 4; ++v) {
				int at = i + 8 * v;
				if (at < n) put(at, NTT::mul8(inv[v], NTT::load8(prefix + at - 32)));
				inv[v] = NTT::mul8(inv[v], factor(x[v], at));
				x[v] = NTT::sub8(x[v], step);
			}
		for (int v = 0; v < 4; ++v) put(8 * v, inv[v]);
	}
	__attribute__((target("avx2"))) void toeplitz(U *dst, U *side, const U *w1, const U *w2, const U *q0, const U *q1, int n) const {
		assert(n >= 64 && n % 64 == 0 && static_cast<int>(plan.roots.size()) >= n / 16);
		I8 p = _mm256_set1_epi32(mod);
		auto canonical = [&](I8 x) __attribute__((target("avx2"))) { return _mm256_min_epu32(x, _mm256_sub_epi32(x, p)); };
		for (int i = 0; i < n; i += 64) {
			alignas(32) U x[64], y[64], z[64];
			U ws[8], qs[8];
			for (int j = 0; j < 8; ++j) {
				int k = i / 8 + j;
				ws[j] = k & 1 ? mod - plan.roots[k / 2] : plan.roots[k / 2];
				qs[j] = k & 1 ? ~plan.root_quotients[k / 2] : plan.root_quotients[k / 2];
				NTT::store8(x + 8 * j, canonical(_mm256_add_epi32(NTT::load8(q0 + i + 8 * j), NTT::load8(q1 + i + 8 * j))));
				NTT::store8(y + 8 * j, canonical(NTT::sub8(NTT::load8(w1 + i + 8 * j), NTT::load8(w2 + i + 8 * j))));
				NTT::store8(z + 8 * j, canonical(NTT::sub8(NTT::load8(side + i + 8 * j), NTT::load8(w2 + i + 8 * j))));
			}
			plan.template inverse_leaf_products8<8, true>(x, const_cast<U *>(w2) + i, ws, qs);
			plan.template inverse_leaf_products8<8, true>(y, const_cast<U *>(q1) + i, ws, qs);
			plan.template inverse_leaf_products8<8, true>(z, const_cast<U *>(q0) + i, ws, qs);
			for (int j = 0; j < 64; j += 8) {
				NTT::store8(dst + i + j, NTT::add8(NTT::load8(x + j), NTT::load8(y + j)));
				NTT::store8(side + i + j, NTT::add8(NTT::load8(x + j), NTT::load8(z + j)));
			}
		}
	}
};
#endif

}

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

#endif

Back to top page