借助Numba和CUDA，用Python编写你的第一个GPU内核

我们在本文中将使用一个向量加法的常见示例，并使用Numba将简单的CPU代码转换成CUDA内核。向量加法是并行机制的理想例子，因为跨单个索引的加法与其他索引无关。这是完美的SIMD场景，因此所有索引可以同时相加，从而在一次运算中完成向量加法。

Python 速度提升80倍？探究如何用一行代码将你的代码变成GPU猛兽！

GPU非常适合处理需要对不同数据执行相同操作的任务。这种方法名为单指令多数据（SIMD）。与只有几个强大核心的CPU不同，GPU拥有数千个较小的核心，它们可以同时运行这些重复性操作。你会在机器学习中经常看到这种模式，比如在对大型向量进行加法或乘法时，因为每个计算都是独立的。这是使用GPU通过并行机制加快处理任务的理想场景。

英伟达创建了CUDA，以便开发者编写在GPU上而不是CPU上运行的程序。它基于C语言，允许你编写名为内核的特殊函数，这些函数可以同时运行多个操作。问题在于，用C 或C++编写CUDA对初学者来说并不友好，必须处理诸如手动分配内存、协调线程以及理解GPU底层工作原理之类的问题。这可能会让初学者感到不知所措，尤其是如果习惯用 Python编写代码的话。

这时候Numba可以助你一臂之力。它允许借助Python编写CUDA内核，并使用LLVM（低级虚拟机）编译器基础架构，将Python代码直接编译成与CUDA兼容的内核。借助即时（JIT）编译，你可以使用装饰器注释函数，其余所有工作交由Numba来处理。

我们在本文中将使用一个向量加法的常见示例，并使用Numba将简单的CPU代码转换成CUDA内核。向量加法是并行机制的理想例子，因为跨单个索引的加法与其他索引无关。这是完美的SIMD场景，因此所有索引可以同时相加，从而在一次运算中完成向量加法。

请注意，你需要一个CUDA GPU才能遵循本文操作。你可以使用Colab的免费T4 GPU或已安装英伟达工具包和NVCC的本地GPU。

搭建环境并安装Numba

Numba以Python软件包的形式提供，你可以使用pip来安装。此外，我们将使用numpy用于向量运算。使用以下命令搭建Python环境：

python3 -m venv venv
source venv/bin/activate
pip install numba-cuda numpy

CPU上的向量加法

我们举一个简单的向量加法例子。对于两个给定的向量，我们将每个索引对应的值相加以获得最终值。我们将使用numpy生成随机的float32向量，并使用for循环生成最终输出。

import numpy as np 
N = 10_000_000 # 10 million elements 
a = np.random.rand(N).astype(np.float32) 
b = np.random.rand(N).astype(np.float32) 
c = np.zeros_like(a) # Output array 
def vector_add_cpu(a, b, c): 
 """Add two vectors on CPU""" 
 for i in range(len(a)): 
 c[i] = a[i] + b[i]

代码分解如下：

• 初始化两个向量，每个向量包含1000万个随机浮点数。
• 我们还创建一个空向量c来存储结果。
• vector_add_cpu函数只是循环遍历每个索引，并将a和b中的元素相加，将结果存储在c中。

这是一个串行操作；每个加法操作都是一个接一个进行。虽然这种方法运行良好，但它并非最高效的方法，尤其是对大型数据集而言。由于每个加法彼此独立，因此非常适合在GPU上并行执行。

在下一节中，你将看到如何使用Numba转换这个相同的操作以便在GPU上运行。通过将每个元素级加法分布到数千个GPU线程上，我们可以显著加快任务完成速度。

借助Numba在GPU上进行向量加法

现在你将使用Numba定义一个可在CUDA上运行的Python函数，并在Python中执行它。我们在执行相同的向量加法运算，但现在它可以针对Numpy数组的每个索引并行运行，从而提高执行速度。

以下是编写内核的代码：

from numba import config
# Required for newer CUDA versions to enable linking tools. 
# Prevents CUDA toolkit and NVCC version mismatches.
config.CUDA_ENABLE_PYNVJITLINK = 1
from numba import cuda, float32
@cuda.jit
def vector_add_gpu(a, b, c):
 """Add two vectors using CUDA kernel"""
 # Thread ID in the current block
 tx = cuda.threadIdx.x
 # Block ID in the grid
 bx = cuda.blockIdx.x
 # Block width (number of threads per block)
 bw = cuda.blockDim.x
 # Calculate the unique thread position
 position = tx + bx * bw
 # Make sure we don't go out of bounds
 if position < len(a):
 c[position] = a[position] + b[position]
def gpu_add(a, b, c):
 # Define the grid and block dimensions
 threads_per_block = 256
 blocks_per_grid = (N + threads_per_block - 1) // threads_per_block
 # Copy data to the device
 d_a = cuda.to_device(a)
 d_b = cuda.to_device(b)
 d_c = cuda.to_device(c)
 # Launch the kernel
 vector_add_gpu[blocks_per_grid, threads_per_block](d_a, d_b, d_c)
 # Copy the result back to the host
 d_c.copy_to_host(c)
