Блог пользователя alexvim

Редукция Монтгомери/Баррета и Теоретико-числовое Преобразование Фурье [Оптимизации]
Автор alexvim, история, 6 месяцев назад, По-русски

$$$\DeclareMathOperator{\ord}{ord}$$$ $$$\DeclareMathOperator{\mod}{mod}$$$

Редукция Монтгомери/Баррета и приложение к Теоретико-числовому Преобразованию Фурье

Здесь будет показано, как применить редукцию Монтгомери/Баррета к оптимизации Теоретико-числового Преобразования Фурье (NTT) без явных операций по модулю и созданию собственного класса модулярной арифметики, применимого в комбинаторных задачах. Будет приведена реализация NTT с минимальным числом операций по модулю, которая выигрывает у классического комплексного FFT не только в точности, но и в производительности на платформах без аппаратного ускорения. Мы сравним производительность различных оптимизаций NTT с комплексным FFT. Также будет показано, как добавить поддержку отрицательных коэффициентов в NTT.

1. Причины использования NTT вместо FFT

DFT (Дискретное Преобразование Фурье) широко используется в задаче умножения многочленов. DFT может быть вычислено в любом поле, где определён корень n-й степени, т.е. элемент с порядком (показателем) n. Например, это может быть поле комплексных чисел (используемое в FFT) или кольцо вычетов по простому модулю (NTT). Известная реализация — Быстрое Преобразование Фурье — вычисляет DFT за $$$O(n\log{n})$$$ в любом поле (отличается только арифметика). Реализация FFT по модулю $$$p$$$ называется Теоретико-числовым Преобразованием Фурье (NTT). Просто FFT будем называть алгоритм в комплексном поле. Базовый алгоритм без оптимизаций рассмотрен во многих статьях, например, на CP-algorithms. NTT почти не используют, потому что его принятая реализация с модулярной арифметикой по большому модулю является очень медленной без аппаратного ускорения. Здесь будет показано, как реализовать быстрый NTT с неявными операциями по модулю.

Конечно, есть удобный FFT с комплесным типом double, который используется в перемножении многочленов (и не только), в т.ч. для решения задач на Codeforces. Но FFT имеет очень плохую точность из-за ограниченной памяти машинной реализации арифметики с плавающей точкой, и он довольно медленный (даже с нерекурсивной реализацией). Следовательно, если даны многочлены с неотрицательными большими коэффициентами для перемножения, есть возможность использования NTT: перемножить многочлены по модулю $$$p$$$ и восстановить реальные коэффициенты до $$$10^{18}$$$ (например, с помощью двух 32-битных модулей и Китайской Теоремы об Остатках для решения систем сравнений по модулям или просто используя большой модуль с кастом к 128-битному целочисленному типу).

Главное преимущество NTT над FFT — это точность: через NTT мы получаем абсолютно точные коэффициенты произведения в диапазоне $$$[0; 10^{18}]$$$, а использование FFT приводит к ошибке 10-100 при использовании double и 1-10 при long double (например, 259759436 c double получает WA43 в задаче, сводимой к умножению многочленов, когда как 259759880 с long double получает OK).

NTT не работает быстрее FFT "из коробки", т.к. стандартная реализация NTT использует много операций по модулю, так что на многих системах (особенно без встроенного ускорения через SIMD) оно работает медленно. Но оптимизации описанные ниже помогут нам избавиться от явных операций взятия остатка, и NTT сможет работать заметно быстрее FFT. Для этого мы можем использовать редукцию Монтгомери и редукцию Баррета.

2. Редукция Монтгомери

Редукция Монтгомери — это подход к организации арифметики по фиксированному модулю без явного деления с остатком (%). Главная идея в том, что мы можем свести деление с остатком на $$$p$$$ к делению с остатком на $$$2^k$$$, которое может быть реализовано через битовые операции. Для ознакомления с самим алгоритмом и доказательством можно в статье. Редукция Монтгомери обязывает перед выполнением операций по модулю приведение к специальной форме пространства Монтгомери. Приведение к этой форме требует явного деления с остатком, но в NTT особый случай: все числа, с которыми мы будем оперировать по модулю в алгоритме — это только коэффициенты исходных многочленов, корень и некоторые обратные (n, корень). Мы можем сделать это за $$$O(n)$$$ перед вызовом NTT и после него (чтобы привести в нормальную форму). Следовательно, можно избавиться от модулярной арифметики внутри NTT. Единственный нюанс в том, что при использовании большого модуля ($$$~10^{18}$$$) мы вынуждены использовать __int128. Из-за изобилия операций с многобитными числами была даже выпущена статья Intel для ускорения редукции Монтгомери через AVX.

