网上搜了一下，发现是驱动问题。通过Ubuntu自带的工具安装显卡驱动：

还是失败：

$ nvidia-smi
NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running.

最后，通过命令行安装驱动，成功解决这个问题：

$ sudo apt install nvidia-driver-470

检查后发现正常了：

$ nvidia-smi 
Wed Dec  7 22:13:49 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0 N/A |                  N/A |
| N/A   51C    P8    N/A /  N/A |      4MiB /  2004MiB |     N/A      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

测试代码也可以跑了。

测试Python代码

打印ID

准备以下代码：

from numba import cuda
import os

def cpu_print():
    print('cpu print')


@cuda.jit
def gpu_print():
    dataIndex = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    print('gpu print ', cuda.threadIdx.x, cuda.blockIdx.x, cuda.blockDim.x, dataIndex)


if __name__ == '__main__':
    gpu_print[4, 4]()
    cuda.synchronize()
    cpu_print()

这个代码主要有两个函数，一个是用CPU执行，一个是用GPU执行，执行打印操作。关键在于@cuda.jit这个注解，让代码在GPU上执行。运行结果如下：

$ /home/larry/anaconda3/bin/python /home/larry/code/pkslow-samples/python/src/main/python/cuda/print_test.py
gpu print  0 3 4 12
gpu print  1 3 4 13
gpu print  2 3 4 14
gpu print  3 3 4 15
gpu print  0 2 4 8
gpu print  1 2 4 9
gpu print  2 2 4 10
gpu print  3 2 4 11
gpu print  0 1 4 4
gpu print  1 1 4 5
gpu print  2 1 4 6
gpu print  3 1 4 7
gpu print  0 0 4 0
gpu print  1 0 4 1
gpu print  2 0 4 2
gpu print  3 0 4 3
cpu print

可以看到GPU总共打印了16次，使用了不同的Thread来执行。这次每次打印的结果都可能不同，因为提交GPU是异步执行的，无法确保哪个单元先执行。同时也需要调用同步函数cuda.synchronize()，确保GPU执行完再继续往下跑。

查看时间

我们通过这个函数来看GPU并行的力量：

from numba import jit, cuda
import numpy as np
# to measure exec time
from timeit import default_timer as timer


# normal function to run on cpu
def func(a):
    for i in range(10000000):
        a[i] += 1


# function optimized to run on gpu
@jit(target_backend='cuda')
def func2(a):
    for i in range(10000000):
        a[i] += 1


if __name__ == "__main__":
    n = 10000000
    a = np.ones(n, dtype=np.float64)

    start = timer()
    func(a)
    print("without GPU:", timer() - start)

    start = timer()
    func2(a)
    print("with GPU:", timer() - start)

结果如下：