体の内部演算

components/field_arithmetic.hpp · GitHub で見る

共通ページ

行列と線形漸化式の復元で共有する、体の演算と連続配列の計算です。

汎用の Field<T> は 体の基本演算 と共有しています。

操作	内容	計算量
field_detail::Field<T>	通常の型 T を使う演算	—
field_detail::ModularField<T>	固定法の modint を uint32_t の標準表現で扱う演算	—
Ops::add(a,b) / Ops::mul(a,b) / Ops::neg(a) / Ops::inv(a)	加算・乗算・符号反転・逆元	型 T の演算による
Ops::dot(a,b,n)	長さ n の内積	\(O(n)\)
Ops::scale(a,n,x)	a の長さ n の領域を x 倍	\(O(n)\)
Ops::axpy(a,b,n,x)	a の長さ n の領域に x*b を加える	\(O(n)\)

Ops::V が配列の要素型です。ModularField<T> は固定法が \(2\le p<2^{31}\) の modint 用で、val() と raw() に対応する型を使います。逆元を求める値は可逆である必要があります。

AVX2・AVX512 は対応する CPU でのみ使用し、非対応環境では通常の計算に切り替えます。FIELD_DISABLE_AVX2 を定義すると両方の SIMD 経路を無効にできます。

体の演算を \(O(1)\) として、配列の各操作は時間 \(O(n)\)、追加空間は \(O(1)\) です。

依存ライブラリ

• 体の基本演算

検証コード

• test/berlekamp_massey.test.cpp
• test/matrix_adjugate.test.cpp
• test/matrix_characteristic_polynomial.test.cpp
• test/matrix_det.test.cpp
• test/matrix_inverse.test.cpp
• test/matrix_pow.test.cpp
• test/matrix_product.test.cpp
• test/matrix_rank.test.cpp
• test/matrix_solve.test.cpp

コード

#ifndef FIELD_ARITHMETIC_HPP
#define FIELD_ARITHMETIC_HPP

#include "field_base.hpp"
#include <algorithm>
#include <cstdint>
#include <type_traits>
#include <utility>
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
#include <immintrin.h>
#endif

namespace field_detail {
template <class T, class = void>
struct SmallModint : std::false_type {};

template <class T>
struct SmallModint<T, std::void_t<decltype(std::integral_constant<std::uint32_t, T::mod()>{}),
	decltype(std::declval<T>().val()), decltype(T::raw(std::uint32_t{}))>>
	: std::bool_constant<(T::mod() > 1 && T::mod() < (std::uint32_t(1) << 31))
		&& std::is_convertible_v<decltype(std::declval<T>().val()), std::uint32_t>
		&& std::is_same_v<decltype(T::raw(std::uint32_t{})), T>> {};

template <class T>
inline constexpr bool small_modint = SmallModint<T>::value;

inline bool avx2() {
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__)) && !defined(FIELD_DISABLE_AVX2) && !defined(MATRIX_DISABLE_AVX2)
	static const bool enabled = __builtin_cpu_supports("avx2");
	return enabled;
#else
	return false;
#endif
}

