python - 矢量化 itertools combination_with

正如评论中提到的，没有办法直接combinations_with_replacement()从itertools库中“矢量化”调用（使用 Numba CUDA）。Numba CUDA不能那样工作。

但是，我相信应该可以使用 Numba CUDA 生成等效的结果数据集，在itertools某些情况下，它的运行速度似乎比库函数快。我想可能有很多方法可以做到这一点，我没有声称我描述的方法在任何方面都是最佳的。它当然不是，当然可以让它跑得更快。itertools然而，根据我的测试，即使是这种不太优化的方法也可以在 V100 GPU 上运行比 python 快约 10 倍的特定测试用例。

作为背景，我认为this和this（或同等材料）是必不可少的阅读材料。

综上所述，有选择、有替换的n项目组合数的公式由下式给出：k

(n-1+k)!
--------   (Formula 1)
(n-1)!k!

在下面的代码中，我将上述计算封装在count_comb_with_repl（设备）和host_count_comb_with_repl（主机）函数中。事实证明，我们可以使用这个基本计算，对和进行级联较小的选择序列n，k以驱动整个计算过程来计算仅给定最终结果数组的索引的组合。为了形象化我们正在做的事情，有一个简单的例子会有所帮助。让我们以 3 个项目和 3 个选项为例。从零开始索引项目，可能性数组如下所示：

n = 3, k = 3
index    choices  first digit calculation
0        0,0,0    -----------------
1        0,0,1
2        0,0,2
3        0,1,1     equivalent to n = 3, k = 2
4        0,1,2
5        0,2,2    -----------------
6        1,1,1    -----------------
7        1,1,2     equivalent to n = 2, k = 2
8        1,2,2    -----------------
9        2,2,2     equivalent to n = 1, k = 2

上面列表的长度是通过将n= 3 和k= 3 的值代入公式 1 得出的。理解我提出的方法的关键观察是，仅给定索引来计算选择结果的第一位，我们可以计算0 和 1 之间的分界点，例如通过观察考虑到第一选择索引为 0 的结果，通过将n= 3 和k= 2 的值代入公式 1 来给出该范围的长度。因此，如果我们给定的索引小于这个值 (6) 那么我们知道第一个数字是 0。如果它大于这个值那么我们知道第一个数字是 1 或 2，并且通过适当的偏移我们可以重新计算下一个范围（对应于第一个数字1) 看看我们的指数是否在这个范围内。

一旦我们知道了第一个数字，我们就可以重复这个过程（使用适当的列表缩减和偏移）来找到下一个数字，下一个数字，等等。

这是实现上述方法的python代码。正如我所提到的，对于一个测试用例，n=260它k=4 在我的 V100 上运行不到 3 秒。

$ cat t2.py
from numba import cuda,jit
import numpy as np

@cuda.jit(device=True)
def get_next_count_comb_with_repl(n,k,prev):
    return int(round((prev*(n))/(n+k)))

@cuda.jit(device=True)
def count_comb_with_repl(n,k):
    mymax = max(n-1,k)
    ans = 1.0
    cnt = 1
    for i in range(mymax+1, n+k):
        ans = ans*i/cnt
        cnt += 1
    return int(round(ans))
#intended to be identical to the previous function
#I just need a version I can call from host code
def host_count_comb_with_repl(n,k):
    mymax = max(n-1,k)
    ans = 1.0
    cnt = 1
    for i in range(mymax+1, n+k):
        ans = ans*i/cnt
        cnt += 1
    return int(round(ans))

@cuda.jit(device=True)
def find_first_digit(n,k,i):
    psum = 0
    count = count_comb_with_repl(n, k-1)
    if (i-psum) < count:
        return 0,psum
    psum += count
    for j in range(1,n):
        count = get_next_count_comb_with_repl(n-j,k-1,count)
        if (i-psum) < count:
            return j,psum
        psum += count
    return -1,0 # error

@cuda.jit
def kernel_count_comb_with_repl(n, k, l, r):
    for i in range(cuda.grid(1), l, cuda.gridsize(1)):
        new_ll = n
        new_cc = k
        new_i = i
        new_digit = 0
        for j in range(k):
            digit,psum = find_first_digit(new_ll, new_cc, new_i)
            new_digit += digit
            new_ll -= digit
            new_cc -= 1
            new_i  -= psum
            r[i+j*l] = new_digit

combo_count = 4
ll = 260
cl = host_count_comb_with_repl(ll, combo_count)
print(cl)
# bug if cl > 2G
if cl < 2**31:
    my_dtype = np.uint8
    if ll > 255:
        my_dtype = np.uint16
    r   = np.empty(cl*combo_count, dtype=my_dtype)
    d_r = cuda.device_array_like(r)
    block = 256
    grid = (cl//block)+1
    #grid = 640
    kernel_count_comb_with_repl[grid,block](ll, combo_count, cl, d_r)
    r = d_r.copy_to_host()
    print(r.reshape(combo_count,cl))
$ time python t2.py
194831715
[[  0   0   0 ... 258 258 259]
 [  0   0   0 ... 258 259 259]
 [  0   0   0 ... 259 259 259]
 [  0   1   2 ... 259 259 259]]

real    0m2.212s
user    0m1.110s
sys     0m1.077s
$

（上面的测试用例：n=260，k = 4，在我的系统上使用 OP 的代码需要大约 30 秒。）

这应该被认为是一个想法的草图。我没有声称它没有缺陷。这种类型的问题会迅速耗尽 GPU 上的内存（对于足够大的n和/或的选择k），您唯一的迹象可能是来自 numba 的粗略的内存不足错误。

是的，上面的代码不会产生aathrough的连接，jz但这只是使用结果的索引练习。您可以根据需要使用结果索引来索引项目数组，将 0,0,0,1 之类的结果转换为aa, aa, aa,之类的结果ab

这不是全面的性能胜利。他们的 python 方法对于较小的测试用例仍然更快，而较大的测试用例（例如 n = 260，k = 5）将超过 GPU 上的可用内存。

python - 矢量化 itertools combination_with_replacement

问题描述

解决方案

推荐阅读