2

我正在尝试在gmpy2使用numba. 此处的示例仅用于说明目的（从理论的角度来看，可以对此类方程或椭圆曲线进行不同/更好的处理）。我的代码似乎溢出了，因为它产生的解决方案不是真正的解决方案：

import numpy as np
from numba import jit
import gmpy2
from gmpy2 import mpz, xmpz

import time
import sys

@jit('void(uint64)')
def findIntegerSolutionsGmpy2(limit: np.uint64):
    for x in np.arange(0, limit+1, dtype=np.uint64):
        y = mpz(x**6-4*x**2+4)
        if gmpy2.is_square(y):
            print([x,gmpy2.sqrt(y),y])

def main() -> int:
    limit = 100000000
    start = time.time()
    findIntegerSolutionsGmpy2(limit)
    end = time.time()
    print("Time elapsed: {0}".format(end - start))
    return 0

if __name__ == '__main__':
    sys.exit(main())

使用limit = 1000000000例程在约内完成。4 秒。我交给修饰函数的限制不会超过 64 位的无符号整数（这在这里似乎不是问题）。

我读到大整数不能与 numba 的 JIT 优化结合使用（请参见此处的示例）。

我的问题： 是否有可能在（GPU）优化代码中使用大整数？

4

2 回答 2

2

错误结果的真正原因很简单，您忘记转换x为mpz，因此将语句x ** 6 - 4 * x ** 2 + 4提升为np.uint64类型并使用溢出计算（因为x在语句中是np.uint64）。修复很简单，只需添加x = mpz(x)：

@jit('void(uint64)', forceobj = True)
def findIntegerSolutionsGmpy2(limit: np.uint64):
    for x in np.arange(0, limit+1, dtype=np.uint64):
        x = mpz(x)
        y = mpz(x**6-4*x**2+4)
        if gmpy2.is_square(y):
            print([x,gmpy2.sqrt(y),y])

你也可能注意到我添加了forceobj = True，这是为了在启动时抑制 Numba 编译警告。

在此修复后，一切正常，您不会看到错误的结果。

如果您的任务是检查表达式是否给出严格的平方，那么我决定为您发明并实施另一种解决方案，代码如下。

它的工作原理如下。您可能会注意到，如果一个数是平方数，那么它也是任意数的平方模数（取模数是x % N运算）。

我们可以取任意数，例如一些素数的乘积K = 2 * 2 * 3 * 5 * 7 * 11 * 13 * 17 * 19。现在我们可以制作一个简单的过滤器，计算所有以 K 为模的正方形，在位向量内标记这个正方形，然后检查在这个过滤器位向量中有哪些以 K 为模的数字。

上面提到的过滤器 K（素数的乘积）只留下 1% 的正方形候选。我们也可以做第二阶段，对其他素数应用相同的过滤器，例如K2 = 23 * 29 * 31 * 37 * 41。这将把它们过滤掉 3%。总的来说，我们将有1% * 3% = 0.03%剩余的初始候选人数量。

经过两次过滤后，只有少数数字需要检查。可以使用 轻松快速检查它们gmpy2.is_square()。

过滤阶段可以很容易地包装到 Numba 函数中，就像我在下面所做的那样，这个函数可以有额外的 Numba 参数parallel = True，这将告诉 Numba 自动在所有 CPU 内核上并行运行所有 Numpy 操作。

在我使用的代码中，这表示要检查limit = 1 << 30的所有数量的限制，而我使用，这表示在并行 Numba 函数中一次要检查多少个数字。如果您有足够的内存，您可以设置更大以更有效地占用所有 CPU 内核。大小的块大约使用大约 1 GB 的内存。xblock = 1 << 26block1 << 26

在使用我的过滤和使用多核 CPU 的想法后，我的代码解决了与您相同的任务，速度提高了数百倍。

在线尝试！

import numpy as np, numba

@numba.njit('u8[:](u8[:], u8, u8, u1[:])', cache = True, parallel = True)
def do_filt(x, i, K, filt):
    x += i; x %= K
    x2 = x
    x2 *= x2;     x2 %= K
    x6 = x2 * x2; x6 %= K
    x6 *= x2;     x6 %= K
    x6 += np.uint64(4 * K + 4)
    x2 <<= np.uint64(2)
    x6 -= x2; x6 %= K
    y = x6
    #del x2
    filt_y = filt[y]
    filt_y_i = np.flatnonzero(filt_y).astype(np.uint64)
    return filt_y_i

