CUDA深度调研

Eager

整体

┌─────────────────────────────────────────────────────────────────────────────┐
│  Python 层                                                                   │
│  ─────────────────────────────────────────────────────────────────────────── │
│  1. 用户代码: y = torch.matmul(x, w) + b                                     │
│     ↓                                                                        │
│  2. Python C-API / pybind11 绑定层                                            │
│     torch/csrc/autograd/python_variable.cpp                                  │
│     ↓                                                                        │
└─────────────────────────────────────────────────────────────────────────────┘
                                    ↓
┌─────────────────────────────────────────────────────────────────────────────┐
│  C++ ATen 层 (A Tensor Library)                                              │
│  ─────────────────────────────────────────────────────────────────────────── │
│  3. ATen 原生函数: at::native::matmul(x, w)                                  │
│     aten/src/ATen/native/LinearAlgebra.cpp                                   │
│     ↓                                                                        │
│  4. TensorImpl: 检查 device/dtype/stride，创建输出 Tensor                     │
│     c10/core/TensorImpl.h                                                    │
│     ↓                                                                        │
└─────────────────────────────────────────────────────────────────────────────┘
                                    ↓
┌─────────────────────────────────────────────────────────────────────────────┐
│  Dispatcher 动态分发层 (核心枢纽)                                             │
│  ─────────────────────────────────────────────────────────────────────────── │
│  5. c10::Dispatcher::call(op, args...)                                        │
│     aten/src/ATen/core/dispatch/Dispatcher.h                                  │
│     ↓                                                                        │
│  6. 根据 Tensor 的 device 类型动态路由:                                       │
│     - CPU: 调用 at::native::cpu_kernel                                        │
│     - CUDA: 调用 at::native::cuda_kernel (或 at::cuda::... )                  │
│     - 自定义 Backend (XPU, MPS, PrivateUse1 等)                              │
│     ↓                                                                        │
└─────────────────────────────────────────────────────────────────────────────┘
                                    ↓
┌─────────────────────────────────────────────────────────────────────────────┐
│  后端执行层 (以 CUDA 为例)                                                    │
│  ─────────────────────────────────────────────────────────────────────────── │
│  7. CUDA 实现层: at::native::matmul_cuda                                     │
│     aten/src/ATen/native/cuda/LinearAlgebra.cu                               │
│     ↓                                                                        │
│  8. 选择算法/Kernel: 根据 shape 选择 cuBLAS / CUTLASS / 手写 CUDA Kernel       │
│     - 大矩阵 → cuBLAS (cublasSgemm / cublasLtMatmul)                         │
│     - 小矩阵 → 手写 CUDA kernel / CUTLASS                                    │
│     ↓                                                                        │
│  9. 启动 CUDA Kernel: <<<grid, block, shared_mem, stream>>>                  │
│     通过 CUDA Runtime API (cudaLaunchKernel) 下发到 GPU                       │
│     ↓                                                                        │
│  10. GPU SASS 指令执行                                                       │
│      GPU 的 SM (Streaming Multiprocessor) 执行实际的乘加运算                   │
│      ↓                                                                       │
│  11. 结果通过 PCIe/NVLink 传回 Host Memory，封装成新的 Tensor 返回 Python     │
└─────────────────────────────────────────────────────────────────────────────┘

PyTorch 的数据绑定

Python 的 Tensor 对象本身不存储任何数组数据，它只存储一个 C++ 指针 cdata。所有数据都在 C++ 侧的 TensorImpl → StorageImpl 里

一个具体例子：加法运算的数据流

# Python
c = a + b

# 具体计算过程
1. Python 传入 a, b
   └─ THPVariable_unpack(a) → at::Tensor(a) [共享 StorageImpl A]
   └─ THPVariable_unpack(b) → at::Tensor(b) [共享 StorageImpl B]

2. C++ Dispatcher 调用 at::add(a, b)
   └─ 可能分配新的 StorageImpl C 存放结果
   └─ 返回新的 at::Tensor(c) [指向 StorageImpl C]

