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

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

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

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

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

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

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

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

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

$$ \begin{alignat*}{1} x = 30k + 0 & \quad \text{// делится на 2, 3, 5} \\ x = 30k + 1 & \quad \text{// не делится на 2, 3, 5} \\ x = 30k + 2 & \quad \text{// делится на 2} \\ x = 30k + 3 & \quad \text{// делится на 3} \\ x = 30k + 4 & \quad \text{// делится на 2} \\ x = 30k + 5 & \quad \text{// делится на 5} \\ x = 30k + 6 & \quad \text{// делится на 2, 3} \\ x = 30k + 7 & \quad \text{// не делится на 2, 3, 5} \\ & \ldots \\ x = 30k + 28 & \quad \text{// делится на 2} \\ 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$ ячеек. Памяти под решето понадобится в $3.75$ раза меньше ($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);
    }
}

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

Что дальше

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