Python基于轮分解实现高效素数生成器的问题排查与优化

我一直尝试编写能够有效利用wheel factorisation（轮分解）的高效素数生成器，期望其性能优于基础sieve of Eratosthenes（埃拉托斯特尼筛法）实现，但我的代码因开销过高实际运行速度慢很多。我希望了解现有实现运行缓慢的原因、可行的改进方法，以及其他素数生成器能够实现更高性能的原理。

模30轮分解实现代码 #

由于需要硬编码模30轮逻辑，代码篇幅较长，具体实现如下：

import numpy as np
def wheel_factorisation(number):
    wheel = [1, 7, 11, 13, 17, 19, 23, 29]
    difference = [2, 6, 4, 2, 4, 2, 4, 6]
    squareroot = math.isqrt((number // 30) * 8)
    isprime = np.ones(squareroot,dtype=bool)
    isprime[0] = False

    multiples = []
    primes = [2,3,5]

    factor = 7
    index = 1
    DIVindex = 0

    while factor * factor <= number:
        if isprime[DIVindex * 8 + index]:
            primes.append(factor)
            if index == 0:  #1%30
                multiples.append([8*DIVindex*(factor+1), 0, 0, DIVindex])
                #offset = [0, 1, 2, 3, 4, 5, 6, 7]
                offset = np.ones(8,dtype = int)
                #DIVoffset = [1, 0, 0, 0, 0, 0, 0, 0]
                    
            if index == 1:  # 7%30
                multiples.append([8*DIVindex * (factor + 7) + 13, 1, 1, DIVindex])

                #offset = [1, 5, 4, 0, 7, 3, 2, 6]
                offset = np.array([3,12,7,4,7,4,7,12])
                #DIVoffset = [1, 1, 1, 1, 0, 1, 1, 1]

            if index == 2:  #11%30
                multiples.append(
                    [8*DIVindex * (factor + 11) + 32, 2, 2, DIVindex])
                #offset = [2, 4, 0, 6, 1, 7, 3, 5]
                offset = np.array([5,18,12,6,11,6,12,18])
                #DIVoffset = [1, 2, 2, 0, 2, 0, 2, 2]

            if index == 3:  #13%30
                multiples.append(
                    [8*DIVindex * (factor + 13) + 45, 3, 3, DIVindex])
                #offset = [3, 0, 6, 5, 2, 1, 7, 4]
                offset = np.array([7,21,14,7,13,7,14,21])
                #DIVoffset = [1, 3, 1, 1, 2, 1, 1, 3]

            if index == 4:  #17%30
                multiples.append(
                    [8*DIVindex * (factor + 17) + 77, 4, 4, DIVindex])
                #offset = [4, 7, 1, 2, 5, 6, 0, 3]
                offset = np.array([9,27,18,9,19,9,18,27])
                #DIVoffset = [1, 3, 3, 1, 2, 1, 3, 3]

            if index == 5:  #19%30
                multiples.append(
                    [8*DIVindex * (factor + 19) + 96, 5, 5, DIVindex])
                #offset = [5, 3, 7, 1, 6, 0, 4, 2]
                offset = np.array([11,30,20,10,21,10,20,30])
                #DIVoffset = [1, 4, 2, 2, 2, 2, 2, 4]

            if index == 6:  #23%30
                multiples.append(
                    [8*DIVindex * (factor + 23) + 141, 6, 6, DIVindex])
                #offset = [6, 2, 3, 7, 0, 4, 5, 1]
                offset=np.array([13,36,25,12,25,12,25,36])
                #DIVoffset = [1, 5, 3, 1, 4, 1, 3, 5]

            if index == 7:
                multiples.append(
                    [8*DIVindex * (factor + 29) + 224, 7, 7, DIVindex])
                #offset = [7, 6, 5, 4, 3, 2, 1, 0]
                offset = np.array([15,47,31,15,31,15,31,47])
                #DIVoffset = [1, 6, 4, 2, 4, 2, 4, 6]

            while multiples[-1][0] < squareroot:
                isprime[multiples[-1][0]] = False
                multiples[-1][1] += 1
                if multiples[-1][1] == 8:
                    multiples[-1][1] = 0
                    
                multiples[-1][0] += offset[multiples[-1][1]] + 8*DIVindex*difference[multiples[-1][1]]

        index += 1
        if index == 8:
            DIVindex += 1
            index = 0
        factor += difference[index]
    prime = [DIVindex, index]
    while prime[0] * 8 + prime[1] < squareroot:
        if isprime[prime[0] * 8 + prime[1]]:
            primes.append(prime[0] * 30 + wheel[prime[1]])
        prime[1] += 1
        if prime[1] == 8:
            prime[1] = 0
            prime[0] += 1
    low = squareroot
    limit = low + low
    while low <= (number * 8) // 30:
        isprime = np.ones(squareroot,dtype=bool)
        if limit > (number * 8) // 30:
            limit = (number * 8) // 30
        factor = 0
        while factor < len(multiples) and primes[factor +
                                                 3] <= math.isqrt(number):
            index = multiples[factor][2]
            DIVindex = multiples[factor][3]
            if multiples[factor][0] < limit:    
                if index == 0:  #1%30
                    #offset = [0,1,2,3,4,5,6,7]
                    offset = np.ones(8,dtype = int)
                    #DIVoffset = [1, 0, 0, 0, 0, 0, 0, 0]

                if index == 1:  # 7%30
                    #offset = [1, 5, 4, 0, 7, 3, 2, 6]
                    offset = np.array([3,12,7,4,7,4,7,12])
                    #DIVoffset = [1, 1, 1, 1, 0, 1, 1, 1]

                if index == 2:  #11%30
                    #offset = [2, 4, 0, 6, 1, 7, 3, 5]
                    offset = np.array([5,18,12,6,11,6,12,18])
                    #DIVoffset = [1, 2, 2, 0, 2, 0, 2, 2]

                if index == 3:  #13%30
                    #offset = [3, 0, 6, 5, 2, 1, 7, 4]
                    offset = np.array([7,21,14,7,13,7,14,21])
                    #DIVoffset = [1, 3, 1, 1, 2, 1, 1, 3]

                if index == 4:  #17%30
                    #offset = [4, 7, 1, 2, 5, 6, 0, 3]
                    offset = np.array([9,27,18,9,19,9,18,27])
                    #DIVoffset = [1, 3, 3, 1, 2, 1, 3, 3]

                if index == 5:  #19%30
                    #offset = [5, 3, 7, 1, 6, 0, 4, 2]
                    offset = np.array([11,30,20,10,21,10,20,30])
                    #DIVoffset = [1, 4, 2, 2, 2, 2, 2, 4]

                if index == 6:  #23%30
                    #offset = [6, 2, 3, 7, 0, 4, 5, 1]
                    offset=np.array([13,36,25,12,25,12,25,36])
                    #DIVoffset = [1, 5, 3, 1, 4, 1, 3, 5]

                if index == 7:  #29%30
                    #offset = [7, 6, 5, 4, 3, 2, 1, 0]
                    offset = np.array([15,47,31,15,31,15,31,47])
                    #DIVoffset = [1, 6, 4, 2, 4, 2, 4, 6]
                while multiples[factor][0] < limit:    
                    isprime[multiples[factor][0]-low] = False
                    multiples[factor][1] +=1
                    if multiples[factor][1]==8:
                        multiples[factor][1]=0
                        
                    multiples[factor][0]+= 8*DIVindex * difference[
                        multiples[factor][1]]
                    multiples[factor][0] += offset[multiples[factor][1]]
                    
            factor += 1
        while prime[0]*8+prime[1]< limit:
            if isprime[prime[0]*8+prime[1] - low]:
                primes.append(prime[0] * 30 + wheel[prime[1]])
            prime[1] += 1
            if prime[1] == 8:
                prime[1] = 0
                prime[0] += 1
        low += squareroot
        limit += squareroot
    return primes

性能瓶颈定位 #

经排查，两个核心代码段占据了绝大多数运行时间：

1. 标记素数倍数的while循环：

while multiples[factor][0] < limit:    
    isprime[multiples[factor][0]-low] = False
    multiples[factor][1] +=1
    if multiples[factor][1]==8:
        multiples[factor][1]=0
        
    multiples[factor][0]+= 8*DIVindex * difference[
        multiples[factor][1]]
    multiples[factor][0] += offset[multiples[factor][1]]

2. 从布尔数组收集素数的代码段：

while prime[0]*8+prime[1]< limit:
    if isprime[prime[0]*8+prime[1] - low]:
        primes.append(prime[0] * 30 + wheel[prime[1]])
    prime[1] += 1
    if prime[1] == 8:
        prime[1] = 0
        prime[0] += 1

性能对比测试结果 #

• 自研模30轮分解实现：生成100万以内素数耗时约7.3秒
• 基础模2轮埃氏筛（纯Python循环）：同任务耗时约2秒，对应实现代码如下：

def normal_sieve(number):
    primes = np.ones(number,dtype=bool)
    for factor in range(3,math.isqrt(number)+1,2):
        if primes[factor - 1]:
            for i in range(factor*factor,number+1,2*factor):
                    primes[i - 1] = False
    return [i + 1 for i in range(2, number, 2) if primes[I]]

• 移除分段逻辑后的轮分解实现：耗时降至约0.05秒，但仍存在明显优化空间
• 优化版numpy模6轮实现：同任务仅需0.003秒，性能差距极大，对应代码如下：

def primesfrom2to(n):
    """ Input n>=6, Returns a array of primes, 2 <= p < n """
    sieve = numpy.ones(n//3 + (n%6==2), dtype=bool)
    for i in range(1,int(n**0.5)//3+1):
        if sieve[i]:
            k=3*i+1|1
            sieve[       k*k//3     ::2*k] = False
            sieve[k*(k-2*(i&1)+4)//3::2*k] = False
    return numpy.r_[2,3,((3*numpy.nonzero(sieve)[0][1:]+1)|1)]

• primesieve库：仅需约1分钟即可完成1000亿以内的素数计数，已知其基于C++开发，希望了解除语言特性外的核心提速逻辑。

待解答疑问 #

• 现有轮分解实现性能远低于上述参考方案的根本原因
• 如何正确实现基于轮分解的高性能素数生成器

内容的提问来源于stack exchange，提问作者creeper_the_cat