防止 (GPU) 优化方法（如 gmpy2 和 numba）中的大整数溢出
Preventing overflow of large integers in (GPU) optimized methods such as gmpy2 and numba
我正在尝试在使用 numba 的 JIT 装饰（优化）例程中使用 gmpy2 检查大整数是否是完美平方。这里的例子只是为了说明目的（从理论上讲，这样的方程或者椭圆曲线是可以处理的differently/better）。我的代码似乎溢出了，因为它产生的解决方案并不是真正的解决方案：
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 例程在大约 3 秒内完成。 4 秒。我交给修饰函数的限制不会超过 64 位的无符号整数（这在这里似乎不是问题）。
我读到大整数不能与 numba 的 JIT 优化结合使用（参见示例 here）。
我的问题：
有没有可能在 (GPU) 优化代码中使用大整数？
我现在可以通过以下代码设法避免精度损失：
@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 这个 ammended/fixed 例程不再在 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)
错误结果的真正原因很简单，你忘记将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% 数量的初始候选人。
经过两次筛选后，只剩下很少的数字需要检查。他们可以很容易地 fast-checked 与 gmpy2.is_square().
过滤阶段可以很容易地包装到 Numba 函数中，正如我在下面所做的那样，这个函数可以有额外的 Numba 参数 parallel = True，这将告诉 Numba 自动 运行 所有 Numpy 操作并行所有 CPU 个核心。
在代码中我使用limit = 1 << 30，这表示要检查的所有x的限制，我使用block = 1 << 26，这表示一次检查多少个数字，在并行 Numba 函数。如果您有足够的内存，您可以将 block 设置得更大，以更有效地占用所有 CPU 个核心。大小 1 << 26 的块大约使用大约 1 GB 的内存。
在将我的想法与过滤结合使用 multi-core 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
[]
...............................
我正在尝试在使用 numba 的 JIT 装饰（优化）例程中使用 gmpy2 检查大整数是否是完美平方。这里的例子只是为了说明目的（从理论上讲，这样的方程或者椭圆曲线是可以处理的differently/better）。我的代码似乎溢出了，因为它产生的解决方案并不是真正的解决方案：
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 例程在大约 3 秒内完成。 4 秒。我交给修饰函数的限制不会超过 64 位的无符号整数（这在这里似乎不是问题）。
我读到大整数不能与 numba 的 JIT 优化结合使用（参见示例 here）。
我的问题： 有没有可能在 (GPU) 优化代码中使用大整数？
我现在可以通过以下代码设法避免精度损失：
@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 这个 ammended/fixed 例程不再在 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)
错误结果的真正原因很简单，你忘记将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% 数量的初始候选人。
经过两次筛选后，只剩下很少的数字需要检查。他们可以很容易地 fast-checked 与 gmpy2.is_square().
过滤阶段可以很容易地包装到 Numba 函数中，正如我在下面所做的那样，这个函数可以有额外的 Numba 参数 parallel = True，这将告诉 Numba 自动 运行 所有 Numpy 操作并行所有 CPU 个核心。
在代码中我使用limit = 1 << 30，这表示要检查的所有x的限制，我使用block = 1 << 26，这表示一次检查多少个数字，在并行 Numba 函数。如果您有足够的内存，您可以将 block 设置得更大，以更有效地占用所有 CPU 个核心。大小 1 << 26 的块大约使用大约 1 GB 的内存。
在将我的想法与过滤结合使用 multi-core 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
[]
...............................