def main():
    import math
    gmpy2 = None
    import gmpy2
    
    Int = lambda x: (int(x) if gmpy2 is None else gmpy2.mpz(x))
    IsSquare = lambda x: gmpy2.is_square(x)
    Sqrt = lambda x: Int(gmpy2.sqrt(x))
    
    Ks = [2 * 2 * 3 * 5 * 7 * 11 * 13 * 17 * 19,    23 * 29 * 31 * 37 * 41]
    filts = []
    for i, K in enumerate(Ks):
        a = np.arange(K, dtype = np.uint64)
        a *= a
        a %= K
        filts.append((K, np.zeros((K,), dtype = np.uint8)))
        filts[-1][1][a] = 1
        print(f'filter {i} ratio', round(len(np.flatnonzero(filts[-1][1])) / K, 4))
    
    limit = 1 << 30
    block = 1 << 26
    
    for i in range(0, limit, block):
        print(f'i block {i // block:>3} (2^{math.log2(i + 1):>6.03f})')
        x = np.arange(0, min(block, limit - i), dtype = np.uint64)
        
        for ifilt, (K, filt) in enumerate(filts):
            len_before = len(x)
            x = do_filt(x, i, K, filt)
            print(f'squares filtered by filter {ifilt}:', round(len(x) / len_before, 4))
        
        x_to_check = x
        print(f'remain to check {len(x_to_check)}')
        
        sq_x = []
        for x0 in x_to_check:
            x = Int(i + x0)
            y = x ** 6 - 4 * x ** 2 + 4
            if not IsSquare(y):
                continue
            yr = Sqrt(y)
            assert yr * yr == y
            sq_x.append((int(x), int(yr)))
        print('squares found', len(sq_x))
        print(sq_x)
        
        del x

if __name__ == '__main__':
    main()

输出：

filter 0 ratio 0.0094
filter 1 ratio 0.0366
i block   0 (2^ 0.000)
squares filtered by filter 0: 0.0211
squares filtered by filter 1: 0.039
remain to check 13803
squares found 2
[(0, 2), (1, 1)]
i block   1 (2^24.000)
squares filtered by filter 0: 0.0211
squares filtered by filter 1: 0.0392
remain to check 13880
squares found 0
[]
i block   2 (2^25.000)
squares filtered by filter 0: 0.0211
squares filtered by filter 1: 0.0391
remain to check 13835
squares found 0
[]
i block   3 (2^25.585)
squares filtered by filter 0: 0.0211
squares filtered by filter 1: 0.0393
remain to check 13907
squares found 0
[]

...............................
于 2022-01-22T07:18:08.073 回答
1

我现在可以通过以下代码设法避免精度损失：

@jit('void(uint64)')
def findIntegerSolutionsGmpy2(limit: np.uint64):
    for x in np.arange(0, limit+1, dtype=np.uint64):
        x_ = mpz(int(x))**2
        y = x_**3-mpz(4)*x_+mpz(4)
        if gmpy2.is_square(y):
            print([x,gmpy2.sqrt(y),y])

但是通过使用limit = 100000000这个修正/固定的例程不再在4 秒内完成。现在花了912 秒。我们很可能在精度和速度之间存在无法逾越的差距。

使用 CUDA 它变得更快，即5 分钟（机器有 128GB RAM，Intel Xeon CPU E5-2630 v4，2.20GHz 处理器和两个 Tesla V100 类型的显卡，每个显卡都有 16GB RAM），但我得到的正确结果也是错误的结果再次。

%%time
from numba import jit, cuda
import numpy as np
from math import sqrt

@cuda.jit
def findIntegerSolutionsCuda(arr):
    i=0
    for x in range(0, 1000000000+1):
        y = float(x**6-4*x**2+4)
        sqr = int(sqrt(y))
        if sqr*sqr == int(y):
            arr[i][0]=x
            arr[i][1]=sqr
            arr[i][2]=y
            i+=1

arr=np.zeros((10,3))
findIntegerSolutionsCuda[128, 255](arr)

print(arr)
于 2022-01-14T11:04:51.823 回答