如何在Colaboratory中使用Numba？遇“library nvvm not found”错误求助

我之前在Colab里折腾过好几次Numba的CUDA配置，这个"library nvvm not found"的坑确实容易踩，给你几个亲测有效的解决步骤：

• 先安装对应版本的CUDA工具包
  nvvm是CUDA工具链里的核心组件，Colab默认环境往往没把它的路径配置到位。先运行下面的命令安装兼容性最好的CUDA工具包版本：

  !apt update && apt install -y cuda-toolkit-11-8
  

  选11.8是因为它和当前Colab的Python版本、GPU环境适配性拉满，别随便换更高版本，容易触发其他兼容问题。

• 手动指定Numba需要的环境变量
  安装完工具包后，得明确告诉Numba去哪找nvvm和libdevice文件，运行这段Python代码设置环境变量：

  import os
  # 设置nvvm库的具体路径
  os.environ['NUMBAPRO_NVVM'] = '/usr/lib/nvidia-cuda-toolkit/nvvm/lib64/libnvvm.so'
  # 设置libdevice的目录路径
  os.environ['NUMBAPRO_LIBDEVICE'] = '/usr/lib/nvidia-cuda-toolkit/nvvm/libdevice/'
  

  要是找不到这个默认路径，你可以用!find / -name libnvvm.so命令在Colab里搜索实际的文件位置，再替换掉上面的路径就行。

• 测试配置是否生效
  写一段极简的Numba CUDA代码验证一下：

  from numba import cuda
  import numpy as np
  
  @cuda.jit
  def add_kernel(x, y, result):
      idx = cuda.grid(1)
      if idx < x.size:
          result[idx] = x[idx] + y[idx]
  
  # 准备测试数据
  x = np.array([1,2,3,4])
  y = np.array([5,6,7,8])
  result = np.zeros_like(x)
  
  # 配置CUDA线程与块的数量
  threads_per_block = 2
  blocks_per_grid = (x.size + threads_per_block - 1) // threads_per_block
  
  # 执行核函数
  add_kernel[blocks_per_grid, threads_per_block](x, y, result)
  print(result)
  

  如果输出[ 6 8 10 12]，说明问题已经彻底解决啦！

• 额外小提示
  要是设置完环境变量还是报错，试试重启Colab的运行时（点击顶部菜单栏的「运行时」→「重启运行时」），有时候环境变量需要重启才能完全生效。

内容的提问来源于stack exchange，提问作者Algis