3. 包装回 Python
   └─ new THPVariable(cdata = TensorImpl_C)
   └─ Python 拿到 torch.Tensor(c) [共享 StorageImpl C]

PyTorch 的函数绑定

Python: torch.matmul(a, b)
    │
    ▼
[Level 1] Python 包装层 (torch/_C/__init__.pyi / _VariableFunctions.pyi)
    │         提供类型提示和 API 签名
    ▼
[Level 2] C++ Python 绑定层 (python_torch_functions.cpp / python_variable_methods.cpp)
    │         pybind11 / Python C-API 注册
    │         PythonArgParser 解析参数
    ▼
[Level 3] ATen Dispatcher (Dispatcher.cpp)
    │         根据 Tensor 的 device/dtype/layout 分发
    ▼
[Level 4] Native Kernel (cpu/ 或 cuda/ 目录下的 .cpp/.cu)
    │         实际计算逻辑
    ▼
硬件执行

Level 1

PyTorch 的核心算子（add, matmul, conv2d 等）都是从 C++ 扩展模块 torch._C 导出到 Python 命名空间的。

# torch/__init__.py (简化)from torch._C import _VariableFunctions as VF
matmul = VF.matmul  # ← 直接来自 C++ 扩展模块

Level 2

2.1 python模块/函数与C++函数绑定

PyTorch 使用 Python C-API（不是 pybind11） 手动注册函数。

核心任务在于将python的函数绑定到C++中ATen库的函数。绑定的完整案例见附录中《Python C-API案例》

// 1. 定义 Python C-API 的包装函数
static PyObject * THPVariable_add(PyObject* self, PyObject* args, PyObject* kwargs) {
    // PythonArgParser 负责把 Python 参数解析成 C++ 对象
    static PythonArgParser parser({
        "add(Tensor input, Tensor other, *, Scalar alpha=1)",
        "add(Tensor input, Scalar other, Scalar alpha=1)",
    });
    
    // 解析参数
    auto parsed = parser.parse(args, kwargs, /*traceable=*/true);
    
    // 提取 C++ Tensor
    auto self_ = parsed.tensor(0);      // 第一个 Tensor 参数
    auto other_ = parsed.tensor(1);     // 第二个参数
    auto alpha = parsed.scalar(2, 1);   // alpha 默认值 1
    
    // 调用 ATen C++ API
    auto result = at::add(self_, other_, alpha);
    
    // 把 C++ Tensor 包装回 Python
    return wrap(std::move(result));
}

2.2 Tensor 方法的绑定

类方法使用_PyTypeObject创建并绑定_

// 在 python_variable.cpp 的类型定义中
PyTypeObject THPVariableType = {
    PyVarObject_HEAD_INIT(NULL, 0)
    .tp_name = "torch._C.TensorBase",
    .tp_methods = variable_methods,  // ← 包含 add_, matmul_, backward 等
    // ...
};

并使用PyType_Ready+PyModule_AddObject初始化。

// 第一步：调用 PyType_Ready 激活类型
// 这相当于类的“构造函数”，它会补全元类型(ob_type)、MRO、继承的方法等
if (PyType_Ready(&THPVariableType) < 0) {
    return NULL; // 初始化失败
}

// 第二步：将准备好的类挂载到模块上
// 这相当于把类放进模块的命名空间，让 Python 能找到它
PyModule_AddObject(module, "_TensorBase", (PyObject *)&THPVariableType);

Level 3

ATen Dispatcher —— 从 at::add 到具体 Kernel

Tensor实体

at::Tensor = c10::TensorBasec10::TensorImpl

1）PyTorch 三层核心的最底层

• c10/：最底层基础设施，不含任何算子实现，只定义“张量长什么样、在哪儿、怎么分发”

• aten/（ATen = A Tensor）：张量算子库（add/matmul/conv 的真正 kernel 在这）