def time_gpu():
 c_gpu = np.zeros_like(a)
 gpu_add(a, b, c_gpu)
 return c_gpu

不妨细述一下上面发生的操作。

理解GPU函数

@cuda.jit装饰器告诉Numba将以下函数视为CUDA内核；这是一个将跨GPU上的多个线程并行运行的特殊函数。在运行时，Numba会将此函数编译成与CUDA兼容的代码，并为你处理C-API转译。

@cuda.jit
defvector_add_gpu(a, b, c):
...

该函数将同时在数千个线程上运行。但我们需要一种方法来确定每个线程应该处理数据的哪个部分。这就是接下来几行代码的作用：

• tx 是线程在其块中的 ID。
• bx 是块在网格中的 ID。
• bw 是块中有多少个线程。

我们将这些数据组合起来计算出独特的位置，该位置告诉每个线程应该添加数组中的哪个元素。请注意，线程和块可能并不总是提供有效的索引，因为它们以2的幂次方进行操作。当向量长度不符合底层架构时，这可能会导致无效索引。因此，我们在执行向量加法之前添加了一个保护条件来验证索引。这可以防止访问数组时出现任何越界运行时错误。

一旦我们知道了这个独特位置，现在可以像在CPU实现中一样添加值。以下代码行与CPU实现一致：

c[position] = a[position] + b[position]

启动内核

gpu_add函数负责进行设置：

• 它定义了要使用的线程和块的数量。你可以尝试块和线程大小的不同值，并在GPU内核中打印输出相应的值。这可以帮助你理解底层GPU索引的工作原理。
• 它将输入数组（a、b 和 c）从CPU内存复制到GPU内存，以便可以在GPU RAM中访问这些向量。
• 它使用vector_add_gpu[blocks_per_grid,threads_per_block]运行GPU内核。
• 最后，它将结果从GPU复制回到c数组，以便我们可以在CPU上访问这些值。

比较实现和潜在加速

我们已有了CPU和GPU版本的向量加法，是时候比较一下它们了。验证结果以及使用CUDA并行机制所能获得的执行提升非常重要。

import timeit
c_cpu = time_cpu()
c_gpu = time_gpu()
print("Results match:", np.allclose(c_cpu, c_gpu))
cpu_time = timeit.timeit("time_cpu()", globals=globals(), number=3) / 3
print(f"CPU implementation: {cpu_time:.6f} seconds")
gpu_time = timeit.timeit("time_gpu()", globals=globals(), number=3) / 3
print(f"GPU implementation: {gpu_time:.6f} seconds")
speedup = cpu_time / gpu_time
print(f"GPU speedup: {speedup:.2f}x")

首先，我们运行两种实现，检查它们的结果是否一致。这对于确保我们的GPU代码正常运行且输出结果与CPU版本一致至关重要。

接下来，我们使用Python内置的timeit模块来测量每个版本的运行时间。我们运行每个函数几次，取平均值以获得可靠的时间。最后，我们计算GPU版本比CPU版本快多少倍。你应该会看到显著的差异，因为GPU可以同时执行多项操作，而CPU在循环中一次只处理一项操作。

以下是Colab上英伟达T4 GPU 的预期输出。请注意，具体的加速效果可能因CUDA版本和底层硬件而异。

Results match: True
CPU implementation: 4.033822 seconds
GPU implementation: 0.047736 seconds
GPU speedup: 84.50x

这个简单的测试有助于展示GPU加速的强大功能，以及它为何对涉及大量数据和并行工作的任务如此有用。

结语

就是这样。你现在已经借助Numba编写了第一个CUDA内核，无需实际编写任何C或CUDA代码。Numba提供了一个简单的接口，可以通过Python使用GPU，这使得Python工程师更容易上手CUDA编程。

现在，你可以使用相同的模板来编写高级CUDA算法，这些算法在机器学习和深度学习中非常流行。如果你发现遵循SIMD范式存在问题，使用GPU来提升执行速度始终是个好主意。

完整的代码可以在Colab笔记本上找到，可以点击此处访问。你可以随意测试并进行一些简单的更改，以更好地理解CUDA索引和执行的内部工作原理。

AI大模型学习福利

作为一名热心肠的互联网老兵，我决定把宝贵的AI知识分享给大家。 至于能学习到多少就看你的学习毅力和能力了 。我已将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

因篇幅有限，仅展示部分资料，需要点击文章最下方名片即可前往获取

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

因篇幅有限，仅展示部分资料，需要点击文章最下方名片即可前往获

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。

因篇幅有限，仅展示部分资料，需要点击文章最下方名片即可前往获

四、AI大模型商业化落地方案

因篇幅有限，仅展示部分资料，需要点击文章最下方名片即可前往获

作为普通人，入局大模型时代需要持续学习和实践，不断提高自己的技能和认知水平，同时也需要有责任感和伦理意识，为人工智能的健康发展贡献力量