用Numba CUDA编写可加速的二叉树算法

在开始我们的算法编写之前，我们需要考虑清楚一件事儿：在CUDA的世界里，一切皆为矩阵！

CUDA中的存储模式为块（Block）存储，这意味着，我们在声明一个数组时需要指定它的大小以及维度，即：有效的CUDA Block X和CUDA Block Y。

不过好在我们这次的算法是对一个一维的哈希结果数组，因此我们只需要用到numba.cuda.grid()这个函数，将其设置为1，并且分配一个numba.cuda.gridDim.x和numba.cuda.blockDimX进行相乘，这样我们就得到了一个startX和一个gridX。

startX = cuda.grid(1)
gridX = cuda.gridDim.x * cuda.blockDim.x

明确了CUDA存储的基本要义后，我们就可以先放一边，去考虑一下生成一下模拟的数据。
这里我使用了faker这个库生成数据，调用fake.Fake()类下的并调用email()函数生成数据。

from faker import Faker
from tqdm import tqdm
def gendata(n_sample:int):
    gener = Faker()
    with open('mail.txt','w') as doc:
        for i in tqdm(range(n_sample),desc='generating...'):
            doc.write(gener.email()+'\n')

if __name__ == "__main__":
    gendata(300000)

这里我们生成了一个三十万的数据当做备用。
有了数据后我们需要对其进行哈希化处理，以比对他们之间哈希值的大小，进行排序。

这里使用md5算法对每条记录进行加密，并使用int(<hash>,16)，将其转换为整型数据，但是为了防止转换值过大超出C long的最大限制，因此我们需要对其进行缩小处理。

def ConvertString():
    IntegerArray = []
    with open('mail.txt','r') as f:
        readResult = f.readlines()
        for line in readResult:
            IntegerArray.append(int(md5(line.encode('utf-8')).hexdigest(),16) / (1024 ** 32))
    return IntegerArray,readResult

确定好加密算法后，我们来编写主要的部分：CUDA加速判断
还记得我们之前提到的startX和gridX了吗，我们现在就可以通过调用他们来实现CUDA加速了

@cuda.jit
def Sort(IntegerArray,SortResult): 
    tempValue = 0 # 指定当前哈希节点值
    startX = cuda.grid(1) # 启动cuda.grid，设置维度为1
    gridX = cuda.gridDim.x * cuda.blockDim.x # 定义块X
    for i in range(startX,len(IntegerArray),gridX):# 遍历grid和block
        if tempValue == 0:
            tempValue = IntegerArray[i]
        else:
            if tempValue > IntegerArray[i]:
                SortResult[i] = 2
            elif tempValue < IntegerArray[i]:
                SortResult[i] = 1
            elif tempValue == IntegerArray[i]:
                SortResult[i] = 0
        tempValue = IntegerArray[i]

编写好CUDA函数后，我们需要需要编写主函数来激活并调用CUDA内核。
正如前文所说，我们需要分别设置block和grid所使用的的线程数量，因此我们这里使用了(16,16)的大小，考虑到笔者的RTX 2060只有6G的显存数量和并不充裕的CUDA核心数配备，我们也只好缩小GPU的开销了。

def main():
    print('Progress Starting...')
    startTime = time.time()
    IntegerArray,readResult = ConvertString()
    # 将哈希化数据复制到VRAM上
    IntegerArray = cuda.to_device(np.array(IntegerArray,dtype=np.int64))
    # 定义TPD
    blockDim = (16,16)
    gridDim = (16,16)
    emptyArray = cuda.device_array([len(IntegerArray),1])
    # 在调用CUDA函数前需要将blockDim和gridDim的线程数量传入
    # 以激活内核
    SortResult = Sort[blockDim,gridDim](IntegerArray,emptyArray)
    SortResult = emptyArray.copy_to_host()
    with open('result.txt','w') as f:
        for i in range(len(SortResult)):
            if SortResult[i] == 0:
                f.write(f'[LastNode]{readResult[i]} == [CurrentNode]{readResult[i]} \n')
            elif SortResult[i] == 1:
                f.write(f'[LastNode]{readResult[i]} L [CurrentNode]{readResult[i]} \n')
            elif SortResult[i] == 2:
                f.write(f'[LastNode]{readResult[i]} R [CurrentNode]{readResult[i]} \n')
    print("\n\n")
    print('----------------------------------------------------')
    print('----------------------------------------------------')
    print('\n')
    print('Progress Finished!')
    print('Duration: ',round(time.time()-startTime,2),'seconds')
    print('Records: '+str(len(IntegerArray)),'records')

我们执行一下

(rapids-22.04) [root@ElinsLaptop binaryTree]# python binTreeCUDA.py 
Progress Starting...



----------------------------------------------------
----------------------------------------------------


Progress Finished!
Duration:  1.0 seconds
Records: 300000 records

可以看到，CUDA用时不到1秒就完成了所有的流程。不过可能是算法的问题，导致SortResult的返回值全是0，因此后期还需要重新写一下。
好了，又水了一篇文章，又是开摆的一天。