template <class T>
struct ModularField {
	using V = std::uint32_t;
	static constexpr V mod = T::mod();
	static V add(V a, V b) { V c = a + b; return c < mod ? c : c - mod; }
	static V mul(V a, V b) { return std::uint64_t(a) * b % mod; }
	static V neg(V a) { return a ? mod - a : 0; }
	static V dot(const V *a, const V *b, int n) {
		constexpr int step = mod <= (V(1) << 30) ? 16 : 4;
		std::uint64_t sum = 0;
		for (int i = 0; i < n; i += step) {
			for (int j = i; j < std::min(n, i + step); ++j) sum += std::uint64_t(a[j]) * b[j];
			sum %= mod;
		}
		return sum;
	}
	static V inv(V a) { return (T(1) / T::raw(a)).val(); }
	static bool zero(V a, long double) { return !a; }
	static bool better(V, V) { return false; }
	static V product(V a, V x, V q) {
		V r = a * x - V((std::uint64_t(a) * q) >> 32) * mod;
		return r < mod ? r : r - mod;
	}
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
	__attribute__((target("avx2"))) static __m256i product8(__m256i a, __m256i x, __m256i q) {
		__m256i p = _mm256_set1_epi32(mod);
		__m256i even = _mm256_srli_epi64(_mm256_mul_epu32(a, q), 32);
		__m256i odd = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), q);
		__m256i high = _mm256_or_si256(even, _mm256_slli_epi64(_mm256_srli_epi64(odd, 32), 32));
		__m256i r = _mm256_sub_epi32(_mm256_mullo_epi32(a, x), _mm256_mullo_epi32(high, p));
		return _mm256_min_epu32(r, _mm256_sub_epi32(r, p));
	}
	__attribute__((target("avx512f"))) static __m512i product16(__m512i a, __m512i x, __m512i q) {
		const auto p = _mm512_set1_epi32(mod);
		auto even = _mm512_srli_epi64(_mm512_mul_epu32(a, q), 32);
		auto odd = _mm512_mul_epu32(_mm512_srli_epi64(a, 32), q);
		auto high = _mm512_mask_blend_epi32(0xaaaa, even, odd);
		auto r = _mm512_sub_epi32(_mm512_mullo_epi32(a, x), _mm512_mullo_epi32(high, p));
		return _mm512_min_epu32(r, _mm512_sub_epi32(r, p));
	}
	template <bool add_to>
	__attribute__((target("avx512f"))) static void apply16(V *a, const V *b, int n, V x, V q) {
		const auto vx = _mm512_set1_epi32(x), vq = _mm512_set1_epi32(q), p = _mm512_set1_epi32(mod);
		int i = 0;
		for (; i + 16 <= n; i += 16) {
			auto r = product16(_mm512_loadu_si512(b + i), vx, vq);
			if constexpr (add_to) {
				r = _mm512_add_epi32(r, _mm512_loadu_si512(a + i));
				r = _mm512_min_epu32(r, _mm512_sub_epi32(r, p));
			}
			_mm512_storeu_si512(a + i, r);
		}
		for (; i < n; ++i) {
			V r = product(b[i], x, q);
			if constexpr (add_to) a[i] = add(a[i], r);
			else a[i] = r;
		}
	}
	template <bool add_to>
	__attribute__((target("avx2"))) static void apply8(V *a, const V *b, int n, V x, V q) {
		__m256i vx = _mm256_set1_epi32(x), vq = _mm256_set1_epi32(q), p = _mm256_set1_epi32(mod);
		int i = 0;
		for (; i + 8 <= n; i += 8) {
			__m256i r = product8(_mm256_loadu_si256(reinterpret_cast<const __m256i *>(b + i)), vx, vq);
			if constexpr (add_to) {
				r = _mm256_add_epi32(r, _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + i)));
				r = _mm256_min_epu32(r, _mm256_sub_epi32(r, p));
			}
			_mm256_storeu_si256(reinterpret_cast<__m256i *>(a + i), r);
		}
		for (; i < n; ++i) {
			V r = product(b[i], x, q);
			if constexpr (add_to) a[i] = add(a[i], r);
			else a[i] = r;
		}
	}
#endif
	template <bool add_to>
	static void apply(V *a, const V *b, int n, V x) {
		V q = (std::uint64_t(x) << 32) / mod;
#if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
		if (n >= 64 && avx2() && __builtin_cpu_supports("avx512f")) { apply16<add_to>(a, b, n, x, q); return; }
		if (n >= 16 && avx2()) { apply8<add_to>(a, b, n, x, q); return; }
#endif
		for (int i = 0; i < n; ++i) {
			V r = product(b[i], x, q);
			if constexpr (add_to) a[i] = add(a[i], r);
			else a[i] = r;
		}
	}
	static void scale(V *a, int n, V x) { apply<false>(a, a, n, x); }
	static void axpy(V *a, const V *b, int n, V x) { apply<true>(a, b, n, x); }
};

}

#endif

Back to top page