• torch/：Python 绑定、autograd、nn 模块

2）c10::TensorImpl 是 struct C10_API TensorImpl : public c10::intrusive_ptr_target，它不是数据容器，而是“描述一块内存怎么被看成张量”的元数据对象。核心字段（在 c10/core/TensorImpl.h）

ATen Dispatcher

// aten/src/ATen/native/BinaryOps.cpp (简化)
Tensor add(const Tensor& self, const Tensor& other, const Scalar& alpha) {
    // 1. 构造 DispatchKeySet（根据 tensor 的 device、dtype、layout 等）
    auto dispatchKeySet = at::DispatchKeySet(c10::DispatchKey::Add) 
                          | self.key_set() 
                          | other.key_set();
    
    // 2. 进入 Dispatcher，查找对应的 kernel
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add", "Tensor")
        .typed<Tensor (const Tensor&, const Tensor&, const Scalar&)>();
    
    return op.callWithDispatchKeySet(dispatchKeySet, self, other, alpha);
}

Dispatcher 维护一张全局算子表，每个算子（如 aten::add.Tensor）对应一个 DispatchKey → Kernel 函数 的映射表：

DispatchKey	Kernel 位置	说明
CPU	aten/src/ATen/native/cpu/	CPU 通用实现
CUDA	aten/src/ATen/native/cuda/	CUDA kernel
Autograd	torch/csrc/autograd/generated/	自动求导包装
Functionalize	aten/src/ATen/	函数化变换
PythonTLSSnapshot	—	Python 自定义算子

分发逻辑：

// Dispatcher.cpp
const KernelFunction& Dispatcher::dispatch(
    const OperatorHandle& op, 
    DispatchKeySet ks
) {
    // 1. 计算最终有效的 dispatch key
    // 优先级：Autograd > Functionalize > CUDA > CPU > CompositeImplicitAutograd
    auto dispatchKey = ks.highestPriorityTypeId();
    
    // 2. 从算子表中查找对应 kernel
    return op.operatorDef_->op.registeredKernels_[dispatchKey];
}

Level 4

Native Kernel —— 真正的计算

// aten/src/ATen/native/cuda/BinaryAddSubKernel.cu
void add_kernel_cuda(TensorIteratorBase& iter, const Scalar& alpha) {
    // 1. 从 iter 中获取输入输出数据的指针
    // 这些指针就是 StorageImpl.data_ptr_，与 Python Tensor 共享内存
    
    // 2. 启动 CUDA kernel
    add_kernel<<<grid, block>>>(
        (float*)iter.data_ptr(0),  // out
        (float*)iter.data_ptr(1),  // a
        (float*)iter.data_ptr(2),  // b
        alpha.to<float>()
    );
}

全链路

Python: torch.add(a, b)
    │
    ▼  [Python C-API 调用]
torch._C._VariableFunctions.add (python_torch_functions.cpp)
    │
    ▼  [PythonArgParser 解析参数]
提取 at::Tensor(a), at::Tensor(b)
    │
    ▼  [at::add() C++ API]
aten/src/ATen/native/BinaryOps.cpp
    │
    ▼  [Dispatcher 分发]
c10::Dispatcher::call("aten::add.Tensor")
    │
    ├── requires_grad=True? ──→ Autograd Kernel ──→ 记录计算图 ──┐
    │                                                              │
    └── requires_grad=False? ──→ 直接分发                         │
    │                                                              │
    ▼  [Device/Backend 分发]                                    │
DispatchKey: CUDA / CPU / XPU                                    │
    │                                                              │
    ▼  [Native Kernel]                                           │
aten/src/ATen/native/cuda/BinaryAddSubKernel.cu                  │
    │                                                              │
    ▼  [实际计算]                                                │
CUDA Kernel 读写 StorageImpl.data_ptr_ ◄───────────────────────────┘
    │
    ▼
