Segmented Sieve of Eratosthenes in Scala

Question

I'm trying to solve PRIME1 in SPOJ with the following Scala snippet:

Peter wants to generate some prime numbers for his cryptosystem. Help him! Your task is to generate all prime numbers between two given numbers!

Input

The input begins with the number \$t\$ of test cases in a single line (\$t \le 0\$). In each of the next t lines there are two numbers \$m\$ and \$n\$ (\$1 \le m \le n \le 1000000000, n-m \le 100000\$) separated by a space.

Output

For every test case print all prime numbers \$p\$ such that \$m <= p <= n\$, one number per line, test cases separated by an empty line.

object _002_PRIME1 extends App {

  import scala.io.StdIn._

  /// ***1***
  // def isPrime(num: Int) = (2 until num).view.takeWhile(i => i * i <= num).forall(i => i * (num / i) < num)

  def sieve(m: Int, n: Int, step: Int, primeArray: Array[Boolean]): Unit = {
    val start = {
      if (m <= step) {
        2 * step
      } else {
        val pre = (m / step) * step
        if (pre == m) {
          pre
        } else {
          pre + step
        }
      }
    }
    Range(start, n + 1, step).foreach(i => primeArray(i - m) = false)
  }

  def simpleSieve(num: Int): IndexedSeq[Int] = {
    val basicPrimes = Array.fill(num - 1)(true)
    for {
      i <- 2 to num
      if basicPrimes(i - 2)
      j <- i * 2 to num by i
    } basicPrimes(j - 2) = false
    for {
      j <- basicPrimes.indices
      if basicPrimes(j)
    } yield (j + 2)
  }

  def sieveAll(m: Int, n: Int): Unit = {
    val primeArray = Array.fill(n - m + 1)(true)

    val basicPrimes = simpleSieve(math.sqrt(n).toInt)
    /// for(p<- basicPrimes)sieve(m, n, p, primeArray)
    basicPrimes.foreach(p => sieve(m, n, p, primeArray))
    if (m == 1) {
      primeArray(0) = false
    }
    for (j <- primeArray.indices; if primeArray(j)) println(j + m)
    /// ***2***
    //    Array.tabulate(primeArray.length) { i => {
    //      if (primeArray(i)) {
    //        println(i + m)
    //      }
    //    }
    //    }
  }

  (1 to readInt()).foreach(_ => {
    val ss = readLine().split("\\s+").map(_.toInt)
    val (m, n) = (ss.head, ss.last)
    sieveAll(m, n)
  })
}

For PRIME1 itself, this is able to output the correct answer within 0.6s. However for another problem, PRINT, which is almost the same except that the input range is bigger: \$2 \le L < U \le 2147483647, U-L \le 1000000\$, and the time limit is more restricted (1.223s).

When I submit this code, it always tells that time limit exceeded. How can I tune this snippet to make it run faster?

Attempts I've tried (the time costs are based on the OJ from PRIME1):

1. Using isPrime for basic primes checking, this would increase the running time to 0.7s.
2. Using Array.tabulate instead of for comprehension for output, decreasing time a little, about 0.01s.

Additionally, from Daniel C. Sobral's answer in the other question, mutable.BitSet may be more memory efficient, however this answer mentions that as to Java, an array might be more CPU efficient. I therefore didn't use BitSet for Scala (the memory used by is about 314MB while the memory limited by SPOJ is 1536MB).

Answer 1

There is a lot of overlap between the concept of windowed Sieve of Eratosthenes and that of segmented one, which is why the two are often confuddled.

A windowed sieve is rigged such that it can sieve a window, that is a range of numbers in the interval [m, n] (in SPOJ PRIME1/PRINT parlance), without having to sieve all the numbers before the start of the interval as well. The main trick here is to use modulo operations to find where the crossing-off sequence for a prime first intersects with the sieve window.

You have already implemented this, but I'm showing it again in pseudo code that makes the logic clearer. It is for an odds-only sieve, which is why there is a division by 2 at the end of the offset calculation, which transforms the starting point from the realm of numbers into the realm of odds-only bit indices.

internal static ushort[] small_odd_primes; // up to sqrt(MAX_N)
internal static bool[] sieve = new bool[1 << 15];

