Оптимизация решета Эратосфена. Wheel factorization

Что было ранее

Мы написали наивную версию решета Эратосфена и при помощи сегментации значительно ее ускорили:

Сейчас мы разберемся, как сделать меньше работы и поэтому выполнить ее быстрее. Также мы значительно уменьшим расход памяти под решето.

Благодарности

Оптимизация wheel factorization

Помните, в наивной версии мы проверяли только нечетные числа, но хранили все. Мы легко могли бы сократить расход памяти вдвое, избавившись от четных. Давайте разовьем эту идею. Возьмем первые простые числа $2$, $3$ и $5$. Их наименьшее общее кратное:

$$ \operatorname{НОК}(2,3,5) = 2 \cdot 3 \cdot 5 = 30 $$

Теперь представим, что мы хотим быстро проверить, делится ли число $x$ на один из этих делителей, если знаем остаток от деления числа $x$ на $30$:

$$ \begin{alignat*}{3} x &= 0 \pmod {30} \implies x = 30k + 0 && \quad \text{// делится на 2} \\ x &= 1 \pmod {30} \implies x = 30k + 1 && \quad \text{// не делится на 2, 3, 5} \\ x &= 2 \pmod {30} \implies x = 30k + 2 && \quad \text{// делится на 2} \\ x &= 3 \pmod {30} \implies x = 30k + 3 && \quad \text{// делится на 3} \\ x &= 4 \pmod {30} \implies x = 30k + 4 && \quad \text{// делится на 2} \\ x &= 5 \pmod {30} \implies x = 30k + 5 && \quad \text{// делится на 5} \\ x &= 6 \pmod {30} \implies x = 30k + 6 && \quad \text{// делится на 2} \\ x &= 7 \pmod {30} \implies x = 30k + 7 && \quad \text{// не делится на 2, 3, 5} \\ & \phantom{{}= 1 \pmod {30} \implies} \ldots \\ x &= 28 \pmod {30} \implies x = 30k + 28 && \quad \text{// делится на 2} \\ x &= 29 \pmod {30} \implies x = 30k + 29 && \quad \text{// не делится на 2, 3, 5} \\ \end{alignat*} $$

Проверив все $30$ остатков, мы можем заключить, что если число при делении на $30$ дает остаток не из списка $\{1, 7, 11, 13, 17, 19, 23, 29\}$, то это точно составное число (за исключением $2$, $3$ и $5$). Что это означает для нас?

Во-первых, мы можем не хранить числа, заведомо являющиеся составными. Для каждого блока из $30$ чисел потребуется всего $8$ ячеек. Расход памяти уменьшится в $30 / 8 = 3.75$ раза. А если хранить значения в битах, а не в байтах, то расход памяти уменьшится уже в $30$ раз.

блок 0:     1   7  11  13  17  19  23  29
блок 1:    31  37  41  43  47  49  53  59
блок 2:    61  67  71  73  77  79  83  89
блок 3:    91  97 101 103 107 109 113 119
 ...

Во-вторых, когда мы будем вычеркивать кратные $p$, то есть числа вида $p \cdot q$, будет достаточно рассматривать только те $q$, чей остаток от деления на $30$ совпадает с одним из $\{1, 7, 11, 13, 17, 19, 23, 29\}$.

Давайте на конкретном примере $p = 41$ разберемся, как использовать эту технику.

Без wheel factorization нам надо было бы вычеркнуть:

Я красным подсветил те числа, которые мы могли бы даже не рассматривать, так как они гарантированно являются составными. Это мы можем понять по остатку от деления второго множителя на $30$.

В итоге нам достаточно вычеркнуть только указанные ниже числа:

С точки зрения процессора умножение более трудоемкое, чем сложение. Давайте попытаемся понять, как уменьшить число умножений, чтобы наша программа выполнялась быстрее. Для этого посмотрите на то, как будет происходить вычеркивание кратных для простого $p = 41$:

• во-первых, тут видно, что информацию о каждых $30$ числах мы храним в блоках по $8$ чисел
• во-вторых, можно заметить, что процесс циклический и повторяется через $8$ шагов
• в-третьих, шаги внутри цикла разные, но мы можем их посчитать заранее, как и то, число на какой позиции вычеркиваем

Если хранить блок в $1$ байте, то кодом мы могли бы описать это так:

// вычеркивание составных для p = 41

let mask_r1 = 1 << 0;     // первому числу соответствует 0-ой бит
let mask_r7 = 1 << 1;     // второму числу соответствует 1-ый бит
let mask_r11 = 1 << 2;    // ...
let mask_r13 = 1 << 3;
let mask_r17 = 1 << 4;
let mask_r19 = 1 << 5;
let mask_r23 = 1 << 6;
let mask_r29 = 1 << 7;

// p = 41
let mut b = p * p / 30;
loop {
    sieve[b] |= mask_r1;
    b += 2;
    sieve[b] |= mask_r23;
    b += 6;
    sieve[b] |= mask_r7;
    b += 2;
    sieve[b] |= mask_r29;
    b += 6;
    sieve[b] |= mask_r13;
    b += 8;
    sieve[b] |= mask_r19;
    b += 3;
    sieve[b] |= mask_r11;
    b += 8;
    sieve[b] |= mask_r17;
    b += 6;
}

Как видно по коду, нам только для самого первого блока нужно произвести умножение ($b = p * p / 30$), номер каждого следующего блока мы получаем тривиальным сложением.

mask_r1  = 1 << 0; // = 00000001 — в бинарной записи
mask_r7  = 1 << 1; // = 00000010
mask_r11 = 1 << 2; // = 00000100
mask_r13 = 1 << 3; // = 00001000 
mask_r17 = 1 << 4; // = 00010000
mask_r19 = 1 << 5; // = 00100000
mask_r23 = 1 << 6; // = 01000000
mask_r29 = 1 << 7; // = 10000000

sieve[b] |= mask;

let res = sieve[b] & mask;

mask_r11 = 1 << 2;
mask_r19 = 1 << 5;
sieve[b] |= mask_r11; // sieve[b]: 00000000 -> 00000100
sieve[b] |= mask_r19; // sieve[b]: 00000100 -> 00100100

let res = sieve[b] & mask_r11; // = 00000100 - помеченное

let res = sieve[b] & mask_r13; // = 00000000 - непомеченное

Но этот код подходит только для $p = 41$, а что если мы рассмотрим произвольное $p = 30k + 11$?

Эти переходы можно описать кодом:

// p = 30k + 11
let mut b = p * p / 30;
let k = p / 30;

loop {
    sieve[b] |= mask[1];
    b += 2*k; 
    sieve[b] |= mask[23];
    b += 4*k + 2; 
    sieve[b] |= mask[7];
    b += 2*k; 
    sieve[b] |= mask[29];
    b += 4*k + 2;
    sieve[b] |= mask[13];
    b += 6*k + 2;
    sieve[b] |= mask[19];
    b += 2*k + 1;
    sieve[b] |= mask[11];
    b += 6*k + 2;
    sieve[b] |= mask[17];
    b += 4*k + 2;
}

Давайте добавим в наш код проверку, не вышли ли мы за границу решета. А чтобы не дублировать однотипные прыжки и вычеркивания чисел, сохраним их в массиве cycle, обработчик напишем один раз:

    const WHEEL_SIZE: usize = 30;
    const RESIDUES_PER_WHEEL: usize = 8;
    
    let mut bit_mask: [u8; WHEEL_SIZE] = [0; WHEEL_SIZE];
    for (i, &r) in [1, 7, 11, 13, 17, 19, 23, 29].iter().enumerate() {
        bit_mask[r] = 1 << i;
    }
    
    for p in small_primes {
        let mut b = (p * p) / WHEEL_SIZE;
        let r = p % WHEEL_SIZE;
        let k = p / WHEEL_SIZE;

        let cycle: [(u8, usize); RESIDUES_PER_WHEEL] = match r {
            11 => {
                [
                    (bit_mask[1],  2 * k),
                    (bit_mask[23], 4 * k + 2),
                    (bit_mask[7],  2 * k),
                    (bit_mask[29], 4 * k + 2),
                    (bit_mask[13], 6 * k + 2),
                    (bit_mask[19], 2 * k + 1),
                    (bit_mask[11], 6 * k + 2),
                    (bit_mask[17], 4 * k + 2),
                ]
            }
            ...
        };
        
        let max_b = n / WHEEL_SIZE;
        while b <= max_b {
            for (bit, jump) in cycle {
                sieve[b] |= bit;
                b += jump;
                if b > max_b {
                    break;
                }
            }
        }
    }  

Исходный код можно посмотреть здесь . Добавим результаты запуска этой версии:

На данный момент мы разработали две разные версии программы, ускоряющие наивное решение. Каждая из этих версий использует свой независимый подход:

• в версии с сегментацией мы говорим: “нам приходится обращаться к данным в оперативной памяти снова и снова, давайте использовать кэш процессора так, чтобы повторные обращения были более эффективными”.
• в версии с wheel factorization идея следующая: “нам не надо обрабатывать весь объем данных, большую часть данных мы можем пропустить, таким образом нам потребуется выполнить значительно меньше операций, и всю работу мы сможем завершить существенно быстрее”. При этом в версии с wheel factorization мы будем “промахиваться” мимо кэша процессора больше, чем могли бы. И об этом поговорим дальше.

Хочу отметить, что так как теперь для блока из $30$ чисел используется $1$ байт вместо $30$, то и выигрыш по памяти соответствующий. Решето в наивной версии занимало $2048$ MiB, в этой версии занимает $2048 / 30 \approx 68$ MiB.

Сегментированное решето с wheel factorization

Мы смогли существенно уменьшить размер выделяемой памяти для решета, применив wheel factorization (до $\approx 68$ MiB). Однако даже такое улучшение не позволяет разместить решето целиком в кэше процессора, нам придется обращаться к данным в оперативной памяти для каждого малого простого снова и снова. Сегментация здесь должна помочь.

Разделим наше решето, как и ранее, на кусочки размером $100\ 000$. В версии, где применялась сегментация, мы для каждого простого $p$ вычисляли, с какой именно позиции будем начинать в новом сегменте. Это было легко сделать, так как прыжки между вычеркиваемыми числами были одинаковыми. В версии с wheel factorization прыжки разные, напомню для $p = 30k + 11$:

$$ \begin{alignat*}{1} j0 &= 2k \\ j1 &= 4k + 2 \\ j2 &= 2k \\ j3 &= 4k + 2 \\ j4 &= 6k + 2 \\ j5 &= 2k + 1 \\ j6 &= 6k + 2 \\ j7 &= 4k + 2 \\ \end{alignat*} $$

Обратите внимание, что:

$$ j0 + j1 + … + j7 = 30k + 11 = p $$

Через $8$ шагов цикл повторяется (длина цикла $=p$), но внутри цикла шаги неодинаковые. Давайте делать так: если во время прыжка мы вышли за границу сегмента, то запомним блок, на который прыгнули, и номер прыжка внутри цикла. Тогда для каждого малого простого мы будем хранить позицию, с которой надо начинать в новом сегменте.

Так как мы начинаем вычеркивания с $p^2$ и в одном блоке храним $30$ чисел, то начинать будем с блока $\lfloor p^2 / 30 \rfloor$ и с $0$-го шага внутри цикла:

const WHEEL_SIZE: usize = 30;

let mut resume_state: Vec<_> = 
    small_primes
        .iter()
        .map(|&p| (p * p / WHEEL_SIZE, 0))
        .collect();

А далее для каждого сегмента:

const RESIDUES_PER_WHEEL: usize = 8;

let sieve_len = (n / WHEEL_SIZE) + 1;

for chunk_start in (0..sieve_len).step_by(CHUNK) {
    let chunk_end = (chunk_start + CHUNK).min(sieve_len - 1);
    for (i, &p) in small_primes.iter().enumerate() {
        let r = p % WHEEL_SIZE;
        let k = p / WHEEL_SIZE;
        
        let cycle: [(u8, usize); RESIDUES_PER_WHEEL] = match r {
            // 1 => { ... }
            // 7 => { ... }
            11 => {
                [
                    (bit_mask[1],  2 * k),
                    (bit_mask[23], 4 * k + 2),
                    (bit_mask[7],  2 * k),
                    (bit_mask[29], 4 * k + 2),
                    (bit_mask[13], 6 * k + 2),
                    (bit_mask[19], 2 * k + 1),
                    (bit_mask[11], 6 * k + 2),
                    (bit_mask[17], 4 * k + 2),
                ]
            }
            // 13 => { ... }
            // 17 => { ... }
            // 19 => { ... }
            // 23 => { ... }
            // 29 => { ... }
        };
        
        let (mut b, mut s) = resume_state[i];
        let mut next_s = s;
        'outer: while b <= chunk_end {
            for (j, &(bit, jump)) in cycle[s..].iter().enumerate() {
                sieve[b] |= bit;
                b += jump;
                if b > chunk_end {
                    next_s = (s + j + 1) % RESIDUES_PER_WHEEL;
                    break 'outer;
                }
            }
            s = 0;
            next_s = 0;
        }
        resume_state[i] = (b, next_s);
    }

Исходный код данной версии можно посмотреть здесь . Добавим результаты запусков:

Что дальше

В следующей части мы, наконец, сравним наш код с оригинальным решением и проведем параллели. Дополнительная оптимизация позволит почти уравнять две версии.