结果 Tensor (新 StorageImpl 或复用)
    │
    ▼  [wrap 回 Python]
新的 THPVariable → Python torch.Tensor

附录一：案例

Python C-API案例

用纯 Python C-API（PyMethodDef 方法表 + PyInit_xxx 模块初始化）写一个 C++ 加法扩展，Python 侧 import 后直接调用。这是最“原生”的方式，不依赖 ctypes / pybind11。

核心数据结构与方法

1.python对象在C中的表示：static PyObject

static PyObject：每一个python对象，在C层面都是一个PyObject，这里定义了一个新的python对象，并使用指针static PyObject*指向对象所在内存位置。

2.方法表（method table）：struct PyMethodDef

PyMethodDef 是 CPython 的一个结构体，定义在 Python.h 中，简化版如下：

struct PyMethodDef {
    const char* ml_name;   // Python 中的函数名
    PyCFunction ml_meth;   // C 函数指针
    int         ml_flags;  // 调用方式（METH_VARARGS 等）
    const char* ml_doc;   // docstring
};

每一个数组元素，就对应一个 Python 可调用****函数。最后一行 {NULL, NULL, 0, NULL}是整张映射表的哨兵（sentinel），用来标记数组结束。因为CPython 会遍历这个数组，直到遇到全 NULL。

3.模块定义结构体：static struct PyModuleDef

它定义了一个 模块对象模板，告诉 Python：

“我要创建一个名为 cppmath 的模块，它有哪些函数、文档是什么、内存怎么管理。”

4.模块初始化

// 模块初始化函数，名字必须是 PyInit_<模块名>
PyMODINIT_FUNC PyInit_cppmath(void) {
    return PyModule_Create(&cppmath_module);
}

编译流程

1.import路径检索

.so 扩展模块和 .py 文件走同一套 sys.path 搜索规则，只要把它放在 sys.path 里某个目录下，且文件名是 cppmath.cpython-311-x86_64-linux-gnu.so，就能 import cppmath

2.手动编译

g++ -shared -fPIC -O2 \
    $(python3-config --includes) \
    cppmath.cpp -o cppmath$(python3-config --extension-suffix)

3.利用python库工具自动编译

from setuptools import setup, Extension
setup(ext_modules=[Extension("cppmath", ["cppmath.cpp"])])

python setup.py build_ext --inplace

CUDA 简单算子案例

该案例实现了一个向量加法的cuda kernel，并编译成python lib。kernel可以使用import引入，并传入torch.tensor调用，执行在GPU侧的向量加法计算。项目结构如下：

cuda kernel

vector_add.cu

#include "vector_add.h"
#include <cuda_runtime.h>

extern "C" {

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

void launch_vector_add(
    const float* a,
    const float* b,
    float* c,
    int n,
    cudaStream_t stream
) {
    int threads = 256;
    int blocks = (n + threads - 1) / threads;
    vector_add_kernel<<<blocks, threads, 0, stream>>>(a, b, c, n);
}

}

vector_add.h

#ifndef VECTOR_ADD_H
#define VECTOR_ADD_H

#ifdef __cplusplus
extern "C" {
#endif

void launch_vector_add(
    const float* a,
    const float* b,
    float* c,
    int n,
    cudaStream_t stream
);

#ifdef __cplusplus
}
#endif

#endif

C-API绑定

cu_module.c

// cu_module.c
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <cuda_runtime.h>
#include "vector_add.h"

/* 声明 CUDA 侧函数 */
void launch_vector_add(
    const float* a,
    const float* b,
    float* c,
    int n,
    cudaStream_t stream
);

/* 从 __cuda_array_interface__ 提取 device ptr */
static void* get_cuda_ptr(PyObject* obj) {
    PyObject* iface = PyObject_GetAttrString(obj, "__cuda_array_interface__");
    if (!iface) return NULL;

    PyObject* data = PyDict_GetItemString(iface, "data");
    if (!data) return NULL;

    PyObject* ptr_obj = PyTuple_GetItem(data, 0);  // (ptr, readonly)
    void* ptr = (void*)PyLong_AsUnsignedLongLong(ptr_obj);

    Py_DECREF(iface);
    return ptr;
}

