第二章 开发环境搭建：Linux开发环境配置、CUDA Toolkit与cuDNN安装、Triton编译器环境搭建、Docker镜像使用

好，咱们直接进入正题。做AI芯片算子库开发，第一步就是把开发环境给整利索了。我见过太多人，代码写得挺溜，结果环境配了一周还没跑通第一个demo。说白了，环境搭建这事儿，看着琐碎，但坑是真不少。这一章，我就把这几年的经验掰开揉碎讲给你听。

2.1 Linux开发环境配置

我个人习惯用Ubuntu 20.04 LTS或22.04 LTS。为什么？因为NVIDIA的官方驱动和CUDA Toolkit对这两个版本支持最稳。你想想看，要是用个冷门发行版，装个驱动都得自己编译内核模块，何必呢？

装好系统后，第一件事是换源。别用默认的国外源，慢得让人抓狂。我一般用清华或者中科大的镜像站。命令很简单：

# 备份原源
sudo cp /etc/apt/sources.list /etc/apt/sources.list.bak

# 替换为清华源（Ubuntu 22.04为例）
sudo sed -i 's/archive.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list
sudo sed -i 's/security.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list

sudo apt update && sudo apt upgrade -y

然后装一些基础工具。嗯，这里要注意，别一股脑全装，按需来就行：

• 编译工具链：build-essential、cmake、g++、gdb
• Python环境：python3-dev、python3-pip、virtualenv
• 版本控制：git（这个必须装，后面拉代码全靠它）
• 性能分析：perf、valgrind、htop

#!/bin/bash
sudo apt install -y build-essential cmake g++ gdb \
                    python3-dev python3-pip virtualenv \
                    git perf valgrind htop

2.2 CUDA Toolkit与cuDNN安装

CUDA Toolkit是绕不过去的坎。我刚开始学的时候，装CUDA踩过不少坑。最典型的就是版本不匹配——显卡驱动、CUDA版本、cuDNN版本，这三者必须对齐。

先确认你的显卡型号：

lspci | grep -i nvidia

然后去NVIDIA官网下载对应的驱动。我个人建议用runfile安装，别用deb包。为什么？runfile可以让你选择不装OpenGL库，避免和桌面环境冲突。我曾经因为装了deb包导致图形界面起不来，折腾了半天才发现是驱动冲突。

安装CUDA Toolkit的步骤：

1. 下载对应版本的runfile（比如cuda_12.1.0_530.30.02_linux.run）
2. 执行安装：sudo sh cuda_12.1.0_530.30.02_linux.run
3. 安装时取消勾选Driver（如果已经装过驱动）
4. 配置环境变量：

echo 'export PATH=/usr/local/cuda-12.1/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

验证安装：

nvcc --version
nvidia-smi

看到版本号就对了。如果nvidia-smi报错，八成是驱动没装好。

接下来装cuDNN。这个简单，下载tar包解压后复制到CUDA目录就行：

tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda12-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

2.3 Triton编译器环境搭建

Triton是OpenAI出的一个好东西，专门用来写高效的GPU算子。说白了，它让你用Python写CUDA代码，不用手动调优就能达到接近手写CUDA的性能。

安装Triton有两种方式：

• pip安装：适合快速体验
• 源码编译：适合二次开发

先看pip安装：

pip install triton

就这么简单？嗯，前提是你的CUDA环境已经配好了。如果报错说找不到CUDA，检查一下CUDA_HOME环境变量有没有设对。

源码编译稍微麻烦点，但能让你看到Triton的内部实现：

git clone https://github.com/openai/triton.git
cd triton
git checkout release/2.1.0  # 选一个稳定版本
pip install -e python/

编译过程中会下载LLVM依赖，大概要十几分钟。我建议用代理或者提前下好LLVM的预编译包，不然网络不好容易断。

验证安装：

python -c "import triton; print(triton.__version__)"

如果输出版本号，恭喜你，环境搭好了。

2.4 Docker镜像使用

Docker是环境搭建的终极解决方案。你想想看，每次换机器都要重新配一遍环境，多烦人？用Docker，一个镜像搞定所有依赖。

先装Docker：

sudo apt install docker.io
sudo systemctl start docker
sudo systemctl enable docker
sudo usermod -aG docker $USER  # 把当前用户加入docker组，免sudo

然后拉取NVIDIA官方提供的CUDA镜像：

docker pull nvidia/cuda:12.1.0-devel-ubuntu22.04

这个镜像里已经装好了CUDA Toolkit和cuDNN，省去了手动安装的麻烦。启动容器时记得加上--gpus all参数：

docker run --gpus all -it --rm \
  -v /home/user/project:/workspace \
  nvidia/cuda:12.1.0-devel-ubuntu22.04 \
  /bin/bash

这里-v是把宿主机的项目目录挂载到容器里，这样代码修改后容器里能直接看到。

如果你需要Triton环境，可以自己写个Dockerfile：

FROM nvidia/cuda:12.1.0-devel-ubuntu22.04

RUN apt update && apt install -y python3 python3-pip git
RUN pip install triton

WORKDIR /workspace
CMD ["/bin/bash"]

构建镜像：

docker build -t triton-dev:latest .

以后每次开发，直接启动这个容器就行。我个人的工作流是这样的：代码在宿主机上用VS Code写，然后在容器里编译和测试。这样宿主机干干净净，所有依赖都在容器里。

2.5 环境验证与常见问题

环境搭好后，跑一个简单的CUDA程序验证一下：

# vector_add.cu
#include <stdio.h>

__global__ void add(int *a, int *b, int *c, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) c[idx] = a[idx] + b[idx];
}

int main() {
    int n = 1024;
    int *a, *b, *c;
    cudaMalloc(&a, n * sizeof(int));
    cudaMalloc(&b, n * sizeof(int));
    cudaMalloc(&c, n * sizeof(int));
    add<<<(n+255)/256, 256>>>(a, b, c, n);
    cudaDeviceSynchronize();
    printf("Success!\n");
    return 0;
}

编译运行：

nvcc vector_add.cu -o vector_add
./vector_add

看到"Success!"就说明CUDA环境没问题。

常见问题汇总：

好了，环境搭建这部分就讲到这里。下一章我们开始写第一个算子——向量加法。到时候你会看到，同样的功能，用CUDA写和用Triton写，差别有多大。