问题描述
我正在尝试使用numba
在JIT修饰(优化)的例程中使用gmpy2
检查大整数是否为完美平方。这里的示例仅用于说明目的(从理论角度来看,可以不同/更好地处理此类方程或椭圆曲线)。我的代码似乎溢出,因为它产生的解决方案并不是真正的解决方案:
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优化结合使用(例如,请参阅here)。
我的问题: 是否可以在(GPU)优化代码中使用大整数?
推荐答案
错误结果的真正原因很简单,您忘记了将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
,这表示要检查的所有x
的限制,我使用block = 1 << 26
,这表示在并行的Numba函数中,一次要检查多少个数字。如果您有足够的内存,您可以将block
设置得更大,以便更有效地占用所有CPU核心。1 << 26
大小的挡路大约占用1 GB左右的内存。
将我的想法用于过滤并使用多核CPU后,我的代码解决相同任务的速度比您的快几百倍。
Try it online!
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
[]
...............................
这篇关于防止gmpy2和numba等(GPU)优化方法中的大整数溢出的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持跟版网!