static PyObject* py_vector_add(PyObject* self, PyObject* args) {
    PyObject *a_obj, *b_obj, *c_obj;

    if (!PyArg_ParseTuple(args, "OOO", &a_obj, &b_obj, &c_obj)) {
        return NULL;
    }

    float* a = (float*)get_cuda_ptr(a_obj);
    float* b = (float*)get_cuda_ptr(b_obj);
    float* c = (float*)get_cuda_ptr(c_obj);

    if (!a || !b || !c) {
        PyErr_SetString(PyExc_RuntimeError, "Input must be CUDA tensors");
        return NULL;
    }

    Py_ssize_t n = PySequence_Length(a_obj);
    cudaStream_t stream = 0;  // 默认 stream

    launch_vector_add(a, b, c, (int)n, stream);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        PyErr_SetString(PyExc_RuntimeError, cudaGetErrorString(err));
        return NULL;
    }

    Py_RETURN_NONE;
}

static PyMethodDef LabMethods[] = {
    {"vector_add", py_vector_add, METH_VARARGS, "Zero-copy vector add (CUDA)"},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef labmodule = {
    PyModuleDef_HEAD_INIT,
    "lab",
    "CUDA zero-copy lab module",
    -1,
    LabMethods
};

PyMODINIT_FUNC PyInit_lab(void) {
    return PyModule_Create(&labmodule);
}

编译

setup.py

# ============================================================
# setup.py
# 用法:
#   python setup.py build_ext --inplace
#   pip install -e .
# ============================================================

import os
import sys

# ---------- 探测 torch ----------
def have_torch_cuda_ext():
    try:
        import torch
        from torch.utils.cpp_extension import CUDAExtension  # noqa: F401
        return torch.cuda.is_available()
    except Exception:
        return False

# ---------- 路径探测 ----------
def find_python_include():
    import sysconfig
    return sysconfig.get_path("include")

def find_numpy_include():
    try:
        import numpy as np
        return np.get_include()
    except ImportError:
        return None

CUDA_HOME = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or "/usr/local/cuda"

SOURCES = [
    "vector_add.cu",
    "cu_module.c",
]

# ============================================================
# 路径 1：有 torch → 用 torch CUDAExtension（强烈推荐）
# ============================================================
import torch
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

extra_compile_args = {
    "cxx": ["-O3", "-Wall", "-std=c++17"],
    "nvcc": [
        "-O3", "--use_fast_math", "-std=c++17",
        "-Xcompiler", "-O3,-Wall",
        "-arch=sm_89",
    ],
}

inc_dirs = []
py_inc = find_python_include()
if py_inc: inc_dirs.append(py_inc)
np_inc = find_numpy_include()
if np_inc: inc_dirs.append(np_inc)
inc_dirs.append(os.path.join(CUDA_HOME, "include"))

ext_modules = [
    CUDAExtension(
        name="lab",
        sources=SOURCES,
        include_dirs=inc_dirs,
        library_dirs=[os.path.join(CUDA_HOME, "lib64")],
        libraries=["cudart"],
        extra_compile_args=extra_compile_args,
    )
]

from setuptools import setup
setup(
    name="lab",
    version="0.1.0",
    description="CUDA vector_add Python extension (sm_89, RTX 4090)",
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension},
    python_requires=">=3.8",
    install_requires=["numpy"],
)
sys.exit(0)

Note：C/C++代码编译部分由BuildExtension、CUDAExtension代工完成。

调用

import torch
import lab
import time

N = 10_000_000

a = torch.rand(N, dtype=torch.float32, device="cuda")
b = torch.rand(N, dtype=torch.float32, device="cuda")
c = torch.empty_like(a)