Но если мы хотим уменьшить число операций с многобитными числами без потери в размерах коэффициентах, можно вычислить NTT по двум модулям, используя только 64-битные числа (параллельно через 1 вызов) и решить систему сравнений по простым модулям через Китайскую Теорему об Остатках. Тогда __int128 будет использован только в КТО.

P.S. В будущем статья будет дополнятся полностью русскоязычными разборами модулярных редукций.

3. Редукция Баррета

Другой способ избежать явное деление с остатком — это редукция Баррета. Она основана на приближении ответа через деление с округлением вниз с предподсчитанным множителем определённого вида. Конкретно, преобразуем выражение для остатка:

$$$ a \mod p = a - \lfloor {a \over p} \rfloor \cdot p $$$

Выберем $$$m$$$, так что:

$$$ {1 \over p} = {m \over 2^k} \longleftrightarrow m = {2^k \over p} $$$

Теперь при делении будем подставлять $$$1 \over p$$$. Если взять достаточно большой множитель $$$2^k$$$ для приближения (порядка $$$p^2$$$) ответ для чисел нашего диапазона (до $$$10^{18}$$$) будет точным. С полным доказательством можно ознакомиться здесь. Она даёт такой же функционал, что и редукция Монтгомери, но не требует приведение к специальной форме. Однако мы в каждой операции должны делать битовый сдвиг на больше, чем $$$p^2$$$ и делать умножения, поэтому при использовании модуля около $$$10^{18}$$$ необходимо использовать 256-битный целочисленный тип, поэтому операции дороже.

4. Выбираем модули

Во первых, мы должны выбрать модуль такой модуль $$$p$$$, что $$$n | (p-1): n=2^k$$$, чтобы существовал корень n-й степени из 1 (порядок элемента делит порядок циклической группы по теореме Лагранжа). Т.к. по модулю $$$p$$$ существует первообразный корень $$$g$$$, мы можем взять $$$g^{p-1 \over n}$$$ в качестве корня n-й степени. Для конкретного $$$p$$$ эти корни могут быть легко найдены с помощью бинарного возведения в степень и алгоритмом поиска первообразного корня за $$$O(Ans \cdot \log^2(p))$$$, $$$Ans$$$ — размер первообразного корня в конечном поле — по предположению гипотезы всегда маленький. Здесь имеется доказательство корректности алгоритма. Например, в нашей задаче для перемножения многочленов размером до $$$2^{24}$$$ с коэффициентами произведения до $$$2 \cdot 10^{18}$$$ можно взять $$$p=2524775926340780033$$$. Первообразный корень 3.

Primitive root searching algorithm

5. Реализации

Будем использовать нерекурсивную реализацию для FFT (NTT) с дополнительной структурой для арифметики с редукцией Монтгомери/Баррета. Удобно использовать беззнаковые типы при работе с редукциями из-за трюка с переполнением: оно не ведёт к неопределённому поведению, а возвращает ответ по модулю $$$2^{64}$$$.

NTT with Montgomery multiplication

Реализация с редукцией Баррета может быть использована только с модулем $$$<2^{32}$$$ без 256-битных чисел. Т.к. редукция Баррета априори имеет более дорогие операции, здесь не будет приведена реализация с uint256. В коде можно использовать uint32 вместо uint64, но ради равенства условий в тестировании будут приведены одинаковые форматы и в реализации с Монтгомери и с Барретом.

NTT with Barret multiplication

Также здесь будет приведена реализация FFT со своей структурой для комплексных чисел. Она может быть использована с double / long double.

FFT with custom cmpls

NTT без оптимизаций:

NTT without opt

6. Производительность

Мы будет тестировать 5 алгоритмов: FFT с double и long double, NTT без и с редукциями Монтгомери и Баррета.

Во-первых, посмотрим на 993E - Никита и порядковая статистика. Важно отметить, что NTT с редукцией Баррета и FFT с double получили wa43 (но прошли тесты на других тестах с большими входными данными), потому что FFT с double имеет плохую точность, а NTT с редукцией Баррета по модулю $$$998244353$$$ не может хранить большой ответ (NTT с Барретом и без uint256 может быть использовано только для небольших коэффициентов).

FFT, doubles FFT, long doubles NTT, Montgomery NTT, Barret NTT, no opt
261258380, 311ms, BAD PRECISION 261256931, 702ms, OK 261256487, 281ms, OK 261257532, 312ms, BAD LIMITS 261257532, 859ms, OK

Теперь протестируем перемножение многочленов размером $$$10^6$$$ со случайными коэффициентами на запуске Codeforces:

FFT, doubles FFT, long doubles NTT, Montgomery NTT, Barret NTT, no opt
1407ms 2434ms 1374ms 1437ms 4122ms
1286ms 2426ms 1256ms 1675ms 3683ms
1052ms 2286ms 1153s 1320ms 3518ms
1079ms 2359ms 1309ms 1380ms 3961ms
1217ms 2791ms 1476ms 1780ms 3400ms
1241ms 2191ms 1125ms 1337ms 3512ms

Теперь локальные тесты на Ryzen 5 5650u (AVX), 8Gb RAM 4266MHz, Debian 12, GCC 12.2.0. Параметры компиляции: g++ -Wall -Wextra -Wconversion -static -DONLINE_JUDODGE -O2 -std=c++20 fftest.cc -o fftest как на Codeforces.

FFT, doubles FFT, long doubles NTT, Montgomery NTT, Barret NTT, no opt
691ms 2339ms 352ms 401ms 398ms
700ms 2328ms 411ms 500ms 487ms
729ms 2036ms 359ms 446ms 396ms
698ms 2164ms 380ms 456ms 472ms
734ms 2284ms 407ms 462ms 440ms
737ms 2148ms 396ms 459ms 434ms

На сервере Codeforces одинаковая производительность у Montgomery + NTT и FFT + doubles, затем идёт NTT + Barret, затем FFT + long doubles, затем NTT без оптимизаций. На локальной системе с современным процессором с поддержкой SIMD лидером является Montgomery + NTT. Причём NTT на локальной системе работает намного быстрее даже без оптимизаций. Я точно не могу определить, какие технологии отвечают за быструю 128-битную модулярную арифметику, но, скорее всего, это часть SIMD. Вообще, NTT хорошо поддаётся аппаратному ускорению, напирмер, в статье описано ускорение NTT на FPGA.

Как мы видим, производительность NTT по большому модулю сильно упирается в платформу из-за наличия большого числа операций с высокобитными числами, а именно деления с остатком 128-битных чисел. Это операция сильно упирается в аппаратную реализацию, поэтому из-за неё на Codeforces NTT без оптимизаций оказался самым медленным, но на CPU с ускорением SIMD он выигрывает у FFT. Тем не менее, разница может быть нивелирована засчёт явной реализации редукции, которая использует только умножение 128-битных чисел. Про проблемы 128-битного деления изложено в статье.

В итоге NTT + Montgomery является довольно универсальным выбором на всех платформах. Единственный нюанс в том, что NTT изначально не поддерживает отрицательные коэффициенты (т.к. всё берётся по модулю), но это можно исправить.

7. Отрицательные коэффициенты

Рассмотрим $$$A(x)$$$ и $$$B(x)$$$ с любыми целыми коэффициентами, и мы хотим получить коэффициенты произведения через NTT. Пусть $$$C(x)$$$ — такой многочлен, что любой коэффициент у $$$A(x)+C(x)$$$, $$$B(x)+C(x)$$$ и $$$A(x)+B(x)+C(x)$$$ является неотрицательным. Заметим, что

$$$ A \cdot B = (A+C) \cdot (B+C) - C \cdot (A+B+C) $$$

Так как каждое слагаемое в сумме имеет неотрицательные коэффициенты, $$$A(x) \cdot B(x)$$$ можно найти через 2 вызова умножения многочленов через NTT. $$$C(x)$$$ можно найти жадно за $$$O(n)$$$.

8. Класс для модулярной арифметики

Удобно создать класс для арифметики в $$$\mathbb{Z}/p\mathbb{Z}$$$ и какой-то редукцией. Ниже приведён класс с редукцией Баррета с методами для НОДа, степени, инверсий и корней (класс можно расширить до арифметики в любом конечном поле). P.S. Если мы хотим получить максимальную производительность и работаем с специфичным множеством чисел, можно использовать редукцию Монтгомери, потому что Монтгомери быстрее.

Modular arithmetic class

9. Заключение

В результате мы сравнили различные редукции для модульной арифметики и выяснили, что редукция Монтгомери имеет наилучшую производительность, если мы заранее знаем набор целых чисел. Следовательно, NTT с редукцией Монтгомери можно использовать вместо FFT в задачах с умножением многочленов (даже для полиномов с отрицательными коэффициентами). Здесь приведена реализация NTT с редукцией Монтгомери, которая показывает лучшую производительность даже по сравнению с FFT с double, особенно с аппаратной оптимизацией современных x64 процессоров (SIMD?). Остается открытым вопрос, какие именно аппаратные оптимизации позволяют NTT работать быстрее на современных процессорах "из коробки".

Полный текст и комментарии »

  • Проголосовать: нравится
  • +72
  • Проголосовать: не нравится

Абсурдный алгоритм нахождения простых чисел, использующий длинную арифметику
Автор alexvim, история, 8 месяцев назад, перевод, По-русски

Абсурдный алгоритм нахождения простых чисел, использующий длинную арифметику

Рассмотрим тривиальный и, наверное, кратчайший (но не эффективный) алгоритм наождения протых чисел $$$p \in \mathbb{N}: p \leq n$$$, основанный на теореме Вильсона и длинной арифметике.

1. Теорема Вильсона

Алгоритм основан на известной в элементарной теории чисел теореме Вильсона. Часто она доказывается только слева направо, но важно доказать критерий.

Теорема (Вильсон).

$$$p \in \mathbb{N}: p \ge 2$$$ — простое тогда и только тогда, когда $$$(p-1)! \equiv -1 \pmod{p}.$$$

Доказательство.

$$$\rightarrow \quad$$$ p — простое, сл-но $$$\forall$$$ a $$$\in \mathbb{Z}/p\mathbb{Z}$$$ $$$\quad$$$ $$$\exists a^{-1} \in \mathbb{Z}/p\mathbb{Z}: a\cdot{a^{-1}}=1$$$. Следовательно, в произведении $$$1 \cdot 2 \cdot \dotso \cdot (p-1)$$$ все элементы разбиваются на пары $$$(a, b)$$$, такие что $$$ab=1$$$.
$$$\leftarrow$$$ $$$\quad$$$ Предположим, что p — составное. Если p — не полный квадрат: $$$p \neq a^2 \quad \forall a \in \mathbb{Z}$$$, тогда $$$\exists a,b \in \mathbb{N}: a,b < p$$$ и $$$ab=p$$$, сл-но $$$p=ab$$$ делит $$$(p-1)!$$$. Иначе если $$$p=a^2: a>1$$$, тогда a делит $$$gcd(p, (p-1)!)$$$, сл-но a делит $$$(p-1)!$$$ mod $$$p$$$ и как результат: $$$(p-1)! \not\equiv -1 \pmod{p}$$$ противоречие.

2. Алгоритм

Из теоремы Вильсона может быть выведен достаточно простой алгоритм нахождения простых чисел. 

a = 1
for i in range(2, n):
    if a%i == i-1:
        pass #prime
    a *= i

3. Длинная арифметика

Легко видеть, что алгоритм требует наличие длинной арифметики, т.к. факториал растёт быстро. Одна из эффективных реализаций длинной арифметики использует FFT (Быстрое Преобразование Фурье) для перемножения чисел (многочленов). FFT умножает $$$a\cdot{b}$$$ за $$$O(n\log{n})$$$, где n — максимальная битовая длина чисел. Мы используем этот факт в анализе. P.S. FFT с эвристиками использован и в python.

4. Асимптотика

Заметим, что на k-й итерации длина $$$a$$$ $$$O(\log{k!})$$$, так что общая асимптотика

$$$O(\sum_{k=2}^n{\log{k!}\cdot{\log{\log{k!}}})}$$$

Утверждение.

$$$\sum_{k=2}^n{\log{k!}\cdot{\log{\log}k!}} = O(n^2\log^2{n})$$$

Доказательство.

Используя формулу Стирлинга для логарифма гамма-функции (факториала):

$$$\log{k!} = k\log{k} — k + O(\log{k}) = k\log{k} + O(k)$$$

получаем:

$$$\sum_{k=2}^n{\log{k!}\cdot{\log{\log{k!}}}} = \sum_{k=2}^n{(k\log{k} + O(k))\log{(k\log{k} + O(k))}}$$$

Используя факт $$$\log{(f + o(f))} = \log{f} + o(\log{f})$$$ (можно явно проверить по правилу Лопиталя) получаем, что сумма асимптотически стремиться к:

$$$\sum_{k=2}^n{k\log^2{k}}$$$

Известна формула суммирования Эйлера-Маклорена, в асимптотическом виде представимая как $$$\sum_{k=a}^b{f(k)} = \int_a^b{f(x)dx} + O(f)$$$ для некоторого класса функций $$$f$$$, таких что $$$f'(x) = O(f)$$$ или $$$f'(x) = o(f)$$$. Поэтому мы можем приблизить сумму интегралом:

$$$\sum_{k=2}^n{k\log^2{k}} = O(\int_2^n{x\log^2{x}dx})$$$

Применяя интегрирование по частям 2 раза, убеждаемся, что значение интеграла — $$$O(n^2\log^2{n})$$$, ч.т.д..

5. Заключение

Алгоритм довольно абсурдный, но показывает некоторые техники работы с суммами, встречающимися в алгоритмах длинной арифметики.

Полный текст и комментарии »

  • Проголосовать: нравится
  • +23
  • Проголосовать: не нравится