internal static void sieve_primes_between (int m, int n)
{
    // ...

    int last_prime_index = Array.BinarySearch(small_odd_primes, (ushort)Math.Sqrt(n));
    if (last_prime_index < 0)
        last_prime_index = -last_prime_index - 2;

    for (int i = 0; i <= last_prime_index; ++i)
    {
        int prime = small_odd_primes[i];
        int start = prime * prime;
        int stride = prime << 1;    

        if (start < m)
            start = (stride - 1) - (m - start - 1) % stride;
        else
            start -= m;

        offset[i] = start >> 1;
    }

    // ...

Of course, this not real pseudocode; it is my PRIME1/PRINT submission from my 'learn C#' tour of the world (fastest PRINT submission in C# and third fastest overall).

The modulo logic is a bit more complicated than what you did, for two reasons. First, I start crossing off at the square of a prime because all lower multiples will already have been crossed off during the runs for smaller primes. The effect minuscule, so don't worry about it. The second reason is that I wanted to minimise the number of branches necessary. Biassing the modulo operation by -1 makes it yield directly the desired result in the range 0 to stride-1, instead of giving 1 to stride where the last value has to be corrected to 0. Again, the effect is minor, but I wanted a clear, correct rendering of the windowed sieve offset logic on record.

Now, on to segmented operation. The purpose here is not to avoid the sieving of unwanted numbers as it is for a windowed sieve. The purpose is to sieve the desired range in small batches that do not exceed the size of the L1 cache of the target CPU (typically 32 KiB). For large sieve windows this can speed up things by orders of magnitude.

    for (int bits_to_sieve = ((n - m) >> 1) + 1; bits_to_sieve > 0; )
    {
        int window_bits = Math.Min(bits_to_sieve, sieve.Length);

        Array.Clear(sieve, 0, window_bits);

        for (int i = first_sieve_prime_index; i <= last_prime_index; ++i)
        {
            int p = small_odd_primes[i], j = offset[i];

            for ( ; j < window_bits; j += p)
                sieve[j] = true;

            offset[i] = j - window_bits;
        }

        append_sieved_numbers_to_result_vector(m, window_bits);

        m += window_bits << 1;
        bits_to_sieve -= window_bits;
    }

I think this makes the core idea of the segmented sieve very clear. In this case I'm reusing the sieve segment over and over, which is why I have to subtract window_bits before stashing an offset back in its slot. The same logic can also be used to process a bigger sieve in small, L1-sized stripes (but there's no advantage to that for PRIME1 and PRINT).

append_sieved_numbers_to_result_vector() just goes through the sieved range and pulls out the primes; I'm not showing it because it is actually twice as long as the sieve code (with mod 6 stepping and unrolling and all that) and still it takes twice as long as the actual sieving. It's core logic is quite simple, though:

    for (int i = 0; i < window_bits; ++i)
        if (!sieve[i])
            result[result_size++] = m + (i << 1);

In my submission I'm blasting a precomputed bit pattern over the sieve instead of calling Array.Clear(), in order to avoid the many iterations necessary for crossing off the smallest primes. However, the effect is not major and nothing to worry about.

What you do have to worry about is efficient extraction of the primes from the sieve, and I'd definitely recommend using an odds-only sieve in order to halve your workload in one fell swoop.

You'll also have to worry - a lot - about efficient output of the sieved primes. My submission used custom conversion code that rendered the numbers directly into the output buffer, digit by digit, in-place without any copying/moving. Using C#'s lame builtin I/O would have added at least a second to the timing, threatening TLE. See SPOJ's INOUTEST for the 'fast bulk output of numbers' thing. There are a lot of problems on SPOJ where fast bulk output of numbers comes in handy, and where the output time can dominate the total timing by a factor of 100 or so if you use the slow builtin I/O that comes with most languages. My first ETF submission got TLE, and after replacing the output code - not the actual processing related to the task - it clocked 0.030 vs. the time limit of 0.161, a nice big reserve.

I cannot comment on Scala's mutable.BitSet but I can assure you that its equivalents in C++ (std::vector<bool>) and in C# (System.Collections.BitArray) are slow as molasses and slow things down by an order of magnitude.

Custom code using a byte array and manual bit indexing (sieve[j>>3] |= 1 << (j & 7)) still takes twice as long as the plain bool[] sieve. It can be made about 33% faster than bool[] by copious unrolling, but the effort required is on the same order of magnitude as that for implementing a mod 30 wheel. To give you an impression, here is one of eight functions, each of which sieves a different residue class mod 8:

static void sieve_residue_class_1 ()
{
    var pao = primes_and_offsets[1];

    for (int i = 0, e = pao.Length; i < e; ++i)
    {
        int p = pao[i].prime, j = pao[i].offset, dj = p << 3;

        if (j + dj + dj <= window_bits)
        {
            int b = j >> 3, db = p >> 3;

            switch (j & 7)
            {
                case 4:  sieve[b] |= 1 << 4;  j += p;  b += db + 0;  goto case 7;
                case 7:  sieve[b] |= 1 << 7;  j += p;  b += db + 1;  goto case 2;
                case 2:  sieve[b] |= 1 << 2;  j += p;  b += db + 0;  goto case 5;
                case 5:  sieve[b] |= 1 << 5;  j += p;  b += db + 1;  goto case 0;
                case 0:  sieve[b] |= 1 << 0;  j += p;  b += db + 0;  goto case 3;
                case 3:  sieve[b] |= 1 << 3;  j += p;  b += db + 0;  goto case 6;
                case 6:  sieve[b] |= 1 << 6;  j += p;  b += db + 1;  break;
            }

            for (int next_j = j + dj; next_j <= window_bits; j = next_j, next_j += dj, b += p)
            {
                sieve[b + 0 * db + 0] |= 1 << 1;
                sieve[b + 1 * db + 0] |= 1 << 4;
                sieve[b + 2 * db + 0] |= 1 << 7;
                sieve[b + 3 * db + 1] |= 1 << 2;
                sieve[b + 4 * db + 1] |= 1 << 5;
                sieve[b + 5 * db + 2] |= 1 << 0;
                sieve[b + 6 * db + 2] |= 1 << 3;
                sieve[b + 7 * db + 2] |= 1 << 6;
            }   
        }

        for ( ; j < window_bits; j += p)
            sieve[j >> 3] |= (byte)(1 << (j & 7));

        pao[i].offset = j - window_bits;
    }
}

The effort just isn't worth it in most cases, and it exceeds the source code size limits for most tasks involving primes. Without such gargantuan efforts, a simple bool[] always wins. At least in C#, C++ and Delphi.

In C# I realised significant performance gains from allocating buffers (sieve, presieve pattern and so on) once in the constructor and using them over and over, instead of allocating and discarding as needed. Of course, this goes against the grain of languages like Scala but that's the reality of tasks like PRIME1 and language runtimes.

Moral: keep things simple and your caches hot.