# warmup
lab.vector_add(a, b, c)
torch.cuda.synchronize()

start = time.time()
lab.vector_add(a, b, c)
torch.cuda.synchronize()
ms = (time.time() - start) * 1000

print(f"GPU time: {ms:.3f} ms")
print(f"Correct: {torch.allclose(c, a + b)}")

附录二：理论

编译

源代码(.c/.cpp)
   ↓
Clang（前端）
   ↓
LLVM IR
   ↓
LLVM 优化器 + 后端
   ↓
机器码(可执行文件 / 目标文件)

GPU微架构：SM

组成结构

SM（Streaming Multiprocessor，流式多处理器） 是 NVIDIA GPU 架构中的核心计算单元，是 GPU 并行计算能力的基础模块。一个典型 SM 包含以下关键组件：

CUDA Cores（CUDA 核心）
最基础的整数/浮点运算单元（做加减乘、逻辑判断），数量最多（比如 Ampere SM 有 64 个 FP32 CUDA Core）。
Tensor Cores（张量核心）
专为矩阵运算设计（深度学习的核心操作），能在一个时钟周期内完成 4x4 矩阵乘加，速度远超 CUDA Core。
SFU（Special Function Unit，特殊函数单元）
处理三角函数、指数、对数等复杂数学函数。
LD/ST（Load/Store 单元）
负责从显存/GPU 缓存读写数据。
Warp Scheduler（线程束调度器）
GPU 以 Warp（线程束，通常 32 个线程） 为单位调度，调度器决定哪个 Warp 执行、哪些暂停（隐藏延迟）。
寄存器文件（Register File）
给 SM 内线程提供高速暂存，容量直接影响并发线程数（比如 Ampere SM 有 256KB 寄存器）。
共享内存（Shared Memory）/ L1 Cache
低延迟片上内存，供 SM 内所有线程块（Thread Block）共享数据，避免反复读显存。

运行时

关键词：grid、block、thread、warp

当你用 CUDA 写程序时，会把任务拆成线程网格（Grid）→ 线程块（Block）→ 线程（Thread****）：

• 一个线程块会被分配到一个 SM 上执行（不会跨 SM 拆分）；

• SM 把线程块里的线程按 32 个一组切成 Warp，由 Warp Scheduler 轮流喂给 CUDA Core/Tensor Core 算；

• 如果 SM 资源够（寄存器、共享内存足够），它可以同时驻留多个线程块，用“多线程切换”掩盖显存访问延迟。

简单说：SM 数量越多 → 能同时并行的工作组越多；单 SM 算力越强（比如带 Tensor Core）→ 单个任务算得越快。

Block

Block 是“资源容器” 
一个 Block 独占：
一组寄存器（Register File）
一块共享内存（Shared Memory）

warp

warp负责组织block内一定数量的thread进行计算，是block计算任务组织的一个单位结构，一个warp一般含有32个thread。

Grid、Block、Thread

Grid
 └── Block (0..N)
      └── Thread (0..M)

Thread

CUDA thread 的本质 ：
一个线程 = 一次 kernel 函数的串行执行实例
逻辑概念，不是硬件资源
极其廉价：一个 kernel 启动几百万个 thread 很常见

Block

Block 是“硬件感知的分组单位” 
一个 block：
包含若干 threads（最多 1024 / 2048，取决于架构）
固定在一个 SM（Streaming Multiprocessor）上
共享 L1 cache / shared memory
可以 __syncthreads() 同步

存储结构

L1 cache: 显存在每个SM中的cache，对编程透明。如果L1 cache未命中，则去L2 cache访存。

shared memory: 程序员显式管理的片上内存。每个 SM 上有一块 Shared Memory，同一个 Block 的所有线程共享，不同 Block 之间互相不可见。

HBM: CUDA显存。

手写算子

场景一：kernel融合

c = a + b

d = c * s

>> fused_add_mul(a, b, s, d)

场景二：高频小算子