feat: 完整中文翻译 maths-cs-ai-compendium（数学·计算机科学·AI 知识大全）
翻译自英文原版 maths-cs-ai-compendium，共 20 章全部完成。 第01章 向量 | 第02章 矩阵 | 第03章 微积分 第04章 统计学 | 第05章 概率论 | 第06章 机器学习 第07章 计算语言学 | 第08章 计算机视觉 | 第09章 音频与语音 第10章 多模态学习 | 第11章 自主系统 | 第12章 图神经网络 第13章 计算与操作系统 | 第14章 数据结构与算法 第15章 生产级软件工程 | 第16章 SIMD与GPU编程 第17章 AI推理 | 第18章 ML系统设计 第19章 应用人工智能 | 第20章 前沿人工智能 翻译说明： - 所有数学公式 $...$ / $$...$$、代码块、图片引用完整保留 - mkdocs.yml 配置中文导航 + language: zh - README.md 已翻译为中文（兼 docs/index.md） - docs/ 目录包含指向各章文件的 symlink - 约 29,000 行中文内容，排除 .cache/ 构建缓存
This commit is contained in:
@@ -0,0 +1,419 @@
|
||||
# 为什么是C++以及ML框架如何工作
|
||||
|
||||
*本书中每一次 `jnp.matmul`、每一次 `torch.nn.Linear`、每一次 `np.dot` 调用，底层都在执行C++和CUDA代码。本文档揭开帷幕：为何ML框架采用这种架构，面向Python工程师的C++快速入门，何时编写自定义C++核函数，以及如何将其绑定到Python——这是连接你所写代码与所运行硬件之间的桥梁。*
|
||||
|
||||
- 你花了15章写Python。你导入了JAX，调用了`jax.grad`，运行了训练循环，构建了模型。一切感觉都像是Python。但事实是：**几乎没有实际计算发生在Python中。**
|
||||
|
||||
- 当你在PyTorch中写 `output = model(input)` 或在JAX中写 `output = jnp.matmul(W, x)` 时，Python几乎什么都不做。它构建一个计算的描述（一个操作图），然后将其交给执行真正工作的C++/CUDA后端。Python是方向盘；C++是引擎。
|
||||
|
||||
## 为什么Python前端搭配C++后端
|
||||
|
||||
- 这种双语言架构的存在是因为Python和C++擅长截然不同的事情：
|
||||
|
||||
| | Python | C++ |
|
||||
|--|--------|-----|
|
||||
| 开发速度 | 快（动态类型、REPL、无需编译） | 慢（静态类型、头文件、编译时间长） |
|
||||
| 执行速度 | 比C慢约100倍（解释型、GIL） | 接近硬件速度（编译型、无开销） |
|
||||
| 内存控制 | 自动（GC），无法控制布局 | 手动，精确控制每一个字节 |
|
||||
| 硬件访问 | 无（无SIMD、无GPU、无自定义内存） | 全面（内联函数、CUDA、内联汇编） |
|
||||
| 生态系统 | ML丰富（笔记本、可视化、数据） | 系统丰富（操作系统、驱动、引擎） |
|
||||
|
||||
- 核心见解：**每种语言发挥其优势**。Python处理人力生产力重要的事务（实验设计、超参数调优、数据探索）。C++处理机器性能重要的事务（矩阵乘法、卷积、注意力核函数）。
|
||||
|
||||
- 一次矩阵乘法 `jnp.matmul(A, B)`，其中 $A$ 为 $4096 \times 4096$，执行约1370亿次浮点运算。在纯Python（嵌套循环）中需要约30分钟。在使用AVX-512 SIMD和多线程优化后的C++中，只需约10毫秒。差距达**180,000倍**。再多的Python技巧也无法弥合这一鸿沟。
|
||||
|
||||
## ML框架的结构
|
||||
|
||||
- 每个主流ML框架都遵循相同的架构：
|
||||
|
||||
```
|
||||
用户代码（Python）
|
||||
↓
|
||||
Python API层（torch.nn、jax.numpy、numpy）
|
||||
↓
|
||||
调度/JIT编译器（torch.compile、XLA、NumPy调度）
|
||||
↓
|
||||
C++核函数库（ATen/PyTorch、XLA、BLAS/LAPACK）
|
||||
↓
|
||||
硬件特定后端（CUDA、cuDNN、MKL、oneDNN、Metal）
|
||||
↓
|
||||
硬件（CPU SIMD单元、GPU核心、TPU MXU）
|
||||
```
|
||||
|
||||
### NumPy
|
||||
|
||||
- NumPy的核心用C编写。当你调用 `np.dot(A, B)` 时，Python调用一个C函数，该函数调用BLAS（基本线性代数子程序），通常是Intel MKL或OpenBLAS。BLAS是手工优化的C和Fortran代码，使用SIMD指令、缓存感知的内存访问模式和多线程。数十年优化致力于让矩阵乘法更快。
|
||||
|
||||
- NumPy仅支持CPU，不使用GPU。但在CPU上，它极其快速，因为它委托给可用的最佳BLAS实现。
|
||||
|
||||
### PyTorch
|
||||
|
||||
- PyTorch的计算引擎是**ATen**（张量库），用C++编写。ATen实现了约2000个张量操作（add、matmul、conv2d、softmax...），每个都有CPU和CUDA后端。
|
||||
|
||||
- 当你调用 `torch.matmul(A, B)` 时：
|
||||
1. Python调度到ATen的C++函数。
|
||||
2. ATen检查设备（CPU或CUDA）和数据类型。
|
||||
3. 在CPU上：调用MKL/OpenBLAS。在GPU上：调用cuBLAS（NVIDIA的GPU优化BLAS）。
|
||||
4. 结果包装在Python张量对象中并返回。
|
||||
|
||||
- **torch.compile**（PyTorch 2.0+）更进一步：它追踪你的Python代码，构建计算图，并使用**Triton**（GPU）或**C++/OpenMP**（CPU）编译。编译后的代码融合操作，消除Python开销，可以比即时模式快2-5倍。
|
||||
|
||||
### JAX
|
||||
|
||||
- JAX将Python函数编译为**XLA**（加速线性代数），Google的ML编译器。当你 `jax.jit` 一个函数时：
|
||||
1. JAX追踪函数，将操作捕获为XLA计算图（HLO——高级操作）。
|
||||
2. XLA优化图：融合操作，消除冗余计算，优化内存布局。
|
||||
3. XLA编译为目标后端：CPU（通过LLVM）、GPU（通过CUDA/PTX）或TPU（通过TPU特定指令）。
|
||||
4. 编译后的代码直接在硬件上运行，零Python参与。
|
||||
|
||||
- 这就是为什么 `jax.jit` 如此重要：没有它，每个操作都是独立的Python→C++往返。有了它，整个函数是一个单一的编译核函数。
|
||||
|
||||
## 面向Python工程师的C++快速入门
|
||||
|
||||
- 你不需要成为C++专家。你需要理解足够的知识来阅读核函数代码、编写简单的扩展以及理解性能讨论。以下是精华内容。
|
||||
|
||||
### 类型和变量
|
||||
|
||||
```cpp
|
||||
// C++需要显式类型（不像Python）
|
||||
int count = 0; // 32位整数
|
||||
float loss = 0.5f; // 32位浮点数
|
||||
double lr = 3e-4; // 64位浮点数
|
||||
bool training = true; // 布尔值
|
||||
|
||||
// 数组（固定大小，栈分配）
|
||||
float weights[1024]; // 1024个浮点数，内存中连续
|
||||
|
||||
// 指针：保存内存地址的变量
|
||||
float* ptr = weights; // ptr指向weights的第一个元素
|
||||
float val = ptr[42]; // 通过指针运算访问元素42
|
||||
// ptr[42] 等价于 *(ptr + 42)
|
||||
```
|
||||
|
||||
- **指针**是与Python最大的概念差异。在Python中，一切都是引用，你从不需要思考内存地址。在C++中，指针让你直接访问内存——强大但危险（悬空指针、缓冲区溢出）。
|
||||
|
||||
### 函数
|
||||
|
||||
```cpp
|
||||
// 函数声明：返回类型 名字(参数类型 参数名)
|
||||
float relu(float x) {
|
||||
return x > 0.0f ? x : 0.0f;
|
||||
}
|
||||
|
||||
// 传引用（避免拷贝大对象）
|
||||
void scale_vector(std::vector<float>& vec, float factor) {
|
||||
for (size_t i = 0; i < vec.size(); i++) {
|
||||
vec[i] *= factor;
|
||||
}
|
||||
}
|
||||
|
||||
// const引用：只读，无拷贝
|
||||
float sum(const std::vector<float>& vec) {
|
||||
float total = 0.0f;
|
||||
for (float x : vec) { // 基于范围的for循环（类似Python的for x in vec）
|
||||
total += x;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
```
|
||||
|
||||
### 内存：栈与堆
|
||||
|
||||
```cpp
|
||||
// 栈分配：快速，自动生命周期（函数返回时释放）
|
||||
float buffer[256]; // 栈上的256个浮点数
|
||||
|
||||
// 堆分配：手动，在函数外仍然存活
|
||||
float* data = new float[n]; // 在堆上分配n个浮点数
|
||||
// ... 使用data ...
|
||||
delete[] data; // 必须手动释放（没有垃圾回收器）
|
||||
|
||||
// 现代C++：智能指针（自动清理，类似Python引用）
|
||||
#include <memory>
|
||||
auto data = std::make_unique<float[]>(n); // 离开作用域时自动释放
|
||||
```
|
||||
|
||||
- **关键规则**：栈快速但有限（通常1-8 MB）。大数组（张量、特征图）必须放在堆上。在Python中，一切都在堆上，GC处理清理。在C++中，你自行管理（或使用智能指针）。
|
||||
|
||||
### 模板（泛型）
|
||||
|
||||
```cpp
|
||||
// 适用于任何数值类型的函数
|
||||
template <typename T>
|
||||
T add(T a, T b) {
|
||||
return a + b;
|
||||
}
|
||||
|
||||
add<float>(1.5f, 2.5f); // 返回 4.0f
|
||||
add<int>(3, 4); // 返回 7
|
||||
```
|
||||
|
||||
- 模板是C++库（如ATen）编写适用于float16、float32、float64等的代码而不重复实现的方式。
|
||||
|
||||
### 标准库精华
|
||||
|
||||
```cpp
|
||||
#include <vector> // 动态数组（类似Python list）
|
||||
#include <string> // 字符串类型
|
||||
#include <unordered_map> // 哈希映射（类似Python dict）
|
||||
#include <algorithm> // sort、find、transform等
|
||||
#include <cmath> // 数学函数
|
||||
|
||||
std::vector<float> vec = {1.0f, 2.0f, 3.0f};
|
||||
vec.push_back(4.0f); // 追加
|
||||
float first = vec[0]; // 索引
|
||||
size_t len = vec.size(); // 长度
|
||||
|
||||
std::unordered_map<std::string, int> counts;
|
||||
counts["hello"] = 5; // 插入
|
||||
if (counts.count("hello")) { } // 检查存在性
|
||||
```
|
||||
|
||||
## 何时编写自定义C++核函数
|
||||
|
||||
- 大多数ML工程师从不需要写C++。框架的内置操作覆盖了99%的用例。仅在以下情况考虑自定义C++：
|
||||
|
||||
1. **框架中不存在你的操作**：新颖的激活函数、自定义注意力模式、无法表示为现有操作组合的特殊损失函数。
|
||||
|
||||
2. **融合操作以提高性能**：你的模型执行 `relu(layernorm(matmul(x, W) + b))`。每个操作启动一个独立的核函数，读写内存，并同步。一个融合核函数在一次遍历中完成所有工作，避免内存往返。这可快2-5倍。
|
||||
|
||||
3. **减少内存使用**：自定义核函数可以在不存储所有中间激活的情况下计算梯度（核函数级别的梯度检查点）。
|
||||
|
||||
4. **针对新型硬件**：新的加速器（如Cerebras、Groq）可能没有框架支持。你需要直接编写核函数。
|
||||
|
||||
- 对于情况1-2，**Triton**（第16章文件05）通常足够且比直接编写CUDA C更简单。只有在Triton无法表达你的需求时才下降到CUDA C。
|
||||
|
||||
## 如何将C++绑定到Python
|
||||
|
||||
- 编写C++只是工作的一半。你还需要从Python调用它。
|
||||
|
||||
### pybind11（通用目的）
|
||||
|
||||
- pybind11用最少的样板代码为C++函数创建Python绑定：
|
||||
|
||||
```cpp
|
||||
// my_ops.cpp
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/numpy.h>
|
||||
namespace py = pybind11;
|
||||
|
||||
// 一个简单的自定义操作
|
||||
py::array_t<float> custom_relu(py::array_t<float> input) {
|
||||
auto buf = input.request();
|
||||
float* ptr = static_cast<float*>(buf.ptr);
|
||||
size_t n = buf.size;
|
||||
|
||||
auto result = py::array_t<float>(n);
|
||||
float* out = static_cast<float*>(result.request().ptr);
|
||||
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
out[i] = ptr[i] > 0 ? ptr[i] : 0;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(my_ops, m) {
|
||||
m.def("custom_relu", &custom_relu, "自定义ReLU操作");
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
# 编译
|
||||
pip install pybind11
|
||||
c++ -O3 -shared -std=c++17 -fPIC $(python3 -m pybind11 --includes) my_ops.cpp -o my_ops$(python3-config --extension-suffix)
|
||||
```
|
||||
|
||||
```python
|
||||
# 从Python使用
|
||||
import my_ops
|
||||
import numpy as np
|
||||
|
||||
x = np.array([-1.0, 2.0, -3.0, 4.0], dtype=np.float32)
|
||||
y = my_ops.custom_relu(x)
|
||||
print(y) # [0. 2. 0. 4.]
|
||||
```
|
||||
|
||||
### PyTorch C++扩展
|
||||
|
||||
- PyTorch提供了一种简化的方式来添加自定义操作：
|
||||
|
||||
```cpp
|
||||
// custom_op.cpp
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor custom_gelu(torch::Tensor x) {
|
||||
return x * 0.5 * (1.0 + torch::erf(x / std::sqrt(2.0)));
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("custom_gelu", &custom_gelu, "自定义GELU激活函数");
|
||||
}
|
||||
```
|
||||
|
||||
```python
|
||||
# 动态加载和编译
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
custom_ops = load(
|
||||
name="custom_ops",
|
||||
sources=["custom_op.cpp"],
|
||||
extra_cflags=["-O3"],
|
||||
)
|
||||
|
||||
x = torch.randn(1000)
|
||||
y = custom_ops.custom_gelu(x)
|
||||
```
|
||||
|
||||
- `torch.utils.cpp_extension.load` 编译C++代码，创建共享库，并将其作为Python模块加载，全在一个调用中完成。这是在PyTorch中实验自定义C++操作的最简单方式。
|
||||
|
||||
### JAX自定义调用
|
||||
|
||||
- JAX使用XLA自定义调用。过程更为复杂（你需要向XLA注册一个C函数），但概念相同：编写C/C++，绑定，从Python调用。
|
||||
|
||||
- 对于大多数JAX用户，**Pallas**（在文件05中介绍）是更好的选择：它让你用类似Python的语法编写GPU核函数，由XLA编译，无需离开JAX生态系统。
|
||||
|
||||
## 大局观
|
||||
|
||||
- 本文解释了Python和硬件之间的层次。本章剩余文件将深入探讨：
|
||||
- **文件01**：硬件本身（CPU架构、GPU架构、内存系统）
|
||||
- **文件02-03**：CPU上的SIMD编程（ARM NEON、x86 AVX）——编写使用CPU向量单元的C++代码
|
||||
- **文件04**：使用CUDA的GPU编程——编写在数千个GPU核心上运行的C++代码
|
||||
- **文件05**：Triton、Pallas和更高级的GPU编程——编写编译为GPU核函数的Python代码
|
||||
|
||||
- 这种递进反映了抽象阶梯：C++内联函数（最低层、最多控制）→ CUDA（GPU专用）→ Triton/Pallas（Python风格、编译型）→ JAX/PyTorch（最高层、自动）。每一层以控制权换取便利性。理解较低层使你成为较高层的更好使用者。
|
||||
|
||||
## 编程任务（用g++或clang++编译）
|
||||
|
||||
1. 编写你的第一个C++程序。分配一个数组，填充数据，计算总和，并测量时间。这介绍了编译、数组、指针和计时。
|
||||
```cpp
|
||||
// task1_basics.cpp
|
||||
// 编译：g++ -O3 -o task1 task1_basics.cpp
|
||||
// 运行：./task1
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
|
||||
int main() {
|
||||
const int N = 10'000'000; // C++允许'作为数字分隔符
|
||||
std::vector<float> data(N);
|
||||
|
||||
// 填充数组
|
||||
for (int i = 0; i < N; i++) {
|
||||
data[i] = static_cast<float>(i) * 0.001f;
|
||||
}
|
||||
|
||||
// 计算总和
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
float sum = 0.0f;
|
||||
for (int i = 0; i < N; i++) {
|
||||
sum += data[i];
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double elapsed = std::chrono::duration<double, std::milli>(end - start).count();
|
||||
|
||||
std::cout << "总和: " << sum << std::endl;
|
||||
std::cout << "时间: " << elapsed << " ms" << std::endl;
|
||||
std::cout << "元素数: " << N << std::endl;
|
||||
std::cout << "吞吐量: " << (N * sizeof(float)) / elapsed / 1e6 << " GB/s" << std::endl;
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
2. 编写一个C++函数在数组上计算ReLU，然后使用pybind11构建Python绑定。从Python调用它并与NumPy比较速度。
|
||||
```cpp
|
||||
// task2_relu.cpp
|
||||
// 编译：c++ -O3 -shared -std=c++17 -fPIC $(python3 -m pybind11 --includes) \
|
||||
// task2_relu.cpp -o my_relu$(python3-config --extension-suffix)
|
||||
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/numpy.h>
|
||||
namespace py = pybind11;
|
||||
|
||||
py::array_t<float> cpp_relu(py::array_t<float> input) {
|
||||
auto buf = input.request();
|
||||
float* ptr = static_cast<float*>(buf.ptr);
|
||||
int n = buf.size;
|
||||
|
||||
auto result = py::array_t<float>(n);
|
||||
float* out = static_cast<float*>(result.request().ptr);
|
||||
|
||||
for (int i = 0; i < n; i++) {
|
||||
out[i] = ptr[i] > 0.0f ? ptr[i] : 0.0f;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(my_relu, m) {
|
||||
m.def("relu", &cpp_relu, "C++ ReLU");
|
||||
}
|
||||
```
|
||||
```python
|
||||
# test_relu.py — 在编译上述C++模块后运行
|
||||
import numpy as np
|
||||
import time
|
||||
import my_relu # 编译后的C++模块
|
||||
|
||||
x = np.random.randn(10_000_000).astype(np.float32)
|
||||
|
||||
# C++ ReLU
|
||||
start = time.time()
|
||||
for _ in range(100):
|
||||
y_cpp = my_relu.relu(x)
|
||||
cpp_time = (time.time() - start) / 100
|
||||
|
||||
# NumPy ReLU
|
||||
start = time.time()
|
||||
for _ in range(100):
|
||||
y_np = np.maximum(x, 0)
|
||||
np_time = (time.time() - start) / 100
|
||||
|
||||
print(f"C++ ReLU: {cpp_time*1000:.2f} ms")
|
||||
print(f"NumPy ReLU: {np_time*1000:.2f} ms")
|
||||
print(f"匹配: {np.allclose(y_cpp, y_np)}")
|
||||
```
|
||||
|
||||
3. 编写一个C++程序，演示为何内存布局很重要。比较行优先与列优先访问模式并测量性能差异。
|
||||
```cpp
|
||||
// task3_layout.cpp
|
||||
// 编译：g++ -O3 -o task3 task3_layout.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
|
||||
int main() {
|
||||
const int N = 4096;
|
||||
std::vector<float> matrix(N * N, 1.0f);
|
||||
|
||||
// 行优先访问：连续内存地址（缓存友好）
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
float sum_row = 0.0f;
|
||||
for (int i = 0; i < N; i++) {
|
||||
for (int j = 0; j < N; j++) {
|
||||
sum_row += matrix[i * N + j]; // 步长1访问
|
||||
}
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double row_ms = std::chrono::duration<double, std::milli>(end - start).count();
|
||||
|
||||
// 列优先访问：步长N访问（缓存不友好）
|
||||
start = std::chrono::high_resolution_clock::now();
|
||||
float sum_col = 0.0f;
|
||||
for (int j = 0; j < N; j++) {
|
||||
for (int i = 0; i < N; i++) {
|
||||
sum_col += matrix[i * N + j]; // 步长N访问（缓存缺失！）
|
||||
}
|
||||
}
|
||||
end = std::chrono::high_resolution_clock::now();
|
||||
double col_ms = std::chrono::duration<double, std::milli>(end - start).count();
|
||||
|
||||
std::cout << "行优先（缓存友好）: " << row_ms << " ms" << std::endl;
|
||||
std::cout << "列优先（缓存不友好）: " << col_ms << " ms" << std::endl;
|
||||
std::cout << "减速比: " << col_ms / row_ms << "x" << std::endl;
|
||||
std::cout << "（两个和: " << sum_row << ", " << sum_col << "）" << std::endl;
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,282 @@
|
||||
# 硬件基础
|
||||
|
||||
*在编写SIMD或GPU代码之前，你需要了解所编程的硬件。本文涵盖为什么并行性取代了时钟速度、现代CPU如何执行指令、什么是SIMD、用于推理性能的屋顶线模型，以及芯片架构的全景*
|
||||
|
||||
- 几十年来，软件免费变快：购买一个时钟频率更高的新CPU，你的程序无需修改一行代码就能运行得更快。这个时代大约在2005年结束。理解它为何结束以及什么替代了它，对任何想编写快速代码的人都至关重要。
|
||||
|
||||
## 免费性能的终结
|
||||
|
||||
- **摩尔定律**（1965年）观察到芯片上的晶体管数量大约每两年翻一番。这一规律维持了60年。更多晶体管意味着更小的晶体管，进而意味着更高的时钟频率，从而意味着更快的程序。
|
||||
|
||||
- 但在2005年左右，时钟频率在大约4 GHz处撞上了墙壁。问题是**功耗**。芯片消耗的功率大约为：
|
||||
|
||||
$$P \propto C \cdot V^2 \cdot f$$
|
||||
|
||||
- 其中 $C$ 是电容（与晶体管数量成正比），$V$ 是电压，$f$ 是时钟频率。要提高频率，必须提高电压（以使晶体管更快地切换）。但功耗与 $V^2 \cdot f$ 成比例，所以频率的小幅增加会导致功耗（和热量）的大幅增加。在4 GHz时，芯片已经达到100+瓦。达到8 GHz需要不切实际的冷却方案。
|
||||
|
||||
- 解决方案：不让单个核心更快，而是在同一芯片上放置**多个核心**。一个4核芯片在3 GHz下使用与单个核心在4.5 GHz下相似的功耗，但可以做4倍的并行工作。这就是为什么每个现代CPU都有多个核心，以及为什么并行性（SIMD、多线程、GPU计算）是获得更高性能的唯一途径。
|
||||
|
||||
- **对ML的影响**：一个在单核上需要10分钟的训练步骤，无法通过购买更快的CPU来加速。只能通过使用更多核心（数据并行性，第6章）、更宽的SIMD单元（本章）或GPU（数千个核心）来加速。
|
||||
|
||||
## 现代CPU如何执行指令
|
||||
|
||||
- 现代CPU核心远比第13章中简单的取指-译码-执行模型复杂。它使用几种技巧来每周期执行更多指令：
|
||||
|
||||
- **超标量执行**：CPU有多个执行单元（ALU、FPU、加载/存储单元），可以同时执行多个独立的指令。如果指令不相互依赖，现代核心每周期可能执行4-6条指令。
|
||||
|
||||
- **乱序执行（OoO）**：CPU不按程序顺序执行指令。它向前看指令流，找到输入已准备好的指令，并立即执行，不论其位置。这隐藏了延迟：当一条指令等待来自内存的数据时（100+周期），CPU执行其他已准备好的指令。
|
||||
|
||||
- **分支预测**：条件分支（`if`语句、循环条件）造成不确定性：CPU在条件被评估之前不知道走哪条路径。为了避免停顿，CPU**预测**结果并沿预测路径投机执行。如果预测正确（使用现代预测器超过95%），则没有时间损失。如果错误，投机工作被丢弃，执行正确路径（约15周期惩罚）。
|
||||
|
||||
- **投机执行**：分支预测的延伸。CPU执行*可能*不需要的指令，赌它们会被需要。这填充了流水线并保持执行单元忙碌。
|
||||
|
||||
- 所有这些都是自动的——CPU无需任何程序员干预即可完成。但它们只帮助**指令级并行性（ILP）**：单条流内相互独立的指令。对于**数据级并行性**（对许多数据元素执行相同操作），我们需要SIMD。
|
||||
|
||||
## SIMD：单指令多数据
|
||||
|
||||
- **SIMD**是将一条指令同时应用于多个数据元素的思想。不是将两个数相加，而是在一条指令中将两个4（或8、或16）元素向量相加。
|
||||
|
||||
- 无SIMD（标量）：
|
||||
|
||||
```cpp
|
||||
// 逐元素相加两数组：4条加法指令
|
||||
for (int i = 0; i < 4; i++) {
|
||||
c[i] = a[i] + b[i]; // 每次迭代一次加法
|
||||
}
|
||||
```
|
||||
|
||||
- 有SIMD（向量化）：
|
||||
|
||||
```cpp
|
||||
// 两数组相加：1条SIMD指令完成所有4次加法
|
||||
#include <immintrin.h> // x86 SIMD内联函数
|
||||
|
||||
__m128 va = _mm_load_ps(a); // 加载4个浮点数到128位寄存器
|
||||
__m128 vb = _mm_load_ps(b); // 加载4个浮点数到另一个寄存器
|
||||
__m128 vc = _mm_add_ps(va, vb); // 同时相加所有4对
|
||||
_mm_store_ps(c, vc); // 存储4个结果
|
||||
```
|
||||
|
||||
- SIMD版本用1/4的指令完成相同工作。这是理论上的4倍加速，通过每条指令处理4个浮点数而非1个实现。
|
||||
|
||||
### 向量寄存器
|
||||
|
||||
- SIMD指令操作**向量寄存器**：保存多个数据元素的宽寄存器。
|
||||
|
||||
| 寄存器宽度 | 浮点数（32位） | 双精度浮点数（64位） | 名称 |
|
||||
|---------------|-----------------|-------------------|------|
|
||||
| 128位 | 4 | 2 | SSE（x86）、NEON（ARM） |
|
||||
| 256位 | 8 | 4 | AVX/AVX2（x86） |
|
||||
| 512位 | 16 | 8 | AVX-512（x86） |
|
||||
| 可变（128-2048） | 可变 | 可变 | SVE/SVE2（ARM） |
|
||||
|
||||
- 更宽的寄存器 = 更多并行性。一条512位AVX-512指令一次处理16个浮点数，是标量代码理论上的16倍加速。实际上，由于内存带宽限制（计算速度可能超过向CPU输送数据的速度），加速比更低。
|
||||
|
||||
- 对于ML：float32值的矩阵乘法从SIMD中获益巨大。内循环（两个向量的点积）直接映射到SIMD乘加指令。这就是为什么BLAS库（NumPy和PyTorch调用的）用SIMD进行了如此深度优化。
|
||||
|
||||
## 屋顶线模型
|
||||
|
||||
- 你如何知道你的代码是否快速？**屋顶线模型**提供了一个框架，根据两个硬件限制来描述性能：
|
||||
|
||||
1. **峰值计算能力**（FLOPS）：每秒最大浮点运算次数。对于一个4 GHz CPU，配备256位AVX（每条指令8个浮点数）和2个FMA单元：$4 \times 10^9 \times 8 \times 2 = 64$ GFLOPS。
|
||||
|
||||
2. **峰值内存带宽**（字节/秒）：数据从内存到CPU的最大传输速度。现代CPU可能有50 GB/s的内存带宽。
|
||||
|
||||
- 代码的**算术强度**是计算与内存访问的比率：
|
||||
|
||||
$$\text{算术强度} = \frac{\text{FLOPS}}{\text{传输的字节数}}$$
|
||||
|
||||
- 如果算术强度低（每加载字节的操作数少），你的代码是**内存受限的**：大部分时间花在等待数据上。让计算更快（更宽的SIMD、更高的时钟）不会有帮助。
|
||||
|
||||
- 如果算术强度高（每字节多次操作），你的代码是**计算受限的**：大部分时间花在计算上。更快的内存不会有帮助。
|
||||
|
||||
- 屋顶线：
|
||||
|
||||
$$\text{可达FLOPS} = \min\left(\text{峰值FLOPS}, \; \text{带宽} \times \text{算术强度}\right)$$
|
||||
|
||||
- **矩阵乘法**具有高算术强度：$O(n^3)$ 次操作作用于 $O(n^2)$ 数据，因此强度 $\approx O(n)$。对于大矩阵，它是计算受限的。这就是为什么GPU（高计算能力）主导矩阵密集型的ML工作负载。
|
||||
|
||||
- **逐元素操作**（ReLU、加法、乘法）具有低算术强度：每加载一个元素1次操作。这些是内存受限的。让GPU更快没有帮助；你需要更快的内存（或者将这些操作与计算密集型操作融合，以避免独立的内存往返）。
|
||||
|
||||
- 屋顶线模型解释了为什么**核函数融合**如此重要：将matmul与偏置加法和ReLU组合成一个核函数，避免了将中间结果写入内存并重新读取，将三个内存受限操作转化为一个计算受限操作。
|
||||
|
||||
## 延迟与吞吐量
|
||||
|
||||
- **延迟**是完成一个操作所需的时间。**吞吐量**是单位时间内完成的操作数量。
|
||||
|
||||
- 打个比方：公交车延迟高（每站都停），但吞吐量高（一次搭载50人）。出租车延迟低（直达你的目的地），但吞吐量低（搭载1-4人）。
|
||||
|
||||
- GPU是公交车：每次操作延迟高（每条指令需要许多周期完成），但吞吐量巨大（数千个核心同时处理）。CPU是出租车：延迟低（乱序执行、分支预测、深层缓存最小化延迟），但吞吐量有限（4-64个核心）。
|
||||
|
||||
- 这就是为什么GPU更适合ML训练（吞吐量重要：处理数百万个样本）而CPU更适合操作系统任务（延迟重要：立即响应按键）。
|
||||
|
||||
- **流水线**将延迟转化为吞吐量。如果一条指令需要5个周期，但流水线每周期开始一条新指令，则吞吐量是每条指令1周期（即使每条指令需要5个周期完成）。这和第13章的CPU流水线是同一原理，但它适用于每个层面：SIMD单元、内存控制器和GPU核心都是流水线化的。
|
||||
|
||||
## 芯片架构全景
|
||||
|
||||
- 你编写代码的硬件决定了哪些SIMD指令可用：
|
||||
|
||||
### x86（Intel, AMD）
|
||||
|
||||
- 主导台式机、笔记本电脑和数据中心CPU。SIMD：SSE（128位）、AVX/AVX2（256位）、AVX-512（512位）。Intel AMX提供专用的矩阵乘法单元用于AI工作负载。
|
||||
|
||||
- **优势**：最高单核性能、最宽SIMD、成熟的软件生态系统（MKL、oneDNN）。
|
||||
- **弱点**：高功耗、复杂指令集、昂贵。
|
||||
|
||||
### ARM
|
||||
|
||||
- 主导移动设备（每一部智能手机），在服务器（AWS Graviton、Ampere Altra）和笔记本电脑（Apple M系列）中增长。SIMD：NEON（128位）、SVE/SVE2（可伸缩，128-2048位）。
|
||||
|
||||
- **优势**：出色的功耗效率（每瓦性能）、自定义核心（Apple M4在单核性能上媲美Intel，功耗仅为其一小部分）。
|
||||
- **弱点**：较窄的SIMD（NEON仅为128位，虽SVE可更宽）、用于HPC的软件生态系统较小。
|
||||
|
||||
### Apple Silicon（M1/M2/M3/M4）
|
||||
|
||||
- 基于ARM并带有自定义扩展。包含**AMX**（Apple矩阵扩展）——未公开的矩阵乘法单元，Accelerate框架将其用于BLAS操作。统一内存架构：CPU和GPU共享同一物理内存，消除了CPU↔GPU拷贝的瓶颈。
|
||||
|
||||
- **对于ML**：Apple的神经网络引擎（16核，专用ML加速器）和统一内存使M系列芯片在本地ML推理和小规模训练方面出奇地强大。不过没有CUDA：你必须使用Metal（Apple的GPU API）或MLX（Apple的ML框架）。
|
||||
|
||||
### RISC-V
|
||||
|
||||
- 开源ISA。无许可费用（不像ARM）。在嵌入式系统、物联网和研究领域增长。SIMD："V"（向量）扩展提供类似于ARM SVE的可伸缩向量处理。
|
||||
|
||||
- **对于ML**：在ML工作负载上尚不能与x86/ARM竞争，但值得关注。几个AI加速器初创公司使用RISC-V核心。
|
||||
|
||||
### GPU（NVIDIA、AMD、Intel）
|
||||
|
||||
- 在文件04-05中深入介绍。数千个为吞吐量优化的简单核心。NVIDIA以CUDA主导ML；AMD以ROCm竞争；Intel以Arc GPU和Gaudi加速器进入市场。
|
||||
|
||||
### TPU（Google）
|
||||
|
||||
- 专门为ML设计的自定义ASIC。为矩阵乘法优化的脉动阵列。在文件05中介绍。
|
||||
|
||||
## 热与功耗约束
|
||||
|
||||
- 性能最终受限于功耗和散热：
|
||||
|
||||
- **TDP**（热设计功耗）：芯片可以持续消耗的最大功率。笔记本电脑CPU可能有15W TDP；服务器CPU 250W；数据中心GPU 700W（NVIDIA B200）。
|
||||
|
||||
- **暗硅**：在任何给定时刻，为了保持在热预算内，必须关闭芯片的相当大一部分晶体管。理论上芯片可以同时使用所有晶体管，但会熔化。
|
||||
|
||||
- **功耗效率**（FLOPS/瓦）日益成为重要指标，而非原始FLOPS。这就是为什么：
|
||||
- ARM正在接管数据中心（相比x86更好的FLOPS/瓦）。
|
||||
- TPU尽管峰值FLOPS较低，但仍与GPU竞争（对于ML工作负载，FLOPS/瓦好得多）。
|
||||
- 量化（INT8、FP8）不仅关乎内存：它也降低了每次操作的功耗。
|
||||
|
||||
- 对于大规模ML：训练前沿LLM数月消耗兆瓦级功率。电费可能超过硬件成本。功耗效率直接影响AI研究的经济性。
|
||||
|
||||
## 实践：在C++中测量性能
|
||||
|
||||
- 要推理性能，你需要测量它。以下是一个最小的C++基准测试设置：
|
||||
|
||||
```cpp
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
|
||||
// 标量加法
|
||||
void add_scalar(const float* a, const float* b, float* c, int n) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 1 << 24; // 约1600万个元素
|
||||
std::vector<float> a(N, 1.0f), b(N, 2.0f), c(N);
|
||||
|
||||
// 预热（填充缓存，触发频率缩放）
|
||||
add_scalar(a.data(), b.data(), c.data(), N);
|
||||
|
||||
// 基准测试
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
|
||||
for (int trial = 0; trial < 100; trial++) {
|
||||
add_scalar(a.data(), b.data(), c.data(), N);
|
||||
}
|
||||
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double elapsed = std::chrono::duration<double>(end - start).count();
|
||||
|
||||
double total_bytes = 3.0 * N * sizeof(float) * 100; // 读a、读b、写c
|
||||
double bandwidth = total_bytes / elapsed / 1e9; // GB/s
|
||||
|
||||
std::cout << "时间: " << elapsed << " s\n";
|
||||
std::cout << "带宽: " << bandwidth << " GB/s\n";
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
# 启用优化编译
|
||||
g++ -O3 -march=native -o bench bench.cpp
|
||||
./bench
|
||||
```
|
||||
|
||||
- **这段代码中的关键C++概念**：
|
||||
- `#include <vector>`：动态数组（`std::vector<float>`）——类似Python的`list`但带类型且在内存中连续。
|
||||
- `a.data()`：返回底层数组的原始指针（`float*`）——SIMD内联函数需要。
|
||||
- `std::chrono`：用于基准测试的高分辨率计时器。
|
||||
- `-O3`：最高编译器优化级别。编译器可能自动向量化你的循环（自动使用SIMD）。`-march=native` 启用你的CPU支持的所有SIMD指令。
|
||||
|
||||
- **为什么需要预热**：首次运行填充缓存并可能触发CPU频率缩放（睿频加速）。后续运行更具代表性。
|
||||
|
||||
- **为什么测量带宽**：对于内存受限的操作（如逐元素加法），有意义的度量是带宽（GB/s），而不是FLOPS。如果你的测量带宽接近硬件极限（DDR5约50 GB/s），你是内存受限的，SIMD不会有多大帮助（瓶颈是内存，而非计算）。
|
||||
|
||||
## 编程任务（使用CoLab或笔记本）
|
||||
|
||||
1. 计算常见ML操作的算术强度，并将它们分类为内存受限或计算受限。
|
||||
```python
|
||||
import jax.numpy as jnp
|
||||
|
||||
def arithmetic_intensity(flops, bytes_transferred):
|
||||
return flops / bytes_transferred
|
||||
|
||||
# 逐元素ReLU：每元素1次比较，读取+写入
|
||||
n = 1024
|
||||
relu_flops = n # 每元素1次操作
|
||||
relu_bytes = 2 * n * 4 # 读取输入+写入输出（float32）
|
||||
print(f"ReLU: {arithmetic_intensity(relu_flops, relu_bytes):.2f} FLOPS/byte → 内存受限")
|
||||
|
||||
# 矩阵乘法：2*n^3次操作，读取2*n^2 + 写入n^2个浮点数
|
||||
matmul_flops = 2 * n**3
|
||||
matmul_bytes = 3 * n**2 * 4 # 读取A + 读取B + 写入C
|
||||
print(f"矩阵乘法 ({n}×{n}): {arithmetic_intensity(matmul_flops, matmul_bytes):.0f} FLOPS/byte → 计算受限")
|
||||
|
||||
# 层归一化：约5n次操作（均值、方差、归一化），读取+写入
|
||||
ln_flops = 5 * n
|
||||
ln_bytes = 2 * n * 4
|
||||
print(f"LayerNorm: {arithmetic_intensity(ln_flops, ln_bytes):.2f} FLOPS/byte → 内存受限")
|
||||
|
||||
# 3x3卷积：2*9*C_in*C_out*H*W，读取卷积核+特征图+写入输出
|
||||
C_in, C_out, H, W = 64, 128, 32, 32
|
||||
conv_flops = 2 * 9 * C_in * C_out * H * W
|
||||
conv_bytes = (9 * C_in * C_out + C_in * H * W + C_out * H * W) * 4
|
||||
print(f"Conv3x3: {arithmetic_intensity(conv_flops, conv_bytes):.0f} FLOPS/byte → 计算受限")
|
||||
```
|
||||
|
||||
2. 演示为什么并行性重要。比较顺序执行与并行（NumPy）执行随数据规模增长的表现。
|
||||
```python
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
for n in [1000, 10000, 100000, 1000000, 10000000]:
|
||||
a = np.random.randn(n).astype(np.float32)
|
||||
b = np.random.randn(n).astype(np.float32)
|
||||
|
||||
# "顺序执行"（Python循环）
|
||||
start = time.time()
|
||||
c = [a[i] * b[i] for i in range(min(n, 100000))] # 上限10万以确保合理
|
||||
seq_time = time.time() - start
|
||||
if n > 100000:
|
||||
seq_time *= n / 100000 # 外推
|
||||
|
||||
# "并行"（NumPy，内部使用SIMD+多线程）
|
||||
start = time.time()
|
||||
c = a * b
|
||||
par_time = time.time() - start
|
||||
|
||||
print(f"n={n:>10,} 顺序={seq_time:.4f}s 并行={par_time:.6f}s "
|
||||
f"加速比={seq_time/par_time:.0f}x")
|
||||
```
|
||||
@@ -0,0 +1,484 @@
|
||||
# ARM与NEON
|
||||
|
||||
*ARM处理器驱动着每一部智能手机、大多数平板电脑、Apple的笔记本电脑以及日益增长的数据中心服务器份额。本文涵盖ARM架构、使用C++内联函数的NEON SIMD编程、用于可伸缩向量处理的SVE/SVE2、Apple Silicon特性以及实际向量化核函数示例*
|
||||
|
||||
- 如果你拥有iPhone、MacBook或使用AWS Graviton实例，你正在运行ARM。ARM的功耗效率使其在移动和嵌入式领域占据主导地位，并在服务器和ML推理方面日益具有竞争力。理解ARM SIMD让你能够编写在大多数人实际使用的硬件上快速运行的代码。
|
||||
|
||||
- 有关生产中ARM SIMD核函数的实际例子，请参见**Cactus**——面向移动设备和可穿戴设备的低延迟AI引擎：[github.com/cactus-compute/cactus](https://github.com/cactus-compute/cactus)。Cactus实现了自定义ARM NEON和NPU加速的注意机制、KV缓存量化和分块预填充核函数，在ARM CPU上实现了最快的推理，且RAM比其它引擎低10倍。其三层架构（引擎→图→核函数）是本文中SIMD概念如何用于构建生产级ML基础设施的具体实例。
|
||||
|
||||
## ARM架构基础
|
||||
|
||||
- ARM是一种**RISC**（精简指令集计算机）架构（第13章）。关键特征：
|
||||
|
||||
- **加载-存储架构**：算术指令只操作寄存器，从不直接操作内存。要对内存中的两个数相加，你必须：(1) 将它们加载到寄存器，(2) 将寄存器相加，(3) 将结果存回内存。这比x86更简单（x86可以在一条指令中加一个寄存器和一个内存位置），但使得流水线更清晰。
|
||||
|
||||
- **定长指令**：每个ARMv8（AArch64）指令恰好32位。这使得解码快速且可预测（不像x86的可变长指令，长度可以是1-15字节）。
|
||||
|
||||
- **32个通用寄存器**（x0-x30，每个64位）加上栈指针（sp）和零寄存器（xzr）。相比之下x86有16个通用寄存器。更多寄存器 = 更少内存访问 = 更快代码。
|
||||
|
||||
- **32个SIMD/浮点寄存器**（v0-v31，每个128位）用于NEON和浮点操作。
|
||||
|
||||
```cpp
|
||||
// ARM汇编（仅感受风格——你将使用内联函数，而非汇编）
|
||||
// 两寄存器相加
|
||||
add x0, x1, x2 // x0 = x1 + x2
|
||||
|
||||
// 从内存加载
|
||||
ldr x0, [x1] // x0 = *x1（从x1中的地址加载64位）
|
||||
|
||||
// NEON：加四个浮点数
|
||||
fadd v0.4s, v1.4s, v2.4s // v0 = v1 + v2（四个32位浮点数）
|
||||
```
|
||||
|
||||
- 你不会写汇编。你将使用**内联函数**：与特定指令一对一映射的C/C++函数。编译器处理寄存器分配、调度和其他底层细节。
|
||||
|
||||
## NEON：128位SIMD
|
||||
|
||||
- **NEON**是ARM的SIMD扩展。每个NEON寄存器宽128位，可容纳：
|
||||
|
||||
| 数据类型 | 每寄存器元素数 | 表示法 |
|
||||
|-----------|---------------|----------|
|
||||
| float32 | 4 | `float32x4_t` |
|
||||
| float16 | 8 | `float16x8_t` |
|
||||
| int32 | 4 | `int32x4_t` |
|
||||
| int16 | 8 | `int16x8_t` |
|
||||
| int8 | 16 | `int8x16_t` |
|
||||
|
||||
- 128位比x86的AVX（256位）或AVX-512（512位）窄。但ARM以出色的功耗效率和广泛的可用性弥补了这一点。
|
||||
|
||||
### NEON内联函数：基础
|
||||
|
||||
- NEON内联函数遵循命名约定：`v[操作][限定符]_[类型]`
|
||||
|
||||
```cpp
|
||||
#include <arm_neon.h>
|
||||
|
||||
// 从内存加载4个浮点数到NEON寄存器
|
||||
float32x4_t a = vld1q_f32(ptr); // vld1q = vector load 1, q = 128位（四字）
|
||||
|
||||
// 从NEON寄存器存储4个浮点数到内存
|
||||
vst1q_f32(out_ptr, a); // vst1q = vector store 1, q = 128位
|
||||
|
||||
// 算术运算
|
||||
float32x4_t c = vaddq_f32(a, b); // c = a + b（4个浮点数）
|
||||
float32x4_t d = vmulq_f32(a, b); // d = a * b（4个浮点数）
|
||||
float32x4_t e = vfmaq_f32(c, a, b); // e = c + a * b（融合乘加，4个浮点数）
|
||||
|
||||
// 比较（返回掩码：若真则全1，若假则全0）
|
||||
uint32x4_t mask = vcgtq_f32(a, b); // mask[i] = (a[i] > b[i]) ? 0xFFFFFFFF : 0
|
||||
|
||||
// 基于掩码选择元素（类似numpy.where）
|
||||
float32x4_t result = vbslq_f32(mask, a, b); // result[i] = mask[i] ? a[i] : b[i]
|
||||
|
||||
// 归约：将所有4个元素求和为标量
|
||||
float total = vaddvq_f32(a); // total = a[0] + a[1] + a[2] + a[3]
|
||||
```
|
||||
|
||||
- **`vfmaq_f32`**（融合乘加）是ML最重要的SIMD指令。它用一次舍入步骤计算 $c = c + a \times b$（比分开乘然后加更精确）。点积、矩阵乘法和卷积都由FMA构建。
|
||||
|
||||
### 实践示例：向量化点积
|
||||
|
||||
- 点积是矩阵乘法的内循环。让我们先用标量C++编写，然后用NEON向量化。
|
||||
|
||||
```cpp
|
||||
#include <arm_neon.h>
|
||||
|
||||
// 标量点积
|
||||
float dot_scalar(const float* a, const float* b, int n) {
|
||||
float sum = 0.0f;
|
||||
for (int i = 0; i < n; i++) {
|
||||
sum += a[i] * b[i];
|
||||
}
|
||||
return sum;
|
||||
}
|
||||
|
||||
// NEON向量化点积
|
||||
float dot_neon(const float* a, const float* b, int n) {
|
||||
float32x4_t sum_vec = vdupq_n_f32(0.0f); // 初始化4个累加器为0
|
||||
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
float32x4_t va = vld1q_f32(a + i); // 从a加载4个元素
|
||||
float32x4_t vb = vld1q_f32(b + i); // 从b加载4个元素
|
||||
sum_vec = vfmaq_f32(sum_vec, va, vb); // sum_vec += va * vb
|
||||
}
|
||||
|
||||
// 将4个累加器归约为单一标量
|
||||
float sum = vaddvq_f32(sum_vec);
|
||||
|
||||
// 处理剩余元素（如果n不是4的倍数）
|
||||
for (; i < n; i++) {
|
||||
sum += a[i] * b[i];
|
||||
}
|
||||
|
||||
return sum;
|
||||
}
|
||||
```
|
||||
|
||||
- **关键C++概念**：
|
||||
- `const float*`：指向只读浮点数据的指针。`const` 承诺我们不会通过此指针修改数据。
|
||||
- `a + i`：指针运算。`a + i` 指向数组的第 $i$ 个元素（等同于 `&a[i]`）。
|
||||
- 末尾的"清理循环"处理 $n$ 不是4的倍数的情况。这是SIMD代码中的通用模式：用向量化块处理主体部分，然后用标量代码处理余数。
|
||||
|
||||
- **为什么 `sum_vec` 中使用4个累加器**：我们使用4个独立的累加器（每个SIMD通道一个），而不是单个标量累加器。这避免了数据依赖：每次迭代的FMA依赖于 `sum_vec`，但有了4个独立通道，CPU可以对FMAs进行流水线处理。最后，我们将4个部分和归约为一个。
|
||||
|
||||
### 实践示例：向量化ReLU
|
||||
|
||||
```cpp
|
||||
#include <arm_neon.h>
|
||||
|
||||
void relu_neon(const float* input, float* output, int n) {
|
||||
float32x4_t zero = vdupq_n_f32(0.0f);
|
||||
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
float32x4_t x = vld1q_f32(input + i);
|
||||
float32x4_t result = vmaxq_f32(x, zero); // max(x, 0) = ReLU
|
||||
vst1q_f32(output + i, result);
|
||||
}
|
||||
|
||||
// 标量清理
|
||||
for (; i < n; i++) {
|
||||
output[i] = input[i] > 0 ? input[i] : 0;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `vmaxq_f32` 计算两个向量的逐元素最大值。由于一个向量全为零，这恰好就是ReLU。无需分支，无需比较——仅一条指令。
|
||||
|
||||
## I8MM：整数矩阵乘法
|
||||
|
||||
- **I8MM**（Int8矩阵乘法）是ARMv8.6扩展，增加了用于INT8矩阵乘法（INT32累加）的专用指令——这正是量化ML推理所需要的。
|
||||
|
||||
- 关键指令是 **`SMMLA`**（有符号矩阵乘加）：它接受两个8×2块的INT8值，并将结果累加到2×2块的INT32中：
|
||||
|
||||
```cpp
|
||||
#include <arm_neon.h>
|
||||
|
||||
// I8MM：将两个8元素INT8向量相乘，累加到4个INT32结果中
|
||||
// 这从2x8 × 8x2输入块计算输出矩阵的一个2x2瓦片
|
||||
void matmul_i8mm_tile(const int8_t* A, const int8_t* B, int32_t* C) {
|
||||
// 从A加载8字节（2行各4元素，打包）
|
||||
int8x16_t va = vld1q_s8(A); // 16字节 = 2行 × 8元素
|
||||
int8x16_t vb = vld1q_s8(B); // 16字节 = 2行 × 8元素
|
||||
|
||||
// 加载现有累加器（2x2 = 4个int32值）
|
||||
int32x4_t acc = vld1q_s32(C);
|
||||
|
||||
// I8MM指令：acc += A_tile × B_tile^T
|
||||
// 从2×8 × 8×2输入计算2×2输出
|
||||
acc = vmmlaq_s32(acc, va, vb); // I8MM指令
|
||||
|
||||
vst1q_s32(C, acc);
|
||||
}
|
||||
```
|
||||
|
||||
- **为什么I8MM重要**：没有I8MM时，NEON上的INT8矩阵乘法需要加宽乘法（`vmull`）后跟成对加法——每个输出元素需要多条指令。有了I8MM，硬件在一条指令中完成8元素点积（2×8 × 8×2 = 2×2）。对于INT8推理工作负载，这比纯NEON快4-8倍。
|
||||
|
||||
- **可用性**：Apple M1+（所有Apple Silicon）、ARM Cortex-A510/A710/X2+（ARMv9）、AWS Graviton3+。用 `#ifdef __ARM_FEATURE_MATMUL_INT8` 检查。
|
||||
|
||||
- 对于ML推理：在ARM服务器（Graviton）或Apple Silicon上运行的INT8量化模型（第18章）从I8MM中获益巨大。ONNX Runtime和llama.cpp等框架在运行时检测I8MM并自动使用优化核函数。
|
||||
|
||||
## SME和SME2：可伸缩矩阵扩展
|
||||
|
||||
- **SME**（可伸缩矩阵扩展）是ARM对Intel AMX和NVIDIA张量核心的回应：用于矩阵操作的专用硬件。SME2（ARMv9.2）进一步扩展了它。
|
||||
|
||||
- SME引入了**ZA瓦片寄存器**：存储在硬件中的2D矩阵，最大可达SVL×SVL字节（其中SVL是流向量长度，通常每维128-512位）。与NEON（1D向量）甚至SVE（1D可伸缩向量）不同，SME原生操作**2D瓦片**。
|
||||
|
||||
- 编程模型有两种模式：
|
||||
- **普通模式**：标准ARM执行（NEON、SVE正常工作）。
|
||||
- **流SVE模式**：通过 `smstart` 进入，启用SME指令。SVE指令在此模式下也可工作，但可能使用不同的寄存器宽度。
|
||||
|
||||
```cpp
|
||||
#include <arm_sme.h>
|
||||
|
||||
// SME2：矩阵乘法的外积累加
|
||||
// 将A_col × B_row 累加到ZA瓦片寄存器中
|
||||
void sme2_matmul_outer(const float* A_col, const float* B_row, int K) {
|
||||
// 进入流模式
|
||||
// smstart; // （通过编译器内联或内联汇编完成）
|
||||
|
||||
// 清零ZA瓦片累加器
|
||||
svzero_za();
|
||||
|
||||
for (int k = 0; k < K; k++) {
|
||||
// 将A的一列和B的一行加载到SVE寄存器中
|
||||
svfloat32_t a = svld1_f32(svptrue_b32(), &A_col[k * SVL]);
|
||||
svfloat32_t b = svld1_f32(svptrue_b32(), &B_row[k * SVL]);
|
||||
|
||||
// 外积：ZA += a × b^T
|
||||
// 这在一个指令中累加一个SVL×SVL瓦片
|
||||
svmopa_za32_f32_m(0, svptrue_b32(), svptrue_b32(), a, b);
|
||||
}
|
||||
|
||||
// 将ZA瓦片存储到内存
|
||||
// svst1_za(...);
|
||||
|
||||
// 退出流模式
|
||||
// smstop;
|
||||
}
|
||||
```
|
||||
|
||||
- **关键概念**：
|
||||
- **`svmopa`**（外积累加）：核心SME指令。它计算两个向量的完整外积并累加到ZA瓦片中。对于SVL=512位（16个浮点数），这是一个16×16外积——一条指令中256次FMA操作。
|
||||
- **ZA瓦片**：在流模式中跨指令持久存在。你将多个外积（每个K迭代一个）累加到同一瓦片中，构建完整的矩阵乘法瓦片。
|
||||
- **流模式**：SME指令仅在流模式下工作。进入/退出流模式的开销意味着SME最适合持续的矩阵计算，而非短时爆发。
|
||||
|
||||
- **SME2新增**：多向量操作（同时处理2或4个SVE向量）、额外的瓦片操作以及与普通模式的改进集成。
|
||||
|
||||
- **可用性**：ARM Neoverse V2（AWS Graviton4）、一些即将推出的移动芯片。截至2026年尚未出现在Apple Silicon上。SME仍处于早期阶段——大多数ML框架还没有SME优化的核函数。
|
||||
|
||||
- **演进脉络**：NEON（128位向量，逐元素）→ I8MM（INT8矩阵瓦片）→ SVE（可伸缩向量）→ SME（可伸缩2D矩阵瓦片）。每一代都更接近硬件原生矩阵操作。
|
||||
|
||||
## SVE和SVE2：可伸缩向量扩展
|
||||
|
||||
- NEON具有固定的128位宽度。**SVE**（可伸缩向量扩展）引入了**向量长度无关（VLA）编程**：你编写一次代码，它在任何向量宽度（128到2048位）的硬件上运行。硬件在运行时确定宽度。
|
||||
|
||||
```cpp
|
||||
#include <arm_sve.h>
|
||||
|
||||
void add_sve(const float* a, const float* b, float* c, int n) {
|
||||
int i = 0;
|
||||
svbool_t pred = svwhilelt_b32(i, n); // 谓词：哪些通道是激活的
|
||||
|
||||
while (svptest_any(svptrue_b32(), pred)) {
|
||||
svfloat32_t va = svld1(pred, a + i);
|
||||
svfloat32_t vb = svld1(pred, b + i);
|
||||
svst1(pred, c + i, svadd_x(pred, va, vb));
|
||||
|
||||
i += svcntw(); // 按硬件向量宽度前进（以32位元素计）
|
||||
pred = svwhilelt_b32(i, n);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- **谓词寄存器**（`svbool_t`）取代了标量清理循环。每个通道有一个谓词位：激活的通道参与，非激活的被屏蔽。`svwhilelt_b32(i, n)` 指令创建一个谓词，其中对应 `i, i+1, ..., n-1` 的通道被激活。这自动处理了尾部。
|
||||
|
||||
- **`svcntw()`** 在运行时返回每个向量寄存器中32位元素的数量。在具有256位SVE的CPU上，返回8。在512位SVE上，返回16。你的代码自动适应。
|
||||
|
||||
- SVE在ARM Neoverse V1/V2上可用（AWS Graviton3/4，一些服务器芯片）。在Apple Silicon上尚不可用。
|
||||
|
||||
## Apple Silicon特性
|
||||
|
||||
- Apple的M系列芯片（M1、M2、M3、M4）是基于ARM的自定义微架构：
|
||||
|
||||
- **性能核心和效率核心**：P核心（Firestorm/Avalanche等）用于重型计算，E核心（Icestorm/Blizzard等）用于后台任务。调度器将线程分配给适当的核心类型。
|
||||
|
||||
- **AMX**（Apple矩阵扩展）：专用矩阵乘法单元，独立于NEON。AMX未公开（Apple不发布ISA），但Accelerate框架内部将其用于BLAS操作。当你在Mac上调用 `np.dot` 时，它通过Accelerate，后者使用AMX。你不能直接对AMX编程（除非逆向工程）。
|
||||
|
||||
- **统一内存**：CPU和GPU共享同一物理RAM。在其他系统上，数据必须从CPU内存拷贝到GPU内存（通过PCIe，约32 GB/s）。在Apple Silicon上，无需拷贝——GPU读取CPU写入的同一内存。这消除了ML工作负载的主要瓶颈。
|
||||
|
||||
- **神经网络引擎**：一个16核专用ML加速器。INT8推理时达到约30 TOPS（每秒万亿次操作）。Core ML将其用于设备端推理。
|
||||
|
||||
- **Apple Silicon上的ML**：使用MLX（Apple的ML框架），它专为统一内存架构设计。PyTorch也有MPS（Metal性能着色器）后端支持，尽管不如CUDA成熟。
|
||||
|
||||
## 自动向量化
|
||||
|
||||
- 编写SIMD内联函数很繁琐。**编译器**能自动向量化你的代码吗？
|
||||
|
||||
- 可以的，但有限制。现代编译器（GCC、Clang）可以自动向量化简单循环：
|
||||
|
||||
```cpp
|
||||
// 编译器可以自动向量化此代码（使用 -O3 -march=native）
|
||||
void add_auto(const float* a, const float* b, float* c, int n) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- **有助于自动向量化的模式**：
|
||||
- 简单的循环，迭代次数已知。
|
||||
- 迭代之间无数据依赖（`c[i]` 不依赖于 `c[i-1]`）。
|
||||
- 连续内存访问（无分散/聚集）。
|
||||
- `const` 和 `restrict` 指针（告知编译器数组不重叠）。
|
||||
|
||||
```cpp
|
||||
// restrict 告诉编译器：a、b、c 指向不重叠的内存
|
||||
void add_restrict(const float* __restrict__ a,
|
||||
const float* __restrict__ b,
|
||||
float* __restrict__ c, int n) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- 没有 `restrict`，编译器必须假设 `c` 可能与 `a` 或 `b` 重叠（写入 `c[i]` 可能改变 `a[i+1]`），从而阻止向量化。
|
||||
|
||||
- **阻止自动向量化的模式**：
|
||||
- 数据依赖：`a[i] = a[i-1] + b[i]`（每次迭代依赖前一次）。
|
||||
- 复杂控制流：循环内的 `if` 语句（除非编译器能转换为谓词化）。
|
||||
- 循环内的函数调用（除非函数被内联）。
|
||||
- 指针别名（数组可能重叠，没有 `restrict`）。
|
||||
|
||||
- **检查自动向量化**：使用编译器标志查看哪些被向量化了：
|
||||
|
||||
```bash
|
||||
# GCC：显示向量化决策
|
||||
g++ -O3 -march=native -fopt-info-vec-optimized code.cpp
|
||||
|
||||
# Clang：显示向量化报告
|
||||
clang++ -O3 -march=native -Rpass=loop-vectorize code.cpp
|
||||
```
|
||||
|
||||
- **何时使用内联函数 vs 自动向量化**：从干净的C++和编译器优化开始。如果编译器向量化了你的循环，很好。如果性能仍不足，检查编译器的向量化报告以理解原因，然后仅为关键内循环编写内联函数。过早使用内联函数会让代码难以阅读而没有确定的收益。
|
||||
|
||||
## 编程任务（在ARM上用g++或clang++编译——Mac M系列或Linux aarch64）
|
||||
|
||||
1. 编写标量点积和NEON向量化点积。对两者进行基准测试并测量加速比。
|
||||
```cpp
|
||||
// task1_neon_dot.cpp
|
||||
// 编译（Mac/ARM Linux）：clang++ -O3 -o task1 task1_neon_dot.cpp
|
||||
// 注意：NEON在AArch64上默认启用，无需特殊标志
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
#include <arm_neon.h>
|
||||
|
||||
float dot_scalar(const float* a, const float* b, int n) {
|
||||
float sum = 0.0f;
|
||||
for (int i = 0; i < n; i++) {
|
||||
sum += a[i] * b[i];
|
||||
}
|
||||
return sum;
|
||||
}
|
||||
|
||||
float dot_neon(const float* a, const float* b, int n) {
|
||||
float32x4_t sum_vec = vdupq_n_f32(0.0f);
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
float32x4_t va = vld1q_f32(a + i);
|
||||
float32x4_t vb = vld1q_f32(b + i);
|
||||
sum_vec = vfmaq_f32(sum_vec, va, vb);
|
||||
}
|
||||
float sum = vaddvq_f32(sum_vec);
|
||||
for (; i < n; i++) sum += a[i] * b[i];
|
||||
return sum;
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 10'000'000;
|
||||
std::vector<float> a(N, 1.0f), b(N, 2.0f);
|
||||
|
||||
// 预热
|
||||
volatile float s1 = dot_scalar(a.data(), b.data(), N);
|
||||
volatile float s2 = dot_neon(a.data(), b.data(), N);
|
||||
|
||||
// 标量基准测试
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) {
|
||||
s1 = dot_scalar(a.data(), b.data(), N);
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double scalar_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
|
||||
// NEON基准测试
|
||||
start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) {
|
||||
s2 = dot_neon(a.data(), b.data(), N);
|
||||
}
|
||||
end = std::chrono::high_resolution_clock::now();
|
||||
double neon_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
|
||||
std::cout << "标量: " << scalar_ms << " ms（结果: " << s1 << "）\n";
|
||||
std::cout << "NEON: " << neon_ms << " ms（结果: " << s2 << "）\n";
|
||||
std::cout << "加速比: " << scalar_ms / neon_ms << "x\n";
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
2. 实现NEON ReLU和softmax最大值查找。练习使用不同操作的加载→计算→存储模式。
|
||||
```cpp
|
||||
// task2_neon_ops.cpp
|
||||
// 编译：clang++ -O3 -o task2 task2_neon_ops.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <cmath>
|
||||
#include <arm_neon.h>
|
||||
|
||||
void relu_neon(const float* in, float* out, int n) {
|
||||
float32x4_t zero = vdupq_n_f32(0.0f);
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
float32x4_t x = vld1q_f32(in + i);
|
||||
vst1q_f32(out + i, vmaxq_f32(x, zero));
|
||||
}
|
||||
for (; i < n; i++) out[i] = in[i] > 0 ? in[i] : 0;
|
||||
}
|
||||
|
||||
float max_neon(const float* data, int n) {
|
||||
float32x4_t max_vec = vdupq_n_f32(-INFINITY);
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
max_vec = vmaxq_f32(max_vec, vld1q_f32(data + i));
|
||||
}
|
||||
float result = vmaxvq_f32(max_vec);
|
||||
for (; i < n; i++) result = result > data[i] ? result : data[i];
|
||||
return result;
|
||||
}
|
||||
|
||||
int main() {
|
||||
std::vector<float> data = {-3, 1, -1, 4, 2, -5, 0, 7, -2, 3};
|
||||
std::vector<float> out(data.size());
|
||||
|
||||
relu_neon(data.data(), out.data(), data.size());
|
||||
std::cout << "ReLU: ";
|
||||
for (float x : out) std::cout << x << " ";
|
||||
std::cout << "\n";
|
||||
|
||||
float mx = max_neon(data.data(), data.size());
|
||||
std::cout << "最大值: " << mx << "（期望值: 7）\n";
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
3. 比较自动向量化代码与手写NEON内联函数。用 `-fopt-info-vec`（GCC）或 `-Rpass=loop-vectorize`（Clang）编译以查看编译器的操作。
|
||||
```cpp
|
||||
// task3_auto_vs_manual.cpp
|
||||
// 编译：clang++ -O3 -Rpass=loop-vectorize -o task3 task3_auto_vs_manual.cpp
|
||||
// （或）：g++ -O3 -fopt-info-vec-optimized -o task3 task3_auto_vs_manual.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
#include <arm_neon.h>
|
||||
|
||||
// 让编译器自动向量化
|
||||
void add_auto(const float* __restrict__ a, const float* __restrict__ b,
|
||||
float* __restrict__ c, int n) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
}
|
||||
|
||||
// 手写NEON
|
||||
void add_neon(const float* a, const float* b, float* c, int n) {
|
||||
int i = 0;
|
||||
for (; i + 4 <= n; i += 4) {
|
||||
vst1q_f32(c + i, vaddq_f32(vld1q_f32(a + i), vld1q_f32(b + i)));
|
||||
}
|
||||
for (; i < n; i++) c[i] = a[i] + b[i];
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 10'000'000;
|
||||
std::vector<float> a(N, 1.0f), b(N, 2.0f), c(N);
|
||||
|
||||
auto bench = [&](auto fn, const char* name) {
|
||||
fn(a.data(), b.data(), c.data(), N); // 预热
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) fn(a.data(), b.data(), c.data(), N);
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
std::cout << name << ": " << ms << " ms\n";
|
||||
};
|
||||
|
||||
bench(add_auto, "自动向量化");
|
||||
bench(add_neon, "手写NEON");
|
||||
// 它们应该非常接近——编译器能很好地自动向量化这个简单循环
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,450 @@
|
||||
# x86与AVX
|
||||
|
||||
*x86处理器来自Intel和AMD，主导着大多数ML训练所在的数据中心服务器。本文涵盖x86 SIMD的演进、AVX/AVX2内联函数编程、AVX-512、用于矩阵操作的Intel AMX、内存对齐、性能陷阱以及性能分析——在全球最常见的服务器CPU上榨取最大性能的工具。*
|
||||
|
||||
- 如果你的训练在云虚拟机（AWS、GCP、Azure）上运行，它几乎肯定运行在x86上。即使是GPU密集训练也有CPU瓶颈：数据加载、预处理、梯度聚合和检查点保存都在CPU上运行。使用x86 SIMD优化这些环节可以有意义地减少端到端训练时间。
|
||||
|
||||
## x86 SIMD演进
|
||||
|
||||
- x86 SIMD经历了越来越宽的向量寄存器：
|
||||
|
||||
| 代次 | 年份 | 寄存器宽度 | 寄存器数量 | 关键特性 |
|
||||
|------|------|---------|----------|----------|
|
||||
| MMX | 1997 | 64位 | 8（mm0-7） | 仅整数，与FPU共享 |
|
||||
| SSE | 1999 | 128位 | 8（xmm0-7） | 4个浮点数，专用寄存器 |
|
||||
| SSE2 | 2001 | 128位 | 8/16 | 2个双精度浮点数，整数操作 |
|
||||
| AVX | 2011 | 256位 | 16（ymm0-15） | 8个浮点数，三操作数指令 |
|
||||
| AVX2 | 2013 | 256位 | 16 | 整数256位，FMA，收集 |
|
||||
| AVX-512 | 2017 | 512位 | 32（zmm0-31） | 16个浮点数，掩码寄存器，分散 |
|
||||
| AMX | 2023 | 瓦片寄存器 | 8个瓦片 | 矩阵乘法（BF16，INT8） |
|
||||
|
||||
- 每一代都将向量化代码的吞吐量翻倍。用SSE内联函数编写的代码可以在2001年以来制造的每一个x86 CPU上运行。AVX2需要2013年以后的CPU。AVX-512需要Intel Xeon和一些消费级芯片。AMX是最新的（Sapphire Rapids及以后）。
|
||||
|
||||
- **向后兼容性**：x86 SSE寄存器（xmm）是AVX寄存器（ymm）的低128位，后者是AVX-512寄存器（zmm）的低256位。旧的SSE代码无需修改即可在新的CPU上运行。
|
||||
|
||||
## AVX2编程
|
||||
|
||||
- AVX2操作256位寄存器（YMM），同时处理8个浮点数或4个双精度浮点数。它是可移植高性能代码的甜点区域：在几乎所有现代x86 CPU（2013+）上可用。
|
||||
|
||||
### 内联函数命名约定
|
||||
|
||||
- 所有x86内联函数遵循模式：`_mm[宽度]_[操作]_[类型]`
|
||||
|
||||
- `_mm` = MMX/SSE（128位），`_mm256` = AVX（256位），`_mm512` = AVX-512（512位）
|
||||
- 操作：`add`、`mul`、`fmadd`、`load`、`store`、`set` 等
|
||||
- 类型：`ps` = 打包单精度（float32），`pd` = 打包双精度（float64），`epi32` = 打包int32，`si256` = 256位整数
|
||||
|
||||
```cpp
|
||||
#include <immintrin.h> // 所有x86 SIMD内联函数
|
||||
|
||||
// 数据类型
|
||||
__m256 a; // 256位寄存器，保存8个float32
|
||||
__m256d b; // 256位寄存器，保存4个float64
|
||||
__m256i c; // 256位寄存器，保存整数（8x32、16x16或32x8）
|
||||
```
|
||||
|
||||
### 加载和存储数据
|
||||
|
||||
```cpp
|
||||
// 从内存加载8个浮点数
|
||||
__m256 v = _mm256_loadu_ps(ptr); // 非对齐加载（适用于任何地址）
|
||||
__m256 v = _mm256_load_ps(ptr); // 对齐加载（ptr必须32字节对齐，更快）
|
||||
|
||||
// 存储8个浮点数到内存
|
||||
_mm256_storeu_ps(out_ptr, v); // 非对齐存储
|
||||
_mm256_store_ps(out_ptr, v); // 对齐存储
|
||||
|
||||
// 将单个值广播到所有8个通道
|
||||
__m256 ones = _mm256_set1_ps(1.0f); // [1, 1, 1, 1, 1, 1, 1, 1]
|
||||
|
||||
// 设置各个值（很少需要）
|
||||
__m256 v = _mm256_set_ps(7,6,5,4,3,2,1,0); // 注意：逆序！
|
||||
|
||||
// 零寄存器
|
||||
__m256 z = _mm256_setzero_ps();
|
||||
```
|
||||
|
||||
### 算术运算
|
||||
|
||||
```cpp
|
||||
__m256 c = _mm256_add_ps(a, b); // c[i] = a[i] + b[i]
|
||||
__m256 d = _mm256_mul_ps(a, b); // d[i] = a[i] * b[i]
|
||||
__m256 e = _mm256_sub_ps(a, b); // e[i] = a[i] - b[i]
|
||||
__m256 f = _mm256_div_ps(a, b); // f[i] = a[i] / b[i]（比乘法慢）
|
||||
|
||||
// 融合乘加：r = a * b + c（一条指令，一次舍入）
|
||||
__m256 r = _mm256_fmadd_ps(a, b, c); // ML最重要的指令
|
||||
|
||||
// 最小值和最大值
|
||||
__m256 mn = _mm256_min_ps(a, b); // min(a[i], b[i]) — 用于裁剪
|
||||
__m256 mx = _mm256_max_ps(a, b); // max(a[i], b[i]) — 用于ReLU
|
||||
```
|
||||
|
||||
### 实践示例：AVX2点积
|
||||
|
||||
```cpp
|
||||
#include <immintrin.h>
|
||||
|
||||
float dot_avx2(const float* a, const float* b, int n) {
|
||||
__m256 sum = _mm256_setzero_ps(); // 8个累加器初始化为0
|
||||
|
||||
int i = 0;
|
||||
for (; i + 8 <= n; i += 8) {
|
||||
__m256 va = _mm256_loadu_ps(a + i);
|
||||
__m256 vb = _mm256_loadu_ps(b + i);
|
||||
sum = _mm256_fmadd_ps(va, vb, sum); // sum += va * vb
|
||||
}
|
||||
|
||||
// 水平归约：将sum的8个元素相加
|
||||
// 步骤1：将上128位加到下128位
|
||||
__m128 hi = _mm256_extractf128_ps(sum, 1);
|
||||
__m128 lo = _mm256_castps256_ps128(sum);
|
||||
__m128 sum128 = _mm_add_ps(hi, lo); // 4个部分和
|
||||
|
||||
// 步骤2：在128位寄存器内水平相加
|
||||
sum128 = _mm_hadd_ps(sum128, sum128); // [a+b, c+d, a+b, c+d]
|
||||
sum128 = _mm_hadd_ps(sum128, sum128); // [a+b+c+d, ...]
|
||||
|
||||
float result = _mm_cvtss_f32(sum128); // 提取标量
|
||||
|
||||
// 标量清理
|
||||
for (; i < n; i++) {
|
||||
result += a[i] * b[i];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
```
|
||||
|
||||
- **为什么水平归约如此丑陋**：SIMD是为垂直操作设计的（通道0与通道0，通道1与通道1）。水平操作（跨通道求和）与硬件对抗。这就是点积在末尾有尴尬归约代码的原因。向量化循环是简洁的；归约是样板代码。
|
||||
|
||||
- **性能**：与NEON版本（文件02）相比，AVX2每次迭代处理8个浮点数，而NEON处理4个。对于长向量，这比NEON快2倍（忽略内存带宽限制）。
|
||||
|
||||
### 实践示例：AVX2 Softmax（简化版）
|
||||
|
||||
- Softmax需要：找到最大值，减去它，求指数，求和，除法。以下是最值查找步骤：
|
||||
|
||||
```cpp
|
||||
float vector_max_avx2(const float* data, int n) {
|
||||
__m256 max_vec = _mm256_set1_ps(-INFINITY);
|
||||
|
||||
int i = 0;
|
||||
for (; i + 8 <= n; i += 8) {
|
||||
__m256 v = _mm256_loadu_ps(data + i);
|
||||
max_vec = _mm256_max_ps(max_vec, v);
|
||||
}
|
||||
|
||||
// 将8个最大值归约为1个
|
||||
__m128 hi = _mm256_extractf128_ps(max_vec, 1);
|
||||
__m128 lo = _mm256_castps256_ps128(max_vec);
|
||||
__m128 max128 = _mm_max_ps(hi, lo);
|
||||
|
||||
// 通过混洗和取最大值找到单一最大值
|
||||
max128 = _mm_max_ps(max128, _mm_shuffle_ps(max128, max128, 0b01001110));
|
||||
max128 = _mm_max_ps(max128, _mm_shuffle_ps(max128, max128, 0b10110001));
|
||||
|
||||
float result = _mm_cvtss_f32(max128);
|
||||
|
||||
for (; i < n; i++) {
|
||||
result = result > data[i] ? result : data[i];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
```
|
||||
|
||||
- `_mm_shuffle_ps` 指令在寄存器内重排元素。二进制常量 `0b01001110` 控制哪些元素去哪里。这称为**置换**，它直接连接到置换矩阵（第2章）：打乱SIMD通道是做硬件级别的乘以置换矩阵。
|
||||
|
||||
## AVX-512
|
||||
|
||||
- AVX-512再次加倍宽度：512位寄存器（ZMM），同时处理16个浮点数。
|
||||
|
||||
```cpp
|
||||
__m512 a = _mm512_loadu_ps(ptr); // 加载16个浮点数
|
||||
__m512 c = _mm512_fmadd_ps(a, b, c); // 16个FMA同时进行
|
||||
float sum = _mm512_reduce_add_ps(a); // 内置水平求和（无需手动归约！）
|
||||
|
||||
// 掩码操作：操作通道子集
|
||||
__mmask16 mask = _mm512_cmpgt_ps_mask(a, zero); // 哪些通道 > 0？
|
||||
__m512 relu = _mm512_maskz_mov_ps(mask, a); // 负通道置零 = ReLU
|
||||
```
|
||||
|
||||
- **掩码寄存器**（`__mmask16`）是AVX-512最强大的功能。每个位控制一个通道是否参与操作。这取代了标量清理循环：最后一次迭代使用掩码，只有有效通道是激活的，处理任何向量长度而无需单独标量循环。
|
||||
|
||||
- **AVX-512频率降频**：在许多Intel CPU上，使用AVX-512指令会导致CPU暂时降低时钟频率（以保持在热限制内）。这意味着对于短时爆发，AVX-512并不总是比AVX2快——频率惩罚可能抵消更宽向量的优势。对于持续工作负载（如矩阵乘法），AVX-512胜出。对于混合代码（部分SIMD、部分标量），频率转换可能造成损失。
|
||||
|
||||
## Intel AMX：矩阵乘法硬件
|
||||
|
||||
- **AMX**（高级矩阵扩展）增加了专用矩阵乘法单元。AMX操作的不是SIMD向量，而是**瓦片**：2D数据块（最多16行 × 每行64字节）。
|
||||
|
||||
```cpp
|
||||
#include <immintrin.h>
|
||||
|
||||
// AMX瓦片乘法：C += A * B（BF16格式）
|
||||
// A为16x32 BF16，B为32x16 BF16，C为16x16 FP32
|
||||
_tile_loadd(0, a_ptr, stride_a); // 从A加载瓦片0
|
||||
_tile_loadd(1, b_ptr, stride_b); // 从B加载瓦片1
|
||||
_tile_dpbf16ps(2, 0, 1); // 瓦片2 += 瓦片0 * 瓦片1（BF16矩阵乘法，FP32累加）
|
||||
_tile_stored(2, c_ptr, stride_c); // 存储瓦片2到C
|
||||
```
|
||||
|
||||
- AMX在一条指令中执行完整的16×32 × 32×16矩阵乘法。这是数百次FMA操作同时进行，专门为Transformer推理中主导的小矩阵乘法设计（注意力得分计算、MLP层）。
|
||||
|
||||
- AMX支持BF16（bfloat16）和INT8，匹配ML推理中使用的精度。结合用于其他操作的AVX-512，配备AMX的CPU（Intel Sapphire Rapids、Emerald Rapids）可以在Transformer推理中与入门级GPU竞争。
|
||||
|
||||
## 内存对齐
|
||||
|
||||
- **对齐内存访问**是指数据地址是向量寄存器宽度的倍数（SSE为16字节、AVX为32字节、AVX-512为64字节）。对齐访问在某些CPU上更快，并且是 `_mm256_load_ps`（相对于 `_mm256_loadu_ps`）的要求。
|
||||
|
||||
```cpp
|
||||
// 分配对齐内存
|
||||
float* data = (float*)aligned_alloc(32, n * sizeof(float)); // AVX的32字节对齐
|
||||
|
||||
// C++对齐分配
|
||||
#include <new>
|
||||
float* data = new (std::align_val_t(32)) float[n];
|
||||
|
||||
// 或者使用编译器属性
|
||||
alignas(32) float data[1024];
|
||||
```
|
||||
|
||||
- **实际上**：在现代CPU（Haswell及以后）上，当数据不跨越缓存行边界时，非对齐加载（`loadu`）几乎与对齐加载一样快。非对齐访问的性能惩罚已基本消失，但缓存行分割（数据跨越两个64字节缓存行）仍可能使特定加载变慢约2倍。对齐分配完全避免了这种情况。
|
||||
|
||||
## 性能陷阱
|
||||
|
||||
- **AVX-SSE转换惩罚**：在较旧的Intel CPU（Skylake之前）上，在AVX（256位）和SSE（128位）指令之间切换会造成惩罚（约70周期）。这就是为什么你应该在从使用AVX的函数返回之前使用 `_mm256_zeroupper()`（或 `vzeroupper` 指令）清除YMM寄存器的上128位。现代CPU（Skylake+）没有此惩罚。
|
||||
|
||||
- **寄存器压力**：AVX2有16个YMM寄存器。如果你的核函数使用太多变量，编译器会将寄存器溢出到栈（内存），从而破坏性能。保持内循环简单，活变量少。
|
||||
|
||||
- **数据依赖**：`sum = _mm256_fmadd_ps(a, b, sum)` 对 `sum` 有依赖：每次迭代必须等待前一个FMA完成（约4-5个周期的延迟）。解决方案：使用多个独立累加器并在结束时归约：
|
||||
|
||||
```cpp
|
||||
// 单累加器：受FMA延迟限制（4-5个周期）
|
||||
__m256 sum = _mm256_setzero_ps();
|
||||
for (...) {
|
||||
sum = _mm256_fmadd_ps(a, b, sum); // 每个依赖前一个
|
||||
}
|
||||
|
||||
// 四个累加器：4倍吞吐量（隐藏延迟）
|
||||
__m256 sum0 = _mm256_setzero_ps();
|
||||
__m256 sum1 = _mm256_setzero_ps();
|
||||
__m256 sum2 = _mm256_setzero_ps();
|
||||
__m256 sum3 = _mm256_setzero_ps();
|
||||
for (...) {
|
||||
sum0 = _mm256_fmadd_ps(a0, b0, sum0); // 独立
|
||||
sum1 = _mm256_fmadd_ps(a1, b1, sum1); // 独立
|
||||
sum2 = _mm256_fmadd_ps(a2, b2, sum2); // 独立
|
||||
sum3 = _mm256_fmadd_ps(a3, b3, sum3); // 独立
|
||||
}
|
||||
sum0 = _mm256_add_ps(sum0, sum1);
|
||||
sum2 = _mm256_add_ps(sum2, sum3);
|
||||
sum0 = _mm256_add_ps(sum0, sum2);
|
||||
```
|
||||
|
||||
- 这是**循环展开**以隐藏延迟。CPU可以背靠背发出FMAs，因为它们写入不同的寄存器。这是数值代码中最有影响力的微优化之一。
|
||||
|
||||
## 性能分析
|
||||
|
||||
- **性能计数器**提供硬件级测量：
|
||||
|
||||
```bash
|
||||
# Linux perf（需要内核支持）
|
||||
perf stat ./my_program # 基本计数器：周期、指令、IPC
|
||||
perf stat -e cache-misses,cache-references ./my_program # 缓存行为
|
||||
perf record -g ./my_program && perf report # 调用图分析
|
||||
|
||||
# Intel VTune（详细的x86性能分析）
|
||||
vtune -collect hotspots -- ./my_program
|
||||
vtune -collect memory-access -- ./my_program # 内存带宽分析
|
||||
```
|
||||
|
||||
- **需要关注什么**：
|
||||
- **IPC**（每周期指令数）：CPU被使用的效率。IPC > 2 良好。IPC < 1 表明内存停顿或分支预测错误。
|
||||
- **缓存缺失率**：高L1/L2缺失率表明数据局部性差。需重构数据访问模式。
|
||||
- **分支预测错误率**：> 5% 表明分支不可预测。如可能，转换为无分支代码（SIMD比较+混合）。
|
||||
- **实际FLOPS vs 屋顶线**：将你的实测FLOPS与屋顶线模型（文件01）比较。如果你低于屋顶线，还有改进空间。
|
||||
|
||||
## 编程任务（在x86——Intel/AMD上用g++或clang++编译）
|
||||
|
||||
1. 编写标量点积和AVX2点积。对两者进行基准测试并测量8路SIMD带来的加速比。
|
||||
```cpp
|
||||
// task1_avx_dot.cpp
|
||||
// 编译：g++ -O3 -mavx2 -mfma -o task1 task1_avx_dot.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
#include <immintrin.h>
|
||||
|
||||
float dot_scalar(const float* a, const float* b, int n) {
|
||||
float sum = 0.0f;
|
||||
for (int i = 0; i < n; i++) sum += a[i] * b[i];
|
||||
return sum;
|
||||
}
|
||||
|
||||
float dot_avx2(const float* a, const float* b, int n) {
|
||||
__m256 sum = _mm256_setzero_ps();
|
||||
int i = 0;
|
||||
for (; i + 8 <= n; i += 8) {
|
||||
__m256 va = _mm256_loadu_ps(a + i);
|
||||
__m256 vb = _mm256_loadu_ps(b + i);
|
||||
sum = _mm256_fmadd_ps(va, vb, sum);
|
||||
}
|
||||
// 归约：上128加到下128，然后水平相加
|
||||
__m128 hi = _mm256_extractf128_ps(sum, 1);
|
||||
__m128 lo = _mm256_castps256_ps128(sum);
|
||||
__m128 r = _mm_add_ps(hi, lo);
|
||||
r = _mm_hadd_ps(r, r);
|
||||
r = _mm_hadd_ps(r, r);
|
||||
float result = _mm_cvtss_f32(r);
|
||||
for (; i < n; i++) result += a[i] * b[i];
|
||||
return result;
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 10'000'000;
|
||||
std::vector<float> a(N, 1.0f), b(N, 2.0f);
|
||||
|
||||
volatile float s1 = dot_scalar(a.data(), b.data(), N);
|
||||
volatile float s2 = dot_avx2(a.data(), b.data(), N);
|
||||
|
||||
auto bench = [&](auto fn, const char* name) {
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
volatile float s;
|
||||
for (int t = 0; t < 100; t++) s = fn(a.data(), b.data(), N);
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
std::cout << name << ": " << ms << " ms（结果: " << s << "）\n";
|
||||
return ms;
|
||||
};
|
||||
|
||||
double t1 = bench(dot_scalar, "标量");
|
||||
double t2 = bench(dot_avx2, "AVX2 ");
|
||||
std::cout << "加速比: " << t1 / t2 << "x\n";
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
2. 使用 `_mm256_max_ps` 实现AVX2 ReLU并与标量循环比较。然后尝试使用多累加器（循环展开）以隐藏FMA延迟。
|
||||
```cpp
|
||||
// task2_avx_relu.cpp
|
||||
// 编译：g++ -O3 -mavx2 -o task2 task2_avx_relu.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
#include <immintrin.h>
|
||||
|
||||
void relu_scalar(const float* in, float* out, int n) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
out[i] = in[i] > 0.0f ? in[i] : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
void relu_avx2(const float* in, float* out, int n) {
|
||||
__m256 zero = _mm256_setzero_ps();
|
||||
int i = 0;
|
||||
for (; i + 8 <= n; i += 8) {
|
||||
__m256 x = _mm256_loadu_ps(in + i);
|
||||
_mm256_storeu_ps(out + i, _mm256_max_ps(x, zero));
|
||||
}
|
||||
for (; i < n; i++) out[i] = in[i] > 0.0f ? in[i] : 0.0f;
|
||||
}
|
||||
|
||||
// 展开：每次迭代处理32个浮点数（4 x 8）
|
||||
void relu_avx2_unrolled(const float* in, float* out, int n) {
|
||||
__m256 zero = _mm256_setzero_ps();
|
||||
int i = 0;
|
||||
for (; i + 32 <= n; i += 32) {
|
||||
__m256 x0 = _mm256_loadu_ps(in + i);
|
||||
__m256 x1 = _mm256_loadu_ps(in + i + 8);
|
||||
__m256 x2 = _mm256_loadu_ps(in + i + 16);
|
||||
__m256 x3 = _mm256_loadu_ps(in + i + 24);
|
||||
_mm256_storeu_ps(out + i, _mm256_max_ps(x0, zero));
|
||||
_mm256_storeu_ps(out + i + 8, _mm256_max_ps(x1, zero));
|
||||
_mm256_storeu_ps(out + i + 16, _mm256_max_ps(x2, zero));
|
||||
_mm256_storeu_ps(out + i + 24, _mm256_max_ps(x3, zero));
|
||||
}
|
||||
for (; i + 8 <= n; i += 8) {
|
||||
_mm256_storeu_ps(out + i, _mm256_max_ps(_mm256_loadu_ps(in + i), zero));
|
||||
}
|
||||
for (; i < n; i++) out[i] = in[i] > 0.0f ? in[i] : 0.0f;
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 16'000'000;
|
||||
std::vector<float> in(N), out(N);
|
||||
for (int i = 0; i < N; i++) in[i] = (float)(i % 200) - 100.0f;
|
||||
|
||||
auto bench = [&](auto fn, const char* name) {
|
||||
fn(in.data(), out.data(), N); // 预热
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) fn(in.data(), out.data(), N);
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
double bw = 2.0 * N * sizeof(float) / ms / 1e6; // 读取+写入
|
||||
std::cout << name << ": " << ms << " ms（" << bw << " GB/s）\n";
|
||||
};
|
||||
|
||||
bench(relu_scalar, "标量 ");
|
||||
bench(relu_avx2, "AVX2 ");
|
||||
bench(relu_avx2_unrolled, "AVX2展开 ");
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
3. 测量内存对齐的效果。比较在大数组上的对齐加载与非对齐加载。
|
||||
```cpp
|
||||
// task3_alignment.cpp
|
||||
// 编译：g++ -O3 -mavx2 -o task3 task3_alignment.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <cstdlib>
|
||||
#include <immintrin.h>
|
||||
|
||||
int main() {
|
||||
const int N = 16'000'000;
|
||||
|
||||
// 对齐分配（AVX2为32字节）
|
||||
float* aligned = (float*)aligned_alloc(32, N * sizeof(float));
|
||||
|
||||
// 非对齐：从对齐边界偏移4字节（1个浮点数）
|
||||
float* raw = (float*)malloc((N + 1) * sizeof(float));
|
||||
float* unaligned = raw + 1; // 保证未对齐
|
||||
|
||||
for (int i = 0; i < N; i++) {
|
||||
aligned[i] = 1.0f;
|
||||
unaligned[i] = 1.0f;
|
||||
}
|
||||
|
||||
auto bench = [&](float* ptr, bool use_aligned, const char* name) {
|
||||
__m256 sum = _mm256_setzero_ps();
|
||||
// 预热
|
||||
for (int i = 0; i + 8 <= N; i += 8) {
|
||||
__m256 v = use_aligned ? _mm256_load_ps(ptr + i) : _mm256_loadu_ps(ptr + i);
|
||||
sum = _mm256_add_ps(sum, v);
|
||||
}
|
||||
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) {
|
||||
sum = _mm256_setzero_ps();
|
||||
for (int i = 0; i + 8 <= N; i += 8) {
|
||||
__m256 v = use_aligned ? _mm256_load_ps(ptr + i) : _mm256_loadu_ps(ptr + i);
|
||||
sum = _mm256_add_ps(sum, v);
|
||||
}
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
double bw = (double)N * sizeof(float) / ms / 1e6;
|
||||
std::cout << name << ": " << ms << " ms（" << bw << " GB/s）\n";
|
||||
};
|
||||
|
||||
bench(aligned, true, "对齐加载 ");
|
||||
bench(unaligned, false, "非对齐加载");
|
||||
|
||||
free(aligned);
|
||||
free(raw);
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,598 @@
|
||||
# GPU架构与CUDA
|
||||
|
||||
*GPU通过提供数千个核心用于大规模并行计算，改变了AI。本文涵盖GPU与CPU的设计哲学对比、GPU存储层次、C++中的CUDA编程、SIMT执行模型、内存访问模式、同步、流、性能分析以及NVIDIA GPU代次——编写和理解GPU核函数所需的知识。*
|
||||
|
||||
- 有关带有完整工作示例的实践CUDA教程，请参见配套仓库：[github.com/HenryNdubuaku/cuda-tutorials](https://github.com/HenryNdubuaku/cuda-tutorials)。
|
||||
|
||||
- 现代NVIDIA GPU有超过10,000个CUDA核心。CPU有4-128个核心。这100-1000倍的核心优势是GPU主导ML的原因：训练一个Transformer需要数万亿次乘加操作，GPU以CPU无法匹敌的规模并行处理它们。
|
||||
|
||||
- 即使你从不自己编写CUDA核函数，理解GPU架构也能解释：为什么批次大小很重要（需要足够的工作来饱和GPU），为什么内存通常是瓶颈（而非计算），以及为什么某些操作（分散、条件分支）在GPU上很慢。
|
||||
|
||||
## GPU vs CPU：根本不同的设计
|
||||
|
||||
- CPU是为**延迟**设计的：最小化完成一个任务的时间。它将其晶体管预算的大部分用于缓存、分支预测器和乱序执行——所有让单一线程快速运行的技巧。
|
||||
|
||||
- GPU是为**吞吐量**设计的：最大化每秒完成的任务数量。它将大部分晶体管用于执行单元（ALU）。单个线程很慢，但有数千个。
|
||||
|
||||
| | CPU | GPU |
|
||||
|--|-----|-----|
|
||||
| 核心 | 4-128（复杂、快速） | 1,000-20,000（简单、慢速） |
|
||||
| 时钟频率 | 3-5 GHz | 1-2.5 GHz |
|
||||
| 缓存 | 大（32 MB+ L3） | 小（每SM共享内存） |
|
||||
| 分支预测 | 精密 | 无（所有线程遵循相同路径） |
|
||||
| 最适合 | 低延迟、复杂控制流 | 高吞吐量、数据并行工作 |
|
||||
| 典型FLOPS（FP32） | 1-5 TFLOPS | 30-80 TFLOPS |
|
||||
| 内存带宽 | 50-100 GB/s | 1-3 TB/s |
|
||||
|
||||
- GPU的内存带宽优势（10-30倍）通常比其计算优势更重要。许多ML操作是内存受限的（逐元素操作、归一化、注意力），GPU的带宽使其能够足够快地向核心输送数据。
|
||||
|
||||
## GPU存储层次
|
||||
|
||||
- 理解GPU内存至关重要，因为**内存访问是主要瓶颈**，而非计算。
|
||||
|
||||
| 内存 | 大小 | 延迟 | 带宽 | 作用域 |
|
||||
|--------|------|---------|-----------|-------|
|
||||
| 寄存器 | 每SM约256 KB | 0周期 | 最高 | 每线程 |
|
||||
| 共享内存 | 每SM 48-228 KB | 约5周期 | 约20 TB/s | 每线程块 |
|
||||
| L1缓存 | 每SM 128-256 KB | 约30周期 | | 每SM |
|
||||
| L2缓存 | 4-96 MB | 约200周期 | 约6 TB/s | 全局 |
|
||||
| 全局内存（HBM） | 24-192 GB | 约400周期 | 1-3.3 TB/s | 全局 |
|
||||
|
||||
- **寄存器**是最快但最有限的。每个线程有一组私有寄存器（通常最多255个）。每线程使用过多寄存器会降低**占用率**（可同时运行的线程更少）。
|
||||
|
||||
- **共享内存**是由程序员管理的缓存，由块中的所有线程共享。它是编写快速CUDA核函数的关键：将数据瓦片从慢速全局内存加载到快速共享内存，然后进行计算。这是主导GPU编程的**分块**模式。
|
||||
|
||||
- **全局内存（HBM）**：主GPU内存（VRAM）。大但慢（400周期延迟）。所有数据起始和结束于此。核函数优化的目标是尽量减少全局内存访问。
|
||||
|
||||
## CUDA编程模型
|
||||
|
||||
- CUDA（统一计算设备架构）是NVIDIA的GPU编程模型。你编写**核函数**：在GPU上运行的函数，由数千个线程同时执行。
|
||||
|
||||
### 层次结构：网格、块、线程
|
||||
|
||||
```
|
||||
网格（整个启动）
|
||||
├── 块 (0,0)
|
||||
│ ├── 线程 (0,0)
|
||||
│ ├── 线程 (1,0)
|
||||
│ ├── 线程 (2,0)
|
||||
│ └── ... （每块最多1024线程）
|
||||
├── 块 (1,0)
|
||||
│ ├── 线程 (0,0)
|
||||
│ └── ...
|
||||
└── ... （可能有数百万个块）
|
||||
```
|
||||
|
||||
- **线程**：最小单位。每个线程在其块内有唯一ID（`threadIdx.x`）。
|
||||
- **块**：一组可以共享内存和同步的线程。块ID：`blockIdx.x`。块大小：`blockDim.x`（最多1024线程）。
|
||||
- **网格**：单个核函数启动的所有块。可以是1D、2D或3D。
|
||||
|
||||
- 每个线程计算其全局索引：`int idx = blockIdx.x * blockDim.x + threadIdx.x;`
|
||||
|
||||
### 你的第一个CUDA核函数
|
||||
|
||||
```cpp
|
||||
// vector_add.cu — CUDA源文件（.cu扩展名）
|
||||
|
||||
#include <stdio.h>
|
||||
|
||||
// __global__ 标记此为GPU核函数（从CPU调用，在GPU上运行）
|
||||
__global__ void vector_add(const float* a, const float* b, float* c, int n) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < n) { // 边界检查（网格可能大于数据）
|
||||
c[idx] = a[idx] + b[idx];
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
int n = 1 << 20; // 约100万个元素
|
||||
size_t bytes = n * sizeof(float);
|
||||
|
||||
// 分配主机（CPU）内存
|
||||
float *h_a = new float[n];
|
||||
float *h_b = new float[n];
|
||||
float *h_c = new float[n];
|
||||
|
||||
// 初始化
|
||||
for (int i = 0; i < n; i++) {
|
||||
h_a[i] = 1.0f;
|
||||
h_b[i] = 2.0f;
|
||||
}
|
||||
|
||||
// 分配设备（GPU）内存
|
||||
float *d_a, *d_b, *d_c;
|
||||
cudaMalloc(&d_a, bytes);
|
||||
cudaMalloc(&d_b, bytes);
|
||||
cudaMalloc(&d_c, bytes);
|
||||
|
||||
// 将数据从CPU拷贝到GPU
|
||||
cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice);
|
||||
|
||||
// 启动核函数：每块256线程，足够的块覆盖n个元素
|
||||
int block_size = 256;
|
||||
int grid_size = (n + block_size - 1) / block_size; // 上取整除法
|
||||
vector_add<<<grid_size, block_size>>>(d_a, d_b, d_c, n);
|
||||
|
||||
// 将结果从GPU拷贝到CPU
|
||||
cudaMemcpy(h_c, d_a, bytes, cudaMemcpyDeviceToHost);
|
||||
|
||||
// 验证
|
||||
printf("c[0] = %f（期望值 3.0）\n", h_c[0]);
|
||||
|
||||
// 释放内存
|
||||
cudaFree(d_a); cudaFree(d_b); cudaFree(d_c);
|
||||
delete[] h_a; delete[] h_b; delete[] h_c;
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
# 用NVIDIA编译器编译
|
||||
nvcc -O3 -o vector_add vector_add.cu
|
||||
./vector_add
|
||||
```
|
||||
|
||||
- **CUDA中的关键C++概念**：
|
||||
- `__global__`：CUDA关键字，标记核函数。从CPU（主机）调用，在GPU（设备）上运行。
|
||||
- `<<<grid_size, block_size>>>`：核函数启动语法。指定使用多少块和线程。
|
||||
- `cudaMalloc` / `cudaFree`：分配/释放GPU内存（类似于`new`/`delete`，但针对GPU）。
|
||||
- `cudaMemcpy`：在CPU和GPU之间拷贝数据。这通常是最大的瓶颈（PCIe带宽约32 GB/s，而GPU内存带宽约3 TB/s）。
|
||||
|
||||
### 线程束与SIMT
|
||||
|
||||
- GPU以32个为一组称为**线程束**的组执行线程。一个线程束中的所有32个线程同时执行**相同指令**（单指令多线程——SIMT）。这是GPU的SIMD等效，但在线程级别。
|
||||
|
||||
- **线程束分歧**发生在同一线程束中的线程在`if`语句中走不同分支时。GPU不能在一个线程束中同时执行两条不同指令，因此它顺序执行两个分支，屏蔽掉不应参与的线程。这使性能减半（或更差）。
|
||||
|
||||
```cpp
|
||||
// 糟糕：线程束分歧（同一线程束中的线程走不同路径）
|
||||
if (threadIdx.x % 2 == 0) {
|
||||
c[idx] = a[idx] + b[idx]; // 偶数线程做这个
|
||||
} else {
|
||||
c[idx] = a[idx] - b[idx]; // 奇数线程做这个（同一线程束，串行化）
|
||||
}
|
||||
|
||||
// 更好：无分支（无分歧）
|
||||
float sign = (threadIdx.x % 2 == 0) ? 1.0f : -1.0f;
|
||||
c[idx] = a[idx] + sign * b[idx]; // 所有线程执行相同指令
|
||||
```
|
||||
|
||||
### 内存合并
|
||||
|
||||
- **合并访问**：当连续的线程访问连续的内存地址时，GPU将它们组合成单个内存事务。这对性能至关重要。
|
||||
|
||||
```cpp
|
||||
// 好：合并——线程0读a[0]，线程1读a[1]，...
|
||||
c[idx] = a[idx] + b[idx];
|
||||
|
||||
// 坏：跨步——线程0读a[0]，线程1读a[步长]，...
|
||||
c[idx] = a[idx * stride] + b[idx * stride]; // 步长 > 1 浪费带宽
|
||||
```
|
||||
|
||||
- 对于一个32线程的线程束，合并访问在单次事务中加载128字节（32 × 4字节用于float32）。跨步访问需要多次事务，每次加载128字节但只使用一小部分。步长为32是最坏情况：每次事务加载128字节，但只有一个线程使用4字节（3%的利用率）。
|
||||
|
||||
### 共享内存与分块
|
||||
|
||||
- **分块模式**是最重要的GPU优化技术。其想法：将数据块从慢速全局内存加载到快速共享内存，进行计算，然后将结果写回。
|
||||
|
||||
```cpp
|
||||
// 使用共享内存分块的矩阵乘法（简化版）
|
||||
__global__ void matmul_tiled(const float* A, const float* B, float* C,
|
||||
int M, int N, int K) {
|
||||
// A的一个瓦片和B的一个瓦片的共享内存
|
||||
__shared__ float tile_A[TILE_SIZE][TILE_SIZE];
|
||||
__shared__ float tile_B[TILE_SIZE][TILE_SIZE];
|
||||
|
||||
int row = blockIdx.y * TILE_SIZE + threadIdx.y;
|
||||
int col = blockIdx.x * TILE_SIZE + threadIdx.x;
|
||||
float sum = 0.0f;
|
||||
|
||||
// 遍历瓦片
|
||||
for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
|
||||
// 将A和B的一个瓦片加载到共享内存
|
||||
if (row < M && t * TILE_SIZE + threadIdx.x < K)
|
||||
tile_A[threadIdx.y][threadIdx.x] = A[row * K + t * TILE_SIZE + threadIdx.x];
|
||||
else
|
||||
tile_A[threadIdx.y][threadIdx.x] = 0.0f;
|
||||
|
||||
if (col < N && t * TILE_SIZE + threadIdx.y < K)
|
||||
tile_B[threadIdx.y][threadIdx.x] = B[(t * TILE_SIZE + threadIdx.y) * N + col];
|
||||
else
|
||||
tile_B[threadIdx.y][threadIdx.x] = 0.0f;
|
||||
|
||||
__syncthreads(); // 等待所有线程完成加载
|
||||
|
||||
// 计算此瓦片的部分点积
|
||||
for (int k = 0; k < TILE_SIZE; k++) {
|
||||
sum += tile_A[threadIdx.y][k] * tile_B[k][threadIdx.x];
|
||||
}
|
||||
|
||||
__syncthreads(); // 在加载下一个瓦片前等待
|
||||
}
|
||||
|
||||
if (row < M && col < N)
|
||||
C[row * N + col] = sum;
|
||||
}
|
||||
```
|
||||
|
||||
- **`__shared__`**：声明块内所有线程共享的内存（快速、片上）。
|
||||
- **`__syncthreads()`**：一个屏障，等待块中所有线程到达此点。在写入共享内存和读取它之间必须使用（否则某些线程读取到过期数据）。
|
||||
- **为什么分块有效**：没有它，每个线程每次乘法都从全局内存加载。有了分块，一个TILE_SIZE × TILE_SIZE的数据块被加载到共享内存一次，并被块中所有线程重用。重用因子为TILE_SIZE，将全局内存流量减少该因子。
|
||||
|
||||
## 流与并发
|
||||
|
||||
- 默认情况下，CUDA操作是顺序的：CPU启动一个核函数，等待它完成，然后启动下一个。**流**允许重叠：
|
||||
|
||||
```cpp
|
||||
cudaStream_t stream1, stream2;
|
||||
cudaStreamCreate(&stream1);
|
||||
cudaStreamCreate(&stream2);
|
||||
|
||||
// 这些操作可以重叠：不同流并发执行
|
||||
cudaMemcpyAsync(d_a, h_a, bytes, cudaMemcpyHostToDevice, stream1);
|
||||
cudaMemcpyAsync(d_b, h_b, bytes, cudaMemcpyHostToDevice, stream2);
|
||||
|
||||
kernel1<<<grid, block, 0, stream1>>>(d_a, d_c);
|
||||
kernel2<<<grid, block, 0, stream2>>>(d_b, d_d);
|
||||
```
|
||||
|
||||
- 流将数据传输与计算重叠：当流1的核函数运行时，流2在拷贝数据。这隐藏了PCIe传输延迟，并保持GPU忙碌。
|
||||
|
||||
## 分析CUDA代码
|
||||
|
||||
```bash
|
||||
# NVIDIA Nsight Compute：核函数级分析
|
||||
ncu --set full ./my_program
|
||||
|
||||
# NVIDIA Nsight Systems：系统级时间线
|
||||
nsys profile ./my_program
|
||||
|
||||
# 快速指标
|
||||
ncu --metrics sm__throughput,dram__throughput ./my_program
|
||||
```
|
||||
|
||||
- **需要关注什么**：
|
||||
- **占用率**：SM容量中被使用的比例。低占用率（< 50%）意味着线程太少，无法隐藏内存延迟。原因：每线程寄存器过多、每块共享内存过多。
|
||||
- **内存吞吐量**：与峰值带宽比较。如果你达到峰值带宽的50%以下，内存访问模式低效（非合并、存储体冲突）。
|
||||
- **计算吞吐量**：与峰值FLOPS比较。如果内存和计算吞吐量都低，核函数是延迟受限的（并行度不够）。
|
||||
|
||||
## 高级优化技术
|
||||
|
||||
- 除了合并和共享内存分块的基础知识外，高性能GPU（和CPU）代码还使用几种高级技术：
|
||||
|
||||
### 数据布局：AoS vs SoA
|
||||
|
||||
- **结构体数组（AoS）**：每个元素将所有字段存储在一起。`[{x,y,z}, {x,y,z}, {x,y,z}]`。
|
||||
- **数组结构体（SoA）**：每个字段存储在自己的连续数组中。`{[x,x,x], [y,y,y], [z,z,z]}`。
|
||||
|
||||
```cpp
|
||||
// AoS：对于SIMD/GPU不好（访问所有x值触及非连续内存）
|
||||
struct Particle { float x, y, z, mass; };
|
||||
Particle particles[N];
|
||||
// particles[0].x, particles[1].x 相隔16字节
|
||||
|
||||
// SoA：对于SIMD/GPU好（所有x值连续）
|
||||
struct Particles {
|
||||
float x[N], y[N], z[N], mass[N];
|
||||
};
|
||||
// x[0], x[1] 相隔4字节——非常适合合并访问和SIMD
|
||||
```
|
||||
|
||||
- SoA对于数据并行工作负载（SIMD、GPU）几乎总是更快。AoS在你总是同时访问一个元素的所有字段时更好（在数值代码中很少见）。PyTorch张量本质上是SoA：每个特征是一个连续维度。
|
||||
|
||||
### 软件预取
|
||||
|
||||
- 可以告诉CPU在需要之前开始加载数据，隐藏内存延迟：
|
||||
|
||||
```cpp
|
||||
#include <xmmintrin.h> // for _mm_prefetch
|
||||
|
||||
for (int i = 0; i < n; i += 4) {
|
||||
_mm_prefetch((char*)(a + i + 64), _MM_HINT_T0); // 预取之前64个元素
|
||||
// 用SIMD处理 a[i:i+4]
|
||||
__m128 va = _mm_load_ps(a + i);
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
- 预取指令是一个提示：如果数据已在缓存中，它是空操作。如果不是，CPU在执行其他指令的同时开始在后台获取数据。预取距离（此示例中向前64个元素）应根据内存延迟和循环迭代时间进行调整。
|
||||
|
||||
### 核函数融合
|
||||
|
||||
- **核函数融合**将多个操作组合成一个核函数，以避免将中间结果写入内存。这是ML中最有影响力的单个GPU优化：
|
||||
|
||||
```
|
||||
// 未融合：3次核函数启动，3次全局内存往返
|
||||
y = matmul(x, W) // 写y到全局内存
|
||||
z = y + bias // 读y，写z
|
||||
out = relu(z) // 读z，写out
|
||||
|
||||
// 融合：1次核函数启动，1次全局内存写入
|
||||
out = fused_matmul_bias_relu(x, W, bias) // y和z永不离开SRAM
|
||||
```
|
||||
|
||||
- 对于内存受限操作（偏置加法、ReLU、层归一化），内存流量主导执行时间。融合完全消除了流量。PyTorch的`torch.compile`和Triton可以自动或通过最少努力实现融合。
|
||||
|
||||
### 混合精度核函数
|
||||
|
||||
- 使用较低精度（FP16、BF16、INT8）进行计算和较高精度（FP32）进行累加，达到两全其美：
|
||||
|
||||
```cpp
|
||||
// 张量核心：乘FP16矩阵，在FP32中累加
|
||||
// 每条张量核心指令：D（FP32）= A（FP16）× B（FP16）+ C（FP32）
|
||||
nvcuda::wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
|
||||
```
|
||||
|
||||
- FP16比FP32小2倍，因此它使内存带宽加倍（通常的瓶颈），并在缓存中容纳2倍的数据。张量核心以FP32 CUDA核心8-16倍的速度处理FP16。这就是为什么混合精度训练（第6章）提供2-3倍加速且精度损失最小。
|
||||
|
||||
### 内存池分配器
|
||||
|
||||
- `cudaMalloc` 很慢（每次调用约1毫秒），因为它与GPU同步。在每次迭代分配临时缓冲区的训练循环中，这会累积起来。
|
||||
|
||||
- **内存池**（PyTorch的缓存分配器、CUDA内存池）预先分配一大块GPU内存，并从其中子分配而无需系统调用：
|
||||
|
||||
```python
|
||||
# PyTorch自动执行此操作——但理解原因很重要
|
||||
# 每个 torch.empty() 从池中重用内存，无需cudaMalloc
|
||||
temp = torch.empty(1024, 1024, device='cuda') # 微秒，而非毫秒
|
||||
```
|
||||
|
||||
- 这就是为什么PyTorch的 `torch.cuda.memory_allocated()` 和 `torch.cuda.max_memory_allocated()` 不同：allocated是当前使用的，max是峰值（池可能持有比当前使用更多的内存）。
|
||||
|
||||
### 分析指导的优化
|
||||
|
||||
- 不要盲目优化。**先分析**，识别瓶颈，优化那个，然后重新分析。屋顶线模型（文件01）告诉你瓶颈是内存还是计算：
|
||||
|
||||
- **内存受限**（低算术强度）：优化数据布局（SoA）、融合核函数、使用较低精度、预取。
|
||||
- **计算受限**（高算术强度）：使用张量核心、增加并行度、使用更快指令（FMA）。
|
||||
- **延迟受限**（并行度不足）：增加占用率、减少寄存器使用、启动更多线程。
|
||||
|
||||
- 大多数ML工作负载是**内存受限的**。令人惊讶的推论：更快的GPU（更多FLOPS）通常没有帮助。更快的内存（HBM3 vs HBM2e）更有帮助。这就是为什么A100→H100升级不只是关于FLOPS——H100也有2倍的内存带宽。
|
||||
|
||||
## NVIDIA GPU代次
|
||||
|
||||
| 代次 | 年份 | 关键创新 | AI相关性 |
|
||||
|---------|------|-------------|--------------|
|
||||
| Pascal（P100） | 2016 | HBM2、NVLink | 第一代严肃的深度学习GPU |
|
||||
| Volta（V100） | 2017 | **张量核心**（混合精度矩阵乘法） | 实现FP16训练，125 TFLOPS TF32 |
|
||||
| Ampere（A100） | 2020 | TF32、稀疏性、第三代张量核心 | 312 TFLOPS TF32，结构性稀疏2:4 |
|
||||
| Hopper（H100） | 2022 | **Transformer引擎**（FP8）、HBM3 | 989 TFLOPS FP8，动态精度切换 |
|
||||
| Blackwell（B200） | 2024 | 第二代Transformer引擎、NVLink 5 | 2.5 PFLOPS FP4，多芯片设计 |
|
||||
|
||||
- **张量核心**是专用的矩阵乘法单元。单个张量核心指令在一个周期内计算4×4矩阵乘法（D = A×B + C）。常规CUDA核心需要64次FMA操作。张量核心就是为什么混合精度训练（float16计算，float32累加）如此快速。
|
||||
|
||||
- **Transformer引擎**（Hopper+）在单层内动态切换FP8和FP16精度，只在需要时选择更高精度。这最大化吞吐量而不牺牲模型质量。它专为Transformer架构（注意力+MLP）设计，后者主导现代AI。
|
||||
|
||||
## 编程任务（用nvcc编译）
|
||||
|
||||
1. 编写一个对数组应用ReLU的CUDA核函数。测量包括内存传输在内的时间。这教授核函数编写、cudaMalloc/cudaMemcpy以及主机↔设备传输瓶颈。
|
||||
```cpp
|
||||
// task1_relu.cu
|
||||
// 编译：nvcc -O3 -o task1_relu task1_relu.cu
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__global__ void relu_kernel(const float* input, float* output, int n) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < n) {
|
||||
output[idx] = input[idx] > 0.0f ? input[idx] : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 1 << 24; // 约1600万元素
|
||||
size_t bytes = N * sizeof(float);
|
||||
|
||||
// 分配主机内存
|
||||
float* h_input = (float*)malloc(bytes);
|
||||
float* h_output = (float*)malloc(bytes);
|
||||
for (int i = 0; i < N; i++) {
|
||||
h_input[i] = (float)(i % 100) - 50.0f; // 正负混合
|
||||
}
|
||||
|
||||
// 分配设备内存
|
||||
float *d_input, *d_output;
|
||||
cudaMalloc(&d_input, bytes);
|
||||
cudaMalloc(&d_output, bytes);
|
||||
|
||||
// 计时完整流水线：拷贝到GPU、计算、拷贝回
|
||||
cudaEvent_t start, stop;
|
||||
cudaEventCreate(&start);
|
||||
cudaEventCreate(&stop);
|
||||
|
||||
cudaEventRecord(start);
|
||||
cudaMemcpy(d_input, h_input, bytes, cudaMemcpyHostToDevice);
|
||||
|
||||
int block_size = 256;
|
||||
int grid_size = (N + block_size - 1) / block_size;
|
||||
relu_kernel<<<grid_size, block_size>>>(d_input, d_output, N);
|
||||
|
||||
cudaMemcpy(h_output, d_output, bytes, cudaMemcpyDeviceToHost);
|
||||
cudaEventRecord(stop);
|
||||
cudaEventSynchronize(stop);
|
||||
|
||||
float ms = 0;
|
||||
cudaEventElapsedTime(&ms, start, stop);
|
||||
|
||||
// 验证
|
||||
int errors = 0;
|
||||
for (int i = 0; i < N; i++) {
|
||||
float expected = h_input[i] > 0.0f ? h_input[i] : 0.0f;
|
||||
if (h_output[i] != expected) errors++;
|
||||
}
|
||||
|
||||
printf("时间（含传输）: %.2f ms\n", ms);
|
||||
printf("带宽: %.1f GB/s\n", 2.0 * bytes / ms / 1e6); // 读取+写入
|
||||
printf("错误: %d / %d\n", errors, N);
|
||||
|
||||
cudaFree(d_input); cudaFree(d_output);
|
||||
free(h_input); free(h_output);
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
2. 在CUDA中使用共享内存编写分块矩阵乘法。将性能与朴素（非分块）版本进行比较。这教授共享内存、`__syncthreads`以及为什么分块重要。
|
||||
```cpp
|
||||
// task2_matmul.cu
|
||||
// 编译：nvcc -O3 -o task2_matmul task2_matmul.cu
|
||||
|
||||
#include <stdio.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define TILE 16
|
||||
|
||||
// 朴素矩阵乘法：每个线程计算C的一个元素
|
||||
__global__ void matmul_naive(const float* A, const float* B, float* C, int N) {
|
||||
int row = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int col = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (row < N && col < N) {
|
||||
float sum = 0.0f;
|
||||
for (int k = 0; k < N; k++) {
|
||||
sum += A[row * N + k] * B[k * N + col];
|
||||
}
|
||||
C[row * N + col] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
// 分块矩阵乘法：使用共享内存减少全局内存访问
|
||||
__global__ void matmul_tiled(const float* A, const float* B, float* C, int N) {
|
||||
__shared__ float sA[TILE][TILE];
|
||||
__shared__ float sB[TILE][TILE];
|
||||
|
||||
int row = blockIdx.y * TILE + threadIdx.y;
|
||||
int col = blockIdx.x * TILE + threadIdx.x;
|
||||
float sum = 0.0f;
|
||||
|
||||
for (int t = 0; t < (N + TILE - 1) / TILE; t++) {
|
||||
sA[threadIdx.y][threadIdx.x] = (row < N && t*TILE+threadIdx.x < N)
|
||||
? A[row * N + t*TILE + threadIdx.x] : 0.0f;
|
||||
sB[threadIdx.y][threadIdx.x] = (t*TILE+threadIdx.y < N && col < N)
|
||||
? B[(t*TILE + threadIdx.y) * N + col] : 0.0f;
|
||||
|
||||
__syncthreads();
|
||||
for (int k = 0; k < TILE; k++)
|
||||
sum += sA[threadIdx.y][k] * sB[k][threadIdx.x];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (row < N && col < N)
|
||||
C[row * N + col] = sum;
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 1024;
|
||||
size_t bytes = N * N * sizeof(float);
|
||||
|
||||
float *d_A, *d_B, *d_C;
|
||||
cudaMalloc(&d_A, bytes); cudaMalloc(&d_B, bytes); cudaMalloc(&d_C, bytes);
|
||||
|
||||
// 初始化为1（容易验证：C应全为N）
|
||||
float* h_A = new float[N*N];
|
||||
for (int i = 0; i < N*N; i++) h_A[i] = 1.0f;
|
||||
cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d_B, h_A, bytes, cudaMemcpyHostToDevice);
|
||||
|
||||
dim3 block(TILE, TILE);
|
||||
dim3 grid((N+TILE-1)/TILE, (N+TILE-1)/TILE);
|
||||
|
||||
// 基准测试朴素版
|
||||
cudaEvent_t start, stop;
|
||||
cudaEventCreate(&start); cudaEventCreate(&stop);
|
||||
|
||||
cudaEventRecord(start);
|
||||
for (int i = 0; i < 10; i++)
|
||||
matmul_naive<<<grid, block>>>(d_A, d_B, d_C, N);
|
||||
cudaEventRecord(stop);
|
||||
cudaEventSynchronize(stop);
|
||||
float naive_ms; cudaEventElapsedTime(&naive_ms, start, stop);
|
||||
|
||||
// 基准测试分块版
|
||||
cudaEventRecord(start);
|
||||
for (int i = 0; i < 10; i++)
|
||||
matmul_tiled<<<grid, block>>>(d_A, d_B, d_C, N);
|
||||
cudaEventRecord(stop);
|
||||
cudaEventSynchronize(stop);
|
||||
float tiled_ms; cudaEventElapsedTime(&tiled_ms, start, stop);
|
||||
|
||||
double gflops_naive = 2.0 * N * N * N * 10 / naive_ms / 1e6;
|
||||
double gflops_tiled = 2.0 * N * N * N * 10 / tiled_ms / 1e6;
|
||||
|
||||
printf("朴素版: %.2f ms, %.1f GFLOPS\n", naive_ms/10, gflops_naive);
|
||||
printf("分块版: %.2f ms, %.1f GFLOPS\n", tiled_ms/10, gflops_tiled);
|
||||
printf("加速比: %.1fx\n", naive_ms / tiled_ms);
|
||||
|
||||
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
|
||||
delete[] h_A;
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
3. 演示线程束分歧。编写一个核函数，其中同一线程束中的线程走不同分支，并与无分支版本比较。
|
||||
```cpp
|
||||
// task3_divergence.cu
|
||||
// 编译：nvcc -O3 -o task3_diverge task3_divergence.cu
|
||||
|
||||
#include <stdio.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 糟糕：线程束分歧——偶数/奇数线程走不同路径
|
||||
__global__ void divergent_kernel(float* data, int n) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < n) {
|
||||
if (idx % 2 == 0) {
|
||||
data[idx] = data[idx] * 2.0f + 1.0f;
|
||||
} else {
|
||||
data[idx] = data[idx] * 0.5f - 1.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 好：无分支——所有线程执行相同指令
|
||||
__global__ void branchless_kernel(float* data, int n) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < n) {
|
||||
float scale = (idx % 2 == 0) ? 2.0f : 0.5f;
|
||||
float offset = (idx % 2 == 0) ? 1.0f : -1.0f;
|
||||
data[idx] = data[idx] * scale + offset;
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 1 << 24;
|
||||
float* d_data;
|
||||
cudaMalloc(&d_data, N * sizeof(float));
|
||||
cudaMemset(d_data, 0, N * sizeof(float));
|
||||
|
||||
int block = 256, grid = (N + block - 1) / block;
|
||||
|
||||
cudaEvent_t start, stop;
|
||||
cudaEventCreate(&start); cudaEventCreate(&stop);
|
||||
|
||||
// 分歧版
|
||||
cudaEventRecord(start);
|
||||
for (int i = 0; i < 100; i++)
|
||||
divergent_kernel<<<grid, block>>>(d_data, N);
|
||||
cudaEventRecord(stop);
|
||||
cudaEventSynchronize(stop);
|
||||
float div_ms; cudaEventElapsedTime(&div_ms, start, stop);
|
||||
|
||||
// 无分支版
|
||||
cudaEventRecord(start);
|
||||
for (int i = 0; i < 100; i++)
|
||||
branchless_kernel<<<grid, block>>>(d_data, N);
|
||||
cudaEventRecord(stop);
|
||||
cudaEventSynchronize(stop);
|
||||
float nodiv_ms; cudaEventElapsedTime(&nodiv_ms, start, stop);
|
||||
|
||||
printf("分歧版: %.2f ms\n", div_ms / 100);
|
||||
printf("无分支版: %.2f ms\n", nodiv_ms / 100);
|
||||
printf("加速比: %.2fx\n", div_ms / nodiv_ms);
|
||||
|
||||
cudaFree(d_data);
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,393 @@
|
||||
# Triton与TPU
|
||||
|
||||
*CUDA C功能强大但冗长。Triton让你用Python编写GPU核函数。TPU提供了GPU之外的选择，具有不同的权衡。本文涵盖Triton核函数编程、以Flash Attention为案例研究、TPU架构与JAX/Pallas，以及如何选择合适的工具。关于Vulkan和跨平台GPU计算，请参见文件07。*
|
||||
|
||||
- 上篇文件教授了CUDA C中的GPU编程。本文更上一层抽象阶梯：Triton以20%的工作量提供CUDA 80%的性能，且用Python。TPU和Vulkan为特定用例提供替代硬件目标。
|
||||
|
||||
## Triton：用Python编写GPU核函数
|
||||
|
||||
- **Triton**（OpenAI）是一种基于Python的GPU核函数编写语言。你不需要思考单个线程（CUDA），而是思考**块**级数据。Triton的编译器自动处理线程映射、内存合并、共享内存管理和许多优化。
|
||||
|
||||
- **为什么Triton重要**：CUDA C需要对线程束调度、共享内存存储体冲突、寄存器压力和合并模式有深入理解。Triton抽象了其中大部分内容，使GPU核函数开发对了解Python但不了解系统编程的ML研究人员可及。
|
||||
|
||||
### 你的第一个Triton核函数
|
||||
|
||||
```python
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
|
||||
@triton.jit
|
||||
def add_kernel(
|
||||
x_ptr, y_ptr, output_ptr,
|
||||
n_elements,
|
||||
BLOCK_SIZE: tl.constexpr, # 编译时常量
|
||||
):
|
||||
# 每个程序实例处理一个BLOCK_SIZE元素的块
|
||||
pid = tl.program_id(axis=0) # 我是哪个块？
|
||||
block_start = pid * BLOCK_SIZE
|
||||
|
||||
# 此块的偏移量
|
||||
offsets = block_start + tl.arange(0, BLOCK_SIZE)
|
||||
|
||||
# 掩码处理n_elements不是BLOCK_SIZE倍数的情况
|
||||
mask = offsets < n_elements
|
||||
|
||||
# 加载数据（带掩码：越界读取返回0）
|
||||
x = tl.load(x_ptr + offsets, mask=mask)
|
||||
y = tl.load(y_ptr + offsets, mask=mask)
|
||||
|
||||
# 计算
|
||||
output = x + y
|
||||
|
||||
# 存储结果
|
||||
tl.store(output_ptr + offsets, output, mask=mask)
|
||||
|
||||
|
||||
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
output = torch.empty_like(x)
|
||||
n_elements = output.numel()
|
||||
|
||||
# 启动：每个块一个程序
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# 使用
|
||||
x = torch.randn(1000000, device='cuda')
|
||||
y = torch.randn(1000000, device='cuda')
|
||||
z = add(x, y)
|
||||
```
|
||||
|
||||
- **与CUDA的关键区别**：
|
||||
- 无需显式线程管理。你思考**块**（程序），而非线程。
|
||||
- `tl.arange(0, BLOCK_SIZE)` 为整个块创建一个偏移向量。此向量上的所有操作都隐式向量化。
|
||||
- `mask` 处理边界条件（类似于AVX-512掩码寄存器，文件03）。无需标量清理循环。
|
||||
- `tl.load` 和 `tl.store` 自动处理合并访问。
|
||||
- `@triton.jit` 在首次调用时将函数编译为PTX（GPU汇编），然后缓存编译后的核函数。
|
||||
|
||||
### Triton Softmax核函数
|
||||
|
||||
- Softmax是一个很好的Triton示例，因为它需要对数据进行多次遍历（最大值、减去、指数、求和、除法），并且受益于在多次遍历之间将数据保留在SRAM（共享内存）中：
|
||||
|
||||
```python
|
||||
@triton.jit
|
||||
def softmax_kernel(
|
||||
output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
# 每个程序处理一行
|
||||
row_idx = tl.program_id(0)
|
||||
row_start = input_ptr + row_idx * input_row_stride
|
||||
|
||||
# 加载该行
|
||||
col_offsets = tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < n_cols
|
||||
row = tl.load(row_start + col_offsets, mask=mask, other=-float('inf'))
|
||||
|
||||
# Softmax：为数值稳定性取最大值，然后exp，然后归一化
|
||||
row_max = tl.max(row, axis=0)
|
||||
numerator = tl.exp(row - row_max)
|
||||
denominator = tl.sum(numerator, axis=0)
|
||||
softmax_output = numerator / denominator
|
||||
|
||||
# 存储结果
|
||||
output_start = output_ptr + row_idx * output_row_stride
|
||||
tl.store(output_start + col_offsets, softmax_output, mask=mask)
|
||||
```
|
||||
|
||||
- 在PyTorch中，`F.softmax(x, dim=-1)` 启动3个独立核函数（最大值、指数-求和、除法），每个都从全局内存读取和写入。Triton版本在一个核函数内完成所有操作，将数据保留在寄存器/SRAM中。这种**核函数融合**就是自定义Triton核函数可以比PyTorch内置操作快2-4倍的原因。
|
||||
|
||||
### Triton自动调优
|
||||
|
||||
- Triton支持**自动调优**：尝试多种配置并选择最快的：
|
||||
|
||||
```python
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}),
|
||||
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}),
|
||||
triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64}),
|
||||
],
|
||||
key=['M', 'N', 'K'], # 当这些变化时重新调优
|
||||
)
|
||||
@triton.jit
|
||||
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
|
||||
...
|
||||
```
|
||||
|
||||
- Triton在实际硬件上对每种配置进行基准测试并选择最快者。最优瓦片大小取决于GPU架构、矩阵维度和内存布局——自动调优无需手动实验即可找到它们。
|
||||
|
||||
### Triton vs CUDA：何时使用
|
||||
|
||||
| | Triton | CUDA C |
|
||||
|--|--------|--------|
|
||||
| 语言 | Python | C/C++ |
|
||||
| 抽象层级 | 块级 | 线程级 |
|
||||
| 开发速度 | 快（每核函数10-50行） | 慢（100-500行） |
|
||||
| 性能天花板 | 手工调优CUDA的约80-95% | 100%（完全硬件控制） |
|
||||
| 共享内存 | 自动 | 手动 |
|
||||
| 合并 | 自动 | 手动 |
|
||||
| 线程束级原语 | 有限 | 完整（shuffle、vote等） |
|
||||
| 硬件支持 | 仅NVIDIA（AMD实验性） | 仅NVIDIA |
|
||||
|
||||
- **使用Triton**对于：融合核函数、自定义注意力模式、激活函数、大多数ML研究核函数需求。
|
||||
- **使用CUDA C**对于：最高性能（最后5-20%）、线程束级原语、复杂数据相关并行性、当Triton无法表达你的模式。
|
||||
|
||||
## 案例研究：Flash Attention
|
||||
|
||||
- **Flash Attention**（Dao等人，2022）是近年来最具影响力的自定义核函数。它以 $O(n)$ 内存而非 $O(n^2)$ 计算注意力，使得更长的序列成为可能。
|
||||
|
||||
- **问题**：标准注意力计算 $\\text{softmax}(QK^T / \\sqrt{d}) \\cdot V$。$QK^T$ 矩阵是 $n \\times n$，其中 $n$ 是序列长度。对于 $n = 128K$，此矩阵为 $128K \\times 128K \\times 4$ 字节 = 64 GB。它无法放入GPU内存。
|
||||
|
||||
- **关键洞察**：你不需要具体化完整的 $n \\times n$ 矩阵。按**瓦片**计算注意力：加载一组 $Q$、一组 $K$，计算它们的部分注意力得分，累加，然后移动到下一个块。$n \\times n$ 矩阵从未完全具体化——每次只有一块存在于SRAM中。
|
||||
|
||||
- **在线softmax**：棘手的部分是softmax，它需要知道整个行上的最大值（为数值稳定性）。Flash Attention使用**在线softmax**技巧：维护一个运行中的最大值，当发现新的最大值时重新缩放先前计算的值。这允许softmax以增量方式逐块计算。
|
||||
|
||||
- 算法：
|
||||
|
||||
```
|
||||
对于每个Q行块：
|
||||
对于每个K列块：
|
||||
1. 将Q_block从HBM加载到SRAM
|
||||
2. 将K_block从HBM加载到SRAM
|
||||
3. 计算S_block = Q_block @ K_block.T（在SRAM中）
|
||||
4. 更新运行中最大值，重新缩放先前结果
|
||||
5. 计算exp(S_block - 运行中最大值)
|
||||
6. 更新运行中求和和输出累加器
|
||||
加载V_block并计算最终输出
|
||||
将输出块写回HBM
|
||||
```
|
||||
|
||||
- **为什么它快**：内循环完全在SRAM（共享内存）中操作。全局内存（HBM）仅用于加载Q、K、V块和写入最终输出。数据重用因子与SRAM大小成正比，而SRAM比HBM快约100倍。
|
||||
|
||||
- Flash Attention在Triton和CUDA C中都有实现。CUDA版本更快（效率高约10%），但Triton版本更具可读性和可修改性，这对研究新的注意力变体很重要。
|
||||
|
||||
## TPU架构
|
||||
|
||||
- **TPU**（张量处理单元）是Google的自定义ML加速器。它们采用与GPU截然不同的方法：
|
||||
|
||||
- **脉动阵列**：TPU的核心计算单元是**矩阵乘法单元（MXU）**，一个128×128或256×256的脉动阵列，通过让数据流经乘加单元网格来计算矩阵乘法。数据从边缘进入并通过阵列传播，每个单元执行一次乘加并将结果传递给下一个。
|
||||
|
||||
- 与GPU（调度数千个独立线程）不同，脉动阵列是单一的确定性数据流。没有线程调度、没有线程束分歧、没有分支预测。这种简朴性使MXU在矩阵乘法方面极其能效高效。
|
||||
|
||||
- **HBM**：TPU使用与GPU相同的高带宽内存。TPU v5e每芯片16 GB HBM2e；TPU v5p每芯片95 GB HBM2e。
|
||||
|
||||
- **ICI**（芯片间互连）：TPU Pod用自定义高速网络连接数百个TPU。JAX原生支持跨TPU Pod的数据并行性和模型并行性（第6章）。
|
||||
|
||||
- **BFloat16**：TPU是首个使用bfloat16的（第13章文件02）。BF16具有与float32相同的指数范围（防止训练期间溢出），尾数精度较低。这种权衡对ML是理想的，其中梯度值范围广但不需要23位精度。
|
||||
|
||||
### 编程TPU：JAX与Pallas
|
||||
|
||||
- TPU通过**JAX**和**XLA**编程。你编写Python/JAX代码，`jax.jit` 将其编译为XLA HLO，XLA将HLO编译为TPU特定的指令。无需CUDA，无需C++。
|
||||
|
||||
```python
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
@jax.jit
|
||||
def matmul(a, b):
|
||||
return jnp.dot(a, b)
|
||||
|
||||
# 这将根据设备在CPU、GPU或TPU上运行
|
||||
a = jnp.ones((1024, 1024))
|
||||
b = jnp.ones((1024, 1024))
|
||||
c = matmul(a, b)
|
||||
```
|
||||
|
||||
- **Pallas**是JAX的核函数编写API——JAX版的Triton。它让你编写低级核函数，XLA将其编译为GPU或TPU：
|
||||
|
||||
```python
|
||||
from jax.experimental import pallas as pl
|
||||
import jax.numpy as jnp
|
||||
|
||||
def add_kernel(x_ref, y_ref, o_ref):
|
||||
o_ref[...] = x_ref[...] + y_ref[...]
|
||||
|
||||
def add_pallas(x, y):
|
||||
return pl.pallas_call(
|
||||
add_kernel,
|
||||
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
|
||||
grid=(x.shape[0] // 128,),
|
||||
in_specs=[pl.BlockSpec((128,), lambda i: (i,)),
|
||||
pl.BlockSpec((128,), lambda i: (i,))],
|
||||
out_specs=pl.BlockSpec((128,), lambda i: (i,)),
|
||||
)(x, y)
|
||||
```
|
||||
|
||||
- Pallas比Triton更新且不太成熟，但它是为TPU编写自定义核函数的唯一方式（因为TPU不支持CUDA）。
|
||||
|
||||
### GPU vs TPU
|
||||
|
||||
| | GPU（NVIDIA） | TPU（Google） |
|
||||
|--|-------------|--------------|
|
||||
| 可用性 | 任何云、本地部署 | 仅Google Cloud |
|
||||
| 编程 | CUDA C、Triton、PyTorch | JAX/XLA、Pallas |
|
||||
| 灵活性 | 通用计算 | 针对矩阵密集型ML优化 |
|
||||
| 峰值矩阵乘法FLOPS | 非常高（张量核心） | 非常高（MXU） |
|
||||
| 非矩阵乘法操作 | 好 | 较慢（通过向量单元路由，而非MXU） |
|
||||
| 多芯片扩展 | NVLink（8个GPU）、InfiniBand | ICI（数千个TPU，更紧密集成） |
|
||||
| 成本效率 | 有竞争力 | 大规模训练通常更便宜 |
|
||||
| 生态系统 | 最大（PyTorch、TensorFlow、JAX） | 面向JAX |
|
||||
|
||||
- **使用GPU**对于：大多数ML工作负载、基于PyTorch的研究、推理服务、有大量非矩阵乘法计算的工作负载。
|
||||
- **使用TPU**对于：大规模JAX训练（数千芯片）、Google Cloud上的成本敏感训练、以矩阵乘法为主的工作负载。
|
||||
|
||||
## 选择合适的工具
|
||||
|
||||
| 工作负载 | 最佳工具 | 为什么 |
|
||||
|----------|---------|-------|
|
||||
| ML训练（PyTorch） | NVIDIA GPU + CUDA/Triton | 最大生态系统、最佳工具链 |
|
||||
| ML训练（JAX，大规模） | TPU或NVIDIA GPU | TPU在Google规模下成本低，GPU灵活 |
|
||||
| 自定义融合核函数 | Triton（Python）或CUDA C | Triton开发速度快，CUDA峰值性能高 |
|
||||
| JAX自定义核函数 | Pallas | TPU唯一选项，也可在GPU上工作 |
|
||||
| 跨平台推理 | Vulkan（文件07）或ONNX Runtime | 运行在任何GPU供应商上 |
|
||||
| 移动/边缘推理 | Metal（Apple）、Vulkan（Android）、NNAPI | 平台特定的加速器 |
|
||||
| 浏览器推理 | WebGPU（文件07） | 浏览器中唯一选项 |
|
||||
| 仅CPU推理 | ONNX Runtime + AVX/NEON | 无需GPU，使用SIMD（文件02-03） |
|
||||
| 新型硬件 | 供应商专用SDK | 每个加速器有自己的工具链 |
|
||||
|
||||
## 编程任务（使用带GPU运行时的CoLab）
|
||||
|
||||
1. 编写并运行向量加法的Triton核函数。将其性能与PyTorch内置加法比较。
|
||||
```python
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
import time
|
||||
|
||||
@triton.jit
|
||||
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
offs = pid * BLOCK + tl.arange(0, BLOCK)
|
||||
mask = offs < n
|
||||
x = tl.load(x_ptr + offs, mask=mask)
|
||||
y = tl.load(y_ptr + offs, mask=mask)
|
||||
tl.store(out_ptr + offs, x + y, mask=mask)
|
||||
|
||||
n = 10_000_000
|
||||
x = torch.randn(n, device='cuda')
|
||||
y = torch.randn(n, device='cuda')
|
||||
|
||||
# Triton
|
||||
out_triton = torch.empty_like(x)
|
||||
grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
|
||||
add_kernel[grid](x, y, out_triton, n, BLOCK=1024)
|
||||
|
||||
# PyTorch
|
||||
out_torch = x + y
|
||||
|
||||
# 验证正确性
|
||||
assert torch.allclose(out_triton, out_torch, atol=1e-5)
|
||||
|
||||
# 基准测试
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(1000):
|
||||
add_kernel[grid](x, y, out_triton, n, BLOCK=1024)
|
||||
torch.cuda.synchronize()
|
||||
triton_time = (time.time() - start) / 1000
|
||||
|
||||
start = time.time()
|
||||
for _ in range(1000):
|
||||
out_torch = x + y
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start) / 1000
|
||||
|
||||
print(f"Triton: {triton_time*1000:.3f} ms")
|
||||
print(f"PyTorch: {torch_time*1000:.3f} ms")
|
||||
print(f"比率: {torch_time/triton_time:.2f}x")
|
||||
```
|
||||
|
||||
2. 编写一个Triton融合核函数，在单次遍历中执行乘法+加法+ReLU。与三个独立的PyTorch操作比较。
|
||||
```python
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
import time
|
||||
|
||||
@triton.jit
|
||||
def fused_mul_add_relu_kernel(x_ptr, w_ptr, b_ptr, out_ptr, n, BLOCK: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
offs = pid * BLOCK + tl.arange(0, BLOCK)
|
||||
mask = offs < n
|
||||
x = tl.load(x_ptr + offs, mask=mask)
|
||||
w = tl.load(w_ptr + offs, mask=mask)
|
||||
b = tl.load(b_ptr + offs, mask=mask)
|
||||
result = tl.maximum(x * w + b, 0.0) # 融合：乘法 + 加法 + relu
|
||||
tl.store(out_ptr + offs, result, mask=mask)
|
||||
|
||||
n = 10_000_000
|
||||
x = torch.randn(n, device='cuda')
|
||||
w = torch.randn(n, device='cuda')
|
||||
b = torch.randn(n, device='cuda')
|
||||
|
||||
# 融合（Triton）
|
||||
out_fused = torch.empty_like(x)
|
||||
grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
|
||||
fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)
|
||||
|
||||
# 未融合（PyTorch）
|
||||
out_unfused = torch.relu(x * w + b)
|
||||
|
||||
assert torch.allclose(out_fused, out_unfused, atol=1e-5)
|
||||
|
||||
# 基准测试
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(1000):
|
||||
fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)
|
||||
torch.cuda.synchronize()
|
||||
fused_time = (time.time() - start) / 1000
|
||||
|
||||
start = time.time()
|
||||
for _ in range(1000):
|
||||
out_unfused = torch.relu(x * w + b)
|
||||
torch.cuda.synchronize()
|
||||
unfused_time = (time.time() - start) / 1000
|
||||
|
||||
print(f"融合（Triton）: {fused_time*1000:.3f} ms")
|
||||
print(f"未融合（PyTorch）: {unfused_time*1000:.3f} ms")
|
||||
print(f"加速比: {unfused_time/fused_time:.2f}x")
|
||||
```
|
||||
|
||||
3. 测量JAX的XLA编译器如何自动融合操作。比较带和不带jit的操作链。
|
||||
```python
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import time
|
||||
|
||||
def chain_ops(x):
|
||||
x = x * 2.0
|
||||
x = x + 1.0
|
||||
x = jnp.maximum(x, 0.0) # ReLU
|
||||
x = x / jnp.sum(x)
|
||||
return x
|
||||
|
||||
chain_jit = jax.jit(chain_ops)
|
||||
x = jax.random.normal(jax.random.PRNGKey(0), (10000, 1000))
|
||||
|
||||
# 预热
|
||||
_ = chain_jit(x)
|
||||
jax.block_until_ready(_)
|
||||
|
||||
# 即时模式（每个操作是独立核函数启动）
|
||||
start = time.time()
|
||||
for _ in range(100):
|
||||
y = chain_ops(x)
|
||||
jax.block_until_ready(y)
|
||||
eager_time = (time.time() - start) / 100
|
||||
|
||||
# JIT（XLA融合操作）
|
||||
start = time.time()
|
||||
for _ in range(100):
|
||||
y = chain_jit(x)
|
||||
jax.block_until_ready(y)
|
||||
jit_time = (time.time() - start) / 100
|
||||
|
||||
print(f"即时: {eager_time*1000:.2f} ms")
|
||||
print(f"JIT: {jit_time*1000:.2f} ms")
|
||||
print(f"加速比: {eager_time/jit_time:.1f}x（XLA将4个操作融合为1个核函数）")
|
||||
```
|
||||
@@ -0,0 +1,428 @@
|
||||
# RISC-V与嵌入式系统
|
||||
|
||||
*RISC-V是正在重塑芯片行业的开源指令集架构。本文涵盖RISC-V哲学、V向量扩展、嵌入式ML推理、微控制器上的TinyML、AI加速器中的RISC-V以及边缘部署约束*
|
||||
|
||||
- 我们之前讨论的每一种芯片架构（x86、ARM）都需要**许可**。Intel和AMD为x86付费。Apple、Qualcomm以及每一家智能手机厂商每年向ARM支付数十亿美元。**RISC-V**则不同：它是一个开放标准。任何人都可以设计、制造和销售RISC-V芯片，无需向任何人支付版税。这正在改变芯片设计的经济性，特别是对于AI。
|
||||
|
||||
## RISC-V哲学
|
||||
|
||||
- RISC-V（发音为"risk five"）于2010年在加州大学伯克利分校创建，作为一个简洁、现代的RISC指令集。关键原则：
|
||||
|
||||
- **开放标准**：ISA规范免费提供。你可以在没有许可费、NDA或法律协议的情况下构建RISC-V CPU。这就像Linux之于操作系统——任何人都可以使用、修改和在此基础上构建。
|
||||
|
||||
- **模块化设计**：基础ISA（RV32I或RV64I）是最小的——仅47条指令。其他一切都是可选的**扩展**：M（乘法/除法）、A（原子操作）、F/D（浮点）、C（压缩指令）、V（向量处理）。你只选择需要的，保持芯片小巧高效。
|
||||
|
||||
- **无遗留包袱**：x86背负着45年的向后兼容性。ARM背负着35年。RISC-V从零开始，融入了从两者中吸取的经验教训。没有仅为与1980年代软件兼容而存在的晦涩指令。
|
||||
|
||||
- **谁在使用RISC-V**：SiFive（通用核心）、阿里巴巴（玄铁服务器核心）、西部数据（存储控制器，已出货数十亿）、乐鑫（ESP32-C3，流行IoT芯片），以及数十家使用RISC-V作为管理其自定义计算单元的控制处理器的AI加速器初创公司。
|
||||
|
||||
## RISC-V基础架构
|
||||
|
||||
- 基础整数ISA（RV64I用于64位）具有：
|
||||
- **32个通用寄存器**（x0-x31，每个64位）。x0硬连接为零（用于在没有特殊指令的情况下实现常见模式）。
|
||||
- **固定32位指令宽度**（C扩展为代码密度添加了16位压缩指令）。
|
||||
- **加载-存储架构**：与ARM一样，算术仅操作寄存器。内存访问通过显式加载/存储指令进行。
|
||||
|
||||
```
|
||||
# RISC-V汇编（感受风格——你将使用C/C++）
|
||||
add x3, x1, x2 # x3 = x1 + x2
|
||||
lw x4, 0(x5) # 从x5中的地址加载字
|
||||
sw x4, 8(x5) # 存储字到地址 x5 + 8
|
||||
beq x1, x2, label # 如果x1 == x2则分支
|
||||
```
|
||||
|
||||
- ISA的简洁性使RISC-V核心小巧且能效高。最小的RV32I核心可以用约10,000个门实现（ARM Cortex-M0约为12,000）。这对于每一毫瓦和每一平方毫米硅片都至关重要的嵌入式系统很重要。
|
||||
|
||||
## V扩展：RISC-V向量处理
|
||||
|
||||
- **V扩展**（RVV）为RISC-V添加了可伸缩向量处理，类似于ARM SVE。向量寄存器具有可配置长度（VLEN），由硬件指定（128到65,536位）。代码编写为**向量长度无关**：无需重新编译可在任何VLEN上工作。
|
||||
|
||||
```c
|
||||
#include <riscv_vector.h>
|
||||
|
||||
// 使用RVV内联函数进行向量加法
|
||||
void vadd_rvv(const float* a, const float* b, float* c, int n) {
|
||||
while (n > 0) {
|
||||
// vsetvl：设置向量长度——处理 min(n, 硬件最大值) 个元素
|
||||
size_t vl = __riscv_vsetvl_e32m1(n);
|
||||
|
||||
// 加载vl个元素
|
||||
vfloat32m1_t va = __riscv_vle32_v_f32m1(a, vl);
|
||||
vfloat32m1_t vb = __riscv_vle32_v_f32m1(b, vl);
|
||||
|
||||
// 相加
|
||||
vfloat32m1_t vc = __riscv_vfadd_vv_f32m1(va, vb, vl);
|
||||
|
||||
// 存储
|
||||
__riscv_vse32_v_f32m1(c, vc, vl);
|
||||
|
||||
// 前进指针
|
||||
a += vl; b += vl; c += vl; n -= vl;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- **`vsetvl`** 是关键指令。它告诉硬件"我想处理这么多元素"，硬件回应"我可以处理这么多"（受VLEN限制）。循环自动适应任何向量宽度，无需标量清理（最后一次迭代只处理较少的元素）。
|
||||
|
||||
- **LMUL**（长度乘数）：RVV可以将多个向量寄存器分组在一起（m1、m2、m4、m8），以每条指令处理更多元素，代价是可用的寄存器更少。`m1` 每个向量操作数使用一个寄存器；`m8` 使用八个，处理8倍元素，但只留下4个寄存器组可用。
|
||||
|
||||
- 与x86 AVX（固定256/512位）和ARM NEON（固定128位）相比，RVV的可伸缩性对于多样化硬件是一个重要优势：相同代码在小型嵌入式核心（VLEN=128）和高性能服务器核心（VLEN=1024+）上运行。
|
||||
|
||||
## 嵌入式ML：TinyML
|
||||
|
||||
- **TinyML**是微控制器上的机器学习——具有千字节RAM、兆赫级CPU和毫瓦功率预算的设备。想想：检测关键词的传感器（"Hey Siri"）、分类手势的加速度计、或计数人数的小型摄像头，所有这些都在一个售价0.50美元、无需互联网连接的芯片上运行。
|
||||
|
||||
- 约束条件极其严苛：
|
||||
|
||||
| 资源 | 服务器GPU | 智能手机 | 微控制器 |
|
||||
|--------|----------|------------|-----------------|
|
||||
| RAM | 80 GB | 6 GB | 256 KB |
|
||||
| 存储 | TB | 128 GB | 1 MB |
|
||||
| 计算能力 | 1000 TFLOPS | 10 TFLOPS | 0.001 TFLOPS |
|
||||
| 功耗 | 700 W | 5 W | 0.001 W |
|
||||
| 成本 | $30,000 | $500 | $1 |
|
||||
|
||||
- 适合服务器GPU的模型（$O(10^{10})$ 参数）无法放入微控制器。TinyML模型有 $O(10^4)$–$O(10^6)$ 参数，并使用INT8甚至INT4量化。
|
||||
|
||||
### TensorFlow Lite Micro（TFLM）
|
||||
|
||||
- **TFLM**是Google的微控制器推理框架。它运行量化的TensorFlow Lite模型，无需动态内存分配、无需操作系统，二进制占用约20 KB。
|
||||
|
||||
```cpp
|
||||
// 微控制器上的TinyML推理（简化版）
|
||||
#include "tensorflow/lite/micro/micro_interpreter.h"
|
||||
#include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
|
||||
|
||||
// 模型编译为C数组（const unsigned char model_data[]）
|
||||
const tflite::Model* model = tflite::GetModel(model_data);
|
||||
|
||||
// 分配固定内存缓冲区（无malloc！）
|
||||
constexpr int kArenaSize = 10 * 1024; // 10 KB
|
||||
uint8_t tensor_arena[kArenaSize];
|
||||
|
||||
// 设置解释器
|
||||
tflite::MicroInterpreter interpreter(model, resolver, tensor_arena, kArenaSize);
|
||||
interpreter.AllocateTensors();
|
||||
|
||||
// 设置输入
|
||||
float* input = interpreter.input(0)->data.f;
|
||||
input[0] = sensor_reading;
|
||||
|
||||
// 运行推理
|
||||
interpreter.Invoke();
|
||||
|
||||
// 读取输出
|
||||
float* output = interpreter.output(0)->data.f;
|
||||
if (output[0] > 0.8f) {
|
||||
trigger_alert();
|
||||
}
|
||||
```
|
||||
|
||||
- **此代码中的关键约束**：
|
||||
- `tensor_arena` 是静态分配的——没有 `malloc`，没有堆。嵌入式系统通常没有动态内存分配器。
|
||||
- 模型是一个 `const` 字节数组，存储在闪存（ROM）中，而非从文件系统加载。
|
||||
- 整个框架+模型+运行时适合几十KB。
|
||||
|
||||
### 边缘模型优化
|
||||
|
||||
- 让模型在微控制器上运行需要激进优化：
|
||||
|
||||
- **量化**（第18章）：将float32权重转换为INT8（小4倍，在纯整数硬件上快2-4倍）。训练后量化简单；量化感知训练保留更多精度。
|
||||
|
||||
- **剪枝**：移除接近零的权重。结构化剪枝（移除整个通道/头）比非结构化剪枝（随机零）对硬件更友好，因为它减少实际计算，而不仅是存储。
|
||||
|
||||
- **知识蒸馏**（第6章）：训练一个小型"学生"模型来模仿大型"教师"模型。学生模型比从头训练获得更高精度，因为它从教师模型的软预测中学习。
|
||||
|
||||
- **神经架构搜索（NAS）**：自动搜索适合硬件预算（延迟、内存、功耗）的高效架构。**MicroNets**和**MCUNet**为特定微控制器寻找优化架构。
|
||||
|
||||
- **算子融合**：将卷积+批归一化+ReLU组合成单个融合操作，消除中间内存写入（与GPU核函数融合同一原则，但在只有256 KB RAM时更加关键）。
|
||||
|
||||
## RISC-V在AI加速器中的应用
|
||||
|
||||
- 许多AI加速器初创公司使用RISC-V并非直接运行ML模型，而是作为管理自定义计算单元的**控制处理器**：
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────┐
|
||||
│ AI加速器 │
|
||||
│ │
|
||||
│ ┌──────────┐ ┌──────────────────┐ │
|
||||
│ │ RISC-V │───→│ 自定义矩阵 │ │
|
||||
│ │ 控制 │ │ 乘法单元 │ │
|
||||
│ │ 核心 │ │ （脉动阵列、 │ │
|
||||
│ │ │ │ 自定义数据流） │ │
|
||||
│ └──────────┘ └──────────────────┘ │
|
||||
│ │ │ │
|
||||
│ ▼ ▼ │
|
||||
│ ┌──────────┐ ┌──────────────────┐ │
|
||||
│ │ 内存 │ │ 片上SRAM │ │
|
||||
│ │ 控制 │ │ （激活缓冲） │ │
|
||||
│ │ │ │ │ │
|
||||
│ └──────────┘ └──────────────────┘ │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
- RISC-V核心处理：从外部内存加载模型权重、调度层执行、管理计算单元之间的数据流以及与主机通信（通过PCIe、USB或SPI）。繁重计算（矩阵乘法、卷积）由自定义硬件完成，而非RISC-V核心。
|
||||
|
||||
- **为什么用RISC-V做控制**：无需许可费用（对初创公司至关重要）、可定制（添加领域特定指令）、小占用空间（控制核心不需要x86的复杂性），以及开放生态系统支持快速原型开发。
|
||||
|
||||
- **示例**：Esperanto Technologies（1000+个RISC-V核心用于ML）、Tenstorrent（RISC-V控制+自定义tensix核心）、SiFive（带向量扩展的RISC-V核心用于边缘ML）。
|
||||
|
||||
## 边缘部署约束
|
||||
|
||||
- 在边缘部署ML（设备端，而非云端）引入了云端部署不需要的约束：
|
||||
|
||||
- **功耗**：电池供电的设备总功耗预算可能为100 mW。运行消耗50 mW的模型只给系统其余部分（传感器、无线电、显示器）留下50 mW。功耗感知推理调度计算以避免热降频并延长电池寿命。
|
||||
|
||||
- **延迟**：边缘推理通常必须是实时的。唤醒词检测器（"Hey Siri"）必须在约200 ms内响应。自动驾驶感知系统（第11章）必须在约30 ms内处理帧。到云端的网络往返（50-200 ms）对这些用例来说太慢了。
|
||||
|
||||
- **隐私**：在设备上处理数据意味着敏感数据（医学图像、语音记录、个人照片）永远不会离开设备。这在某些司法管辖区是法律要求（GDPR），在所有地方都是用户信任的要求。
|
||||
|
||||
- **连接性**：边缘设备可能间歇性或完全没有互联网连接。在火星车（第11章）、潜艇或农村农田传感器上运行的模型必须完全离线工作。
|
||||
|
||||
- **规模成本**：将ML部署到十亿部智能手机每台成本为$0（硬件已经存在）。将ML部署到十亿个IoT传感器意味着每个传感器的ML硬件预算只有几分钱。RISC-V的零许可成本在这个规模下意义重大。
|
||||
|
||||
## 编程任务（用g++或riscv64-gcc交叉编译器编译）
|
||||
|
||||
1. 编写一个C程序，模拟TinyML推理流水线：静态分配模型缓冲区，运行模拟前向传播，并测量资源使用。这教授嵌入式约束（无malloc、固定内存缓冲区）。
|
||||
```cpp
|
||||
// task1_tinyml_sim.cpp
|
||||
// 编译：g++ -O2 -o task1 task1_tinyml_sim.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
// 模拟微控制器：固定内存缓冲区，无动态分配
|
||||
static constexpr int ARENA_SIZE = 32 * 1024; // 32 KB总RAM预算
|
||||
static uint8_t arena[ARENA_SIZE];
|
||||
|
||||
// 简单的2层MLP：784 -> 64 -> 10（类似MNIST，INT8权重）
|
||||
struct TinyModel {
|
||||
int8_t w1[784 * 64]; // 层1权重：50,176字节
|
||||
int8_t b1[64]; // 层1偏置
|
||||
int8_t w2[64 * 10]; // 层2权重：640字节
|
||||
int8_t b2[10]; // 层2偏置
|
||||
// 总计：约51 KB → 必须放在闪存（ROM），而非RAM
|
||||
};
|
||||
|
||||
// 检查模型是否适合闪存
|
||||
void check_model_fit(int flash_kb) {
|
||||
int model_bytes = sizeof(TinyModel);
|
||||
std::cout << "模型大小: " << model_bytes << " 字节（"
|
||||
<< model_bytes / 1024 << " KB）\n";
|
||||
std::cout << "闪存: " << flash_kb << " KB → "
|
||||
<< (model_bytes <= flash_kb * 1024 ? "适合" : "太大") << "\n";
|
||||
}
|
||||
|
||||
// 使用固定缓冲区进行激活的模拟推理
|
||||
void mock_inference(const int8_t* input, int8_t* output) {
|
||||
// 激活值放在缓冲区（RAM）中，而非动态分配
|
||||
int8_t* act1 = (int8_t*)arena; // 层1输出64字节
|
||||
int8_t* act2 = (int8_t*)(arena + 64); // 层2输出10字节
|
||||
|
||||
// 层1：简化版矩阵乘法（不是真正的量化矩阵乘法，仅结构演示）
|
||||
for (int j = 0; j < 64; j++) {
|
||||
int32_t sum = 0; // 用int32累加避免溢出
|
||||
for (int i = 0; i < 784; i++) {
|
||||
sum += (int32_t)input[i] * 1; // 模拟：权重=1
|
||||
}
|
||||
act1[j] = (int8_t)std::max(-128, std::min(127, sum / 784)); // 量化回
|
||||
act1[j] = act1[j] > 0 ? act1[j] : 0; // ReLU
|
||||
}
|
||||
|
||||
// 层2
|
||||
for (int j = 0; j < 10; j++) {
|
||||
int32_t sum = 0;
|
||||
for (int i = 0; i < 64; i++) {
|
||||
sum += (int32_t)act1[i] * 1;
|
||||
}
|
||||
act2[j] = (int8_t)std::max(-128, std::min(127, sum / 64));
|
||||
}
|
||||
|
||||
std::memcpy(output, act2, 10);
|
||||
}
|
||||
|
||||
int main() {
|
||||
std::cout << "=== TinyML资源预算 ===\n";
|
||||
std::cout << "缓冲区（RAM）: " << ARENA_SIZE << " 字节（"
|
||||
<< ARENA_SIZE / 1024 << " KB）\n";
|
||||
check_model_fit(256); // 典型MCU闪存
|
||||
|
||||
// 激活内存使用
|
||||
int activation_bytes = 64 + 10; // 层1 + 层2输出
|
||||
std::cout << "激活内存: " << activation_bytes
|
||||
<< " 字节 / " << ARENA_SIZE << " 可用\n\n";
|
||||
|
||||
// 基准测试推理
|
||||
int8_t input[784];
|
||||
int8_t output[10];
|
||||
std::memset(input, 1, 784);
|
||||
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int i = 0; i < 10000; i++) {
|
||||
mock_inference(input, output);
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double us = std::chrono::duration<double, std::micro>(end - start).count() / 10000;
|
||||
|
||||
std::cout << "推理延迟: " << us << " us\n";
|
||||
std::cout << "在160 MHz MCU（约6.25 ns/周期）下：约"
|
||||
<< (int)(us * 160) << " 周期\n";
|
||||
|
||||
std::cout << "输出logits: ";
|
||||
for (int i = 0; i < 10; i++) std::cout << (int)output[i] << " ";
|
||||
std::cout << "\n";
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
2. 编写一个C++程序，将float32权重量化为INT8，并测量压缩比和量化误差。
|
||||
```cpp
|
||||
// task2_quantise.cpp
|
||||
// 编译：g++ -O3 -o task2 task2_quantise.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <cmath>
|
||||
#include <algorithm>
|
||||
#include <numeric>
|
||||
|
||||
// 对称量化：将浮点范围 [-max, +max] 映射到 [-127, +127]
|
||||
void quantise_symmetric(const float* input, int8_t* output, int n, float& scale) {
|
||||
float max_val = 0.0f;
|
||||
for (int i = 0; i < n; i++) {
|
||||
max_val = std::max(max_val, std::abs(input[i]));
|
||||
}
|
||||
scale = max_val / 127.0f;
|
||||
for (int i = 0; i < n; i++) {
|
||||
float scaled = input[i] / scale;
|
||||
output[i] = (int8_t)std::max(-127.0f, std::min(127.0f, std::round(scaled)));
|
||||
}
|
||||
}
|
||||
|
||||
// 反量化：INT8转回float
|
||||
void dequantise(const int8_t* input, float* output, int n, float scale) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
output[i] = (float)input[i] * scale;
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int N = 100000;
|
||||
|
||||
// 模拟随机权重（大致正态分布）
|
||||
std::vector<float> weights(N);
|
||||
for (int i = 0; i < N; i++) {
|
||||
// 简单的伪随机正态值
|
||||
float u1 = (float)(i * 7 % 997 + 1) / 998.0f;
|
||||
float u2 = (float)(i * 13 % 991 + 1) / 992.0f;
|
||||
weights[i] = std::sqrt(-2.0f * std::log(u1)) * std::cos(6.2832f * u2) * 0.1f;
|
||||
}
|
||||
|
||||
// 量化
|
||||
std::vector<int8_t> quantised(N);
|
||||
float scale;
|
||||
quantise_symmetric(weights.data(), quantised.data(), N, scale);
|
||||
|
||||
// 反量化并测量误差
|
||||
std::vector<float> reconstructed(N);
|
||||
dequantise(quantised.data(), reconstructed.data(), N, scale);
|
||||
|
||||
float max_error = 0.0f, total_error = 0.0f;
|
||||
for (int i = 0; i < N; i++) {
|
||||
float err = std::abs(weights[i] - reconstructed[i]);
|
||||
max_error = std::max(max_error, err);
|
||||
total_error += err;
|
||||
}
|
||||
|
||||
std::cout << "=== 量化结果 ===\n";
|
||||
std::cout << "原始: " << N * 4 << " 字节（float32）\n";
|
||||
std::cout << "量化: " << N * 1 << " 字节（int8）+ 4 字节（缩放因子）\n";
|
||||
std::cout << "压缩比: " << 4.0f << "x\n";
|
||||
std::cout << "缩放因子: " << scale << "\n";
|
||||
std::cout << "平均绝对误差: " << total_error / N << "\n";
|
||||
std::cout << "最大绝对误差: " << max_error << "\n";
|
||||
std::cout << "最大绝对误差/缩放因子: " << max_error / scale
|
||||
<< "（应 <= 0.5 量化级别）\n";
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
3. 编写一个C++程序，执行INT8矩阵乘法（INT32累加）——这是在嵌入式ML加速器上运行的实际计算。
|
||||
```cpp
|
||||
// task3_int8_matmul.cpp
|
||||
// 编译：g++ -O3 -o task3 task3_int8_matmul.cpp
|
||||
|
||||
#include <iostream>
|
||||
#include <chrono>
|
||||
#include <vector>
|
||||
#include <cstdint>
|
||||
|
||||
// INT8矩阵乘法（INT32累加）——张量核心和MCU加速器的实际工作方式
|
||||
void matmul_int8(const int8_t* A, const int8_t* B, int32_t* C,
|
||||
int M, int N, int K) {
|
||||
for (int i = 0; i < M; i++) {
|
||||
for (int j = 0; j < N; j++) {
|
||||
int32_t sum = 0;
|
||||
for (int k = 0; k < K; k++) {
|
||||
sum += (int32_t)A[i * K + k] * (int32_t)B[k * N + j];
|
||||
}
|
||||
C[i * N + j] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 用于比较的Float32矩阵乘法
|
||||
void matmul_f32(const float* A, const float* B, float* C,
|
||||
int M, int N, int K) {
|
||||
for (int i = 0; i < M; i++) {
|
||||
for (int j = 0; j < N; j++) {
|
||||
float sum = 0.0f;
|
||||
for (int k = 0; k < K; k++) {
|
||||
sum += A[i * K + k] * B[k * N + j];
|
||||
}
|
||||
C[i * N + j] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
const int M = 128, N = 128, K = 128;
|
||||
|
||||
std::vector<int8_t> A_i8(M * K, 1), B_i8(K * N, 1);
|
||||
std::vector<int32_t> C_i32(M * N);
|
||||
|
||||
std::vector<float> A_f32(M * K, 1.0f), B_f32(K * N, 1.0f);
|
||||
std::vector<float> C_f32(M * N);
|
||||
|
||||
// 基准测试INT8
|
||||
auto start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) {
|
||||
matmul_int8(A_i8.data(), B_i8.data(), C_i32.data(), M, N, K);
|
||||
}
|
||||
auto end = std::chrono::high_resolution_clock::now();
|
||||
double i8_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
|
||||
// 基准测试FP32
|
||||
start = std::chrono::high_resolution_clock::now();
|
||||
for (int t = 0; t < 100; t++) {
|
||||
matmul_f32(A_f32.data(), B_f32.data(), C_f32.data(), M, N, K);
|
||||
}
|
||||
end = std::chrono::high_resolution_clock::now();
|
||||
double f32_ms = std::chrono::duration<double, std::milli>(end - start).count() / 100;
|
||||
|
||||
double gops_i8 = 2.0 * M * N * K / i8_ms / 1e6;
|
||||
double gflops_f32 = 2.0 * M * N * K / f32_ms / 1e6;
|
||||
|
||||
std::cout << "INT8矩阵乘法: " << i8_ms << " ms（" << gops_i8 << " GOPS）\n";
|
||||
std::cout << "FP32矩阵乘法: " << f32_ms << " ms（" << gflops_f32 << " GFLOPS）\n";
|
||||
std::cout << "INT8加速比: " << f32_ms / i8_ms << "x\n";
|
||||
std::cout << "内存: INT8 = " << M*K + K*N << " 字节 vs FP32 = "
|
||||
<< (M*K + K*N) * 4 << " 字节（小4倍）\n";
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,668 @@
|
||||
# Vulkan Compute 与跨平台 GPU
|
||||
|
||||
*Vulkan 是唯一能在所有主要平台上运行的 GPU 计算 API：NVIDIA、AMD、Intel、Apple（通过 MoltenVK）、Android，甚至浏览器（通过 WebGPU）。本文涵盖 Vulkan 架构、计算管线、使用 GLSL 编写计算着色器、GPU 计算程序的完整 C++ 设置、共享内存与同步、用于浏览器的 WebGPU，以及实际的机器学习推理示例。*
|
||||
|
||||
- CUDA 在 NVIDIA 硬件上主导着 ML 训练。但并非每个部署目标都有 NVIDIA GPU。移动应用运行在 Qualcomm Adreno 或 ARM Mali GPU 上。Web 应用运行在浏览器中。游戏引擎需要同时支持 AMD、Intel 和 NVIDIA。对于所有这些场景，**Vulkan** 就是答案。
|
||||
|
||||
- Vulkan 很冗长——一个"hello world"计算程序大约有 300 行 C++ 代码。但这种冗长是 **显式控制** 的代价：你需要自己管理每一个 GPU 资源（内存、管线、命令缓冲区）。这种控制带来了最大性能和可移植性，代价是开发速度。
|
||||
|
||||
## Vulkan 架构概述
|
||||
|
||||
- Vulkan 是由 Khronos Group（OpenGL 背后的同一组织）创建的低级 GPU API。与 CUDA（它隐藏了 GPU 资源管理）不同，Vulkan 要求你显式地管理：
|
||||
|
||||
- **实例与设备**：创建 Vulkan 实例，枚举可用 GPU，并选择一个。
|
||||
- **内存**：显式分配 GPU 内存，指定内存类型（设备本地内存用于速度，主机可见内存用于 CPU 访问）。
|
||||
- **缓冲区**：创建引用已分配内存的缓冲区对象。
|
||||
- **描述符集**：将缓冲区绑定到着色器输入（类似于计算着色器的函数参数）。
|
||||
- **计算管线**：编译着色器并创建管线对象。
|
||||
- **命令缓冲区**：记录一系列 GPU 命令（绑定管线、绑定描述符、调度计算）。
|
||||
- **队列提交**：将命令缓冲区提交给 GPU 执行。
|
||||
- **同步**：使用栅栏和屏障确保正确的执行顺序。
|
||||
|
||||
- 这与 CUDA 的 `cudaMalloc` + 内核启动模型截然不同。在 CUDA 中，驱动程序在幕后处理大部分工作。在 Vulkan 中，你需要自己做这一切。
|
||||
|
||||
### 为什么如此冗长？
|
||||
|
||||
- Vulkan 的显式性存在有两方面原因：
|
||||
|
||||
1. **驱动简化**：OpenGL 驱动极其复杂（它们必须猜测应用程序的意图并进行相应优化）。Vulkan 将该责任转移给应用程序，使驱动更精简、更可预测，并且更容易在各厂商间正确实现。
|
||||
|
||||
2. **性能**：对内存布局、同步和命令批处理的显式控制使应用程序能够做出最优决策。在 CUDA 中，驱动可能会插入不必要的同步。在 Vulkan 中，你只在需要时才进行同步。
|
||||
|
||||
## GLSL 中的计算着色器
|
||||
|
||||
- **计算着色器** 是在 GPU 上运行的程序，类似于 CUDA 内核。它使用 **GLSL**（OpenGL 着色语言）编写，并编译为 **SPIR-V** 字节码（一种可移植的二进制格式）。
|
||||
|
||||
### 向量加法
|
||||
|
||||
```glsl
|
||||
// add.comp — 编译命令: glslangValidator -V add.comp -o add.spv
|
||||
#version 450
|
||||
|
||||
// 工作组大小：每个工作组有 256 个调用（= CUDA 中每块的线程数）
|
||||
layout(local_size_x = 256) in;
|
||||
|
||||
// 缓冲区绑定（类似于内核参数）
|
||||
layout(set = 0, binding = 0) buffer InputA { float a[]; };
|
||||
layout(set = 0, binding = 1) buffer InputB { float b[]; };
|
||||
layout(set = 0, binding = 2) buffer Output { float c[]; };
|
||||
|
||||
// 推送常量：小的统一数据（类似于内核参数）
|
||||
layout(push_constant) uniform PushConstants {
|
||||
uint n; // 元素数量
|
||||
};
|
||||
|
||||
void main() {
|
||||
uint idx = gl_GlobalInvocationID.x; // 全局线程索引
|
||||
if (idx < n) {
|
||||
c[idx] = a[idx] + b[idx];
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- **与 CUDA 概念的映射**：
|
||||
|
||||
| Vulkan | CUDA | 含义 |
|
||||
|--------|------|------|
|
||||
| 工作组 (Workgroup) | 块 (Block) | 可以共享内存的线程组 |
|
||||
| 调用 (Invocation) | 线程 (Thread) | 单个执行单元 |
|
||||
| `gl_GlobalInvocationID` | `blockIdx * blockDim + threadIdx` | 全局线程索引 |
|
||||
| `gl_LocalInvocationID` | `threadIdx` | 工作组内的线程索引 |
|
||||
| `gl_WorkGroupID` | `blockIdx` | 工作组索引 |
|
||||
| `local_size_x` | `blockDim.x` | 每工作组的线程数 |
|
||||
| 存储缓冲区 | 全局内存 | 可读写的 GPU 内存 |
|
||||
| 共享内存 (`shared`) | `__shared__` | 每工作组的高速内存 |
|
||||
| 推送常量 | 内核参数 | 小的统一数据 |
|
||||
|
||||
### 使用共享内存的 ReLU
|
||||
|
||||
```glsl
|
||||
// relu_shared.comp
|
||||
#version 450
|
||||
|
||||
layout(local_size_x = 256) in;
|
||||
|
||||
layout(set = 0, binding = 0) buffer Input { float input_data[]; };
|
||||
layout(set = 0, binding = 1) buffer Output { float output_data[]; };
|
||||
|
||||
layout(push_constant) uniform PushConstants { uint n; };
|
||||
|
||||
// 共享内存（等同于 CUDA 的 __shared__）
|
||||
shared float tile[256];
|
||||
|
||||
void main() {
|
||||
uint gid = gl_GlobalInvocationID.x;
|
||||
uint lid = gl_LocalInvocationID.x;
|
||||
|
||||
// 加载到共享内存
|
||||
if (gid < n) {
|
||||
tile[lid] = input_data[gid];
|
||||
}
|
||||
|
||||
// 屏障：等待工作组中所有调用完成加载
|
||||
barrier(); // 等同于 CUDA 的 __syncthreads()
|
||||
|
||||
// 计算 ReLU
|
||||
if (gid < n) {
|
||||
output_data[gid] = max(tile[lid], 0.0);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- 对于 ReLU，共享内存并非严格必要（该操作是按元素进行的）。但这演示了基本模式：加载到共享内存 → 屏障 → 计算 → 存储。对于需要相邻线程数据的操作（卷积、归约、softmax），共享内存是必不可少的。
|
||||
|
||||
### 并行归约（求和）
|
||||
|
||||
```glsl
|
||||
// reduce_sum.comp
|
||||
#version 450
|
||||
|
||||
layout(local_size_x = 256) in;
|
||||
|
||||
layout(set = 0, binding = 0) buffer Input { float input_data[]; };
|
||||
layout(set = 0, binding = 1) buffer Output { float partial_sums[]; };
|
||||
|
||||
layout(push_constant) uniform PushConstants { uint n; };
|
||||
|
||||
shared float sdata[256];
|
||||
|
||||
void main() {
|
||||
uint gid = gl_GlobalInvocationID.x;
|
||||
uint lid = gl_LocalInvocationID.x;
|
||||
uint wgid = gl_WorkGroupID.x;
|
||||
|
||||
// 加载到共享内存
|
||||
sdata[lid] = (gid < n) ? input_data[gid] : 0.0;
|
||||
barrier();
|
||||
|
||||
// 工作组内的树形归约
|
||||
for (uint stride = 128; stride > 0; stride >>= 1) {
|
||||
if (lid < stride) {
|
||||
sdata[lid] += sdata[lid + stride];
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
|
||||
// 线程 0 写入工作组的局部和
|
||||
if (lid == 0) {
|
||||
partial_sums[wgid] = sdata[0];
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- 这是经典的并行归约模式（与 CUDA 相同）。每个工作组产生一个局部和。第二次调度将这些局部和归约为最终结果。树形归约每一步将活跃线程减半：256 → 128 → 64 → ... → 1。
|
||||
|
||||
### 使用分块的矩阵乘法
|
||||
|
||||
```glsl
|
||||
// matmul_tiled.comp
|
||||
#version 450
|
||||
|
||||
#define TILE_SIZE 16
|
||||
|
||||
layout(local_size_x = TILE_SIZE, local_size_y = TILE_SIZE) in;
|
||||
|
||||
layout(set = 0, binding = 0) buffer MatA { float A[]; };
|
||||
layout(set = 0, binding = 1) buffer MatB { float B[]; };
|
||||
layout(set = 0, binding = 2) buffer MatC { float C[]; };
|
||||
|
||||
layout(push_constant) uniform PushConstants {
|
||||
uint M, N, K;
|
||||
};
|
||||
|
||||
shared float tileA[TILE_SIZE][TILE_SIZE];
|
||||
shared float tileB[TILE_SIZE][TILE_SIZE];
|
||||
|
||||
void main() {
|
||||
uint row = gl_GlobalInvocationID.y;
|
||||
uint col = gl_GlobalInvocationID.x;
|
||||
uint lr = gl_LocalInvocationID.y;
|
||||
uint lc = gl_LocalInvocationID.x;
|
||||
|
||||
float sum = 0.0;
|
||||
|
||||
for (uint t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
|
||||
// 将 A 和 B 的分块加载到共享内存中
|
||||
uint aCol = t * TILE_SIZE + lc;
|
||||
uint bRow = t * TILE_SIZE + lr;
|
||||
|
||||
tileA[lr][lc] = (row < M && aCol < K) ? A[row * K + aCol] : 0.0;
|
||||
tileB[lr][lc] = (bRow < K && col < N) ? B[bRow * N + col] : 0.0;
|
||||
|
||||
barrier();
|
||||
|
||||
// 计算部分点积
|
||||
for (uint k = 0; k < TILE_SIZE; k++) {
|
||||
sum += tileA[lr][k] * tileB[k][lc];
|
||||
}
|
||||
|
||||
barrier();
|
||||
}
|
||||
|
||||
if (row < M && col < N) {
|
||||
C[row * N + col] = sum;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- 这与 CUDA 版本（文件 04）中的分块算法相同，只是用了 GLSL 语法。概念完全一样：将分块加载到共享内存，屏障，计算，屏障，重复。
|
||||
|
||||
## C++ Vulkan 设置
|
||||
|
||||
- 计算着色器是简单的部分。困难的部分是创建 Vulkan 实例、分配内存、绑定缓冲区和提交命令的 C++ 样板代码。以下是完整管线的精简版本：
|
||||
|
||||
```cpp
|
||||
// vulkan_compute.cpp — 一个最小但完整的 Vulkan 计算示例
|
||||
// 编译命令: g++ -O3 -o vulkan_compute vulkan_compute.cpp -lvulkan
|
||||
// 要求: 已安装 Vulkan SDK，已从 add.comp 编译 add.spv
|
||||
|
||||
#include <vulkan/vulkan.h>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <fstream>
|
||||
#include <cassert>
|
||||
|
||||
// 辅助函数：读取 SPIR-V 文件
|
||||
std::vector<uint32_t> readSPIRV(const std::string& filename) {
|
||||
std::ifstream file(filename, std::ios::ate | std::ios::binary);
|
||||
size_t fileSize = file.tellg();
|
||||
std::vector<uint32_t> buffer(fileSize / sizeof(uint32_t));
|
||||
file.seekg(0);
|
||||
file.read(reinterpret_cast<char*>(buffer.data()), fileSize);
|
||||
return buffer;
|
||||
}
|
||||
|
||||
int main() {
|
||||
const uint32_t N = 1024;
|
||||
const size_t bufferSize = N * sizeof(float);
|
||||
|
||||
// ========== 1. 创建 Vulkan 实例 ==========
|
||||
VkApplicationInfo appInfo{};
|
||||
appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
|
||||
appInfo.apiVersion = VK_API_VERSION_1_2;
|
||||
|
||||
VkInstanceCreateInfo instanceInfo{};
|
||||
instanceInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
|
||||
instanceInfo.pApplicationInfo = &appInfo;
|
||||
|
||||
VkInstance instance;
|
||||
vkCreateInstance(&instanceInfo, nullptr, &instance);
|
||||
|
||||
// ========== 2. 选择物理设备 (GPU) ==========
|
||||
uint32_t deviceCount = 0;
|
||||
vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
|
||||
std::vector<VkPhysicalDevice> devices(deviceCount);
|
||||
vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());
|
||||
VkPhysicalDevice physicalDevice = devices[0]; // 使用第一个 GPU
|
||||
|
||||
// 打印 GPU 名称
|
||||
VkPhysicalDeviceProperties props;
|
||||
vkGetPhysicalDeviceProperties(physicalDevice, &props);
|
||||
std::cout << "使用的 GPU: " << props.deviceName << "\n";
|
||||
|
||||
// ========== 3. 查找计算队列族 ==========
|
||||
uint32_t queueFamilyCount = 0;
|
||||
vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);
|
||||
std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
|
||||
vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilies.data());
|
||||
|
||||
uint32_t computeFamily = 0;
|
||||
for (uint32_t i = 0; i < queueFamilyCount; i++) {
|
||||
if (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
|
||||
computeFamily = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 4. 创建逻辑设备和队列 ==========
|
||||
float queuePriority = 1.0f;
|
||||
VkDeviceQueueCreateInfo queueInfo{};
|
||||
queueInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
|
||||
queueInfo.queueFamilyIndex = computeFamily;
|
||||
queueInfo.queueCount = 1;
|
||||
queueInfo.pQueuePriorities = &queuePriority;
|
||||
|
||||
VkDeviceCreateInfo deviceInfo{};
|
||||
deviceInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
|
||||
deviceInfo.queueCreateInfoCount = 1;
|
||||
deviceInfo.pQueueCreateInfos = &queueInfo;
|
||||
|
||||
VkDevice device;
|
||||
vkCreateDevice(physicalDevice, &deviceInfo, nullptr, &device);
|
||||
|
||||
VkQueue computeQueue;
|
||||
vkGetDeviceQueue(device, computeFamily, 0, &computeQueue);
|
||||
|
||||
// ========== 5. 分配缓冲区 (A, B, C) ==========
|
||||
// 为简洁起见，这里使用主机可见内存（较慢但更简单）
|
||||
auto createBuffer = [&](VkBuffer& buffer, VkDeviceMemory& memory) {
|
||||
VkBufferCreateInfo bufInfo{};
|
||||
bufInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
|
||||
bufInfo.size = bufferSize;
|
||||
bufInfo.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
|
||||
vkCreateBuffer(device, &bufInfo, nullptr, &buffer);
|
||||
|
||||
VkMemoryRequirements memReqs;
|
||||
vkGetBufferMemoryRequirements(device, buffer, &memReqs);
|
||||
|
||||
// 查找主机可见的内存类型
|
||||
VkPhysicalDeviceMemoryProperties memProps;
|
||||
vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProps);
|
||||
uint32_t memType = 0;
|
||||
for (uint32_t i = 0; i < memProps.memoryTypeCount; i++) {
|
||||
if ((memReqs.memoryTypeBits & (1 << i)) &&
|
||||
(memProps.memoryTypes[i].propertyFlags &
|
||||
(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))) {
|
||||
memType = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
VkMemoryAllocateInfo allocInfo{};
|
||||
allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
|
||||
allocInfo.allocationSize = memReqs.size;
|
||||
allocInfo.memoryTypeIndex = memType;
|
||||
vkAllocateMemory(device, &allocInfo, nullptr, &memory);
|
||||
vkBindBufferMemory(device, buffer, memory, 0);
|
||||
};
|
||||
|
||||
VkBuffer bufA, bufB, bufC;
|
||||
VkDeviceMemory memA, memB, memC;
|
||||
createBuffer(bufA, memA);
|
||||
createBuffer(bufB, memB);
|
||||
createBuffer(bufC, memC);
|
||||
|
||||
// ========== 6. 填充输入缓冲区 ==========
|
||||
float* ptrA;
|
||||
vkMapMemory(device, memA, 0, bufferSize, 0, (void**)&ptrA);
|
||||
for (uint32_t i = 0; i < N; i++) ptrA[i] = 1.0f;
|
||||
vkUnmapMemory(device, memA);
|
||||
|
||||
float* ptrB;
|
||||
vkMapMemory(device, memB, 0, bufferSize, 0, (void**)&ptrB);
|
||||
for (uint32_t i = 0; i < N; i++) ptrB[i] = 2.0f;
|
||||
vkUnmapMemory(device, memB);
|
||||
|
||||
// ========== 7. 创建计算管线 ==========
|
||||
auto spirvCode = readSPIRV("add.spv");
|
||||
VkShaderModuleCreateInfo shaderInfo{};
|
||||
shaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
|
||||
shaderInfo.codeSize = spirvCode.size() * sizeof(uint32_t);
|
||||
shaderInfo.pCode = spirvCode.data();
|
||||
VkShaderModule shaderModule;
|
||||
vkCreateShaderModule(device, &shaderInfo, nullptr, &shaderModule);
|
||||
|
||||
// 描述符集布局（告诉 Vulkan 缓冲区绑定的信息）
|
||||
VkDescriptorSetLayoutBinding bindings[3] = {};
|
||||
for (int i = 0; i < 3; i++) {
|
||||
bindings[i].binding = i;
|
||||
bindings[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
|
||||
bindings[i].descriptorCount = 1;
|
||||
bindings[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
|
||||
}
|
||||
|
||||
VkDescriptorSetLayoutCreateInfo layoutInfo{};
|
||||
layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
|
||||
layoutInfo.bindingCount = 3;
|
||||
layoutInfo.pBindings = bindings;
|
||||
VkDescriptorSetLayout descLayout;
|
||||
vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descLayout);
|
||||
|
||||
// 推送常量范围
|
||||
VkPushConstantRange pushRange{};
|
||||
pushRange.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
|
||||
pushRange.offset = 0;
|
||||
pushRange.size = sizeof(uint32_t);
|
||||
|
||||
// 管线布局
|
||||
VkPipelineLayoutCreateInfo pipeLayoutInfo{};
|
||||
pipeLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
|
||||
pipeLayoutInfo.setLayoutCount = 1;
|
||||
pipeLayoutInfo.pSetLayouts = &descLayout;
|
||||
pipeLayoutInfo.pushConstantRangeCount = 1;
|
||||
pipeLayoutInfo.pPushConstantRanges = &pushRange;
|
||||
VkPipelineLayout pipelineLayout;
|
||||
vkCreatePipelineLayout(device, &pipeLayoutInfo, nullptr, &pipelineLayout);
|
||||
|
||||
// 计算管线
|
||||
VkComputePipelineCreateInfo pipeInfo{};
|
||||
pipeInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
|
||||
pipeInfo.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
|
||||
pipeInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
|
||||
pipeInfo.stage.module = shaderModule;
|
||||
pipeInfo.stage.pName = "main";
|
||||
pipeInfo.layout = pipelineLayout;
|
||||
VkPipeline pipeline;
|
||||
vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeInfo, nullptr, &pipeline);
|
||||
|
||||
// ========== 8. 描述符集（将缓冲区绑定到着色器） ==========
|
||||
VkDescriptorPoolSize poolSize{};
|
||||
poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
|
||||
poolSize.descriptorCount = 3;
|
||||
|
||||
VkDescriptorPoolCreateInfo poolInfo{};
|
||||
poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
|
||||
poolInfo.maxSets = 1;
|
||||
poolInfo.poolSizeCount = 1;
|
||||
poolInfo.pPoolSizes = &poolSize;
|
||||
VkDescriptorPool descPool;
|
||||
vkCreateDescriptorPool(device, &poolInfo, nullptr, &descPool);
|
||||
|
||||
VkDescriptorSetAllocateInfo descAllocInfo{};
|
||||
descAllocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
|
||||
descAllocInfo.descriptorPool = descPool;
|
||||
descAllocInfo.descriptorSetCount = 1;
|
||||
descAllocInfo.pSetLayouts = &descLayout;
|
||||
VkDescriptorSet descSet;
|
||||
vkAllocateDescriptorSets(device, &descAllocInfo, &descSet);
|
||||
|
||||
// 将缓冲区引用写入描述符集
|
||||
VkDescriptorBufferInfo bufInfos[3] = {
|
||||
{bufA, 0, bufferSize}, {bufB, 0, bufferSize}, {bufC, 0, bufferSize}
|
||||
};
|
||||
VkWriteDescriptorSet writes[3] = {};
|
||||
for (int i = 0; i < 3; i++) {
|
||||
writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
|
||||
writes[i].dstSet = descSet;
|
||||
writes[i].dstBinding = i;
|
||||
writes[i].descriptorCount = 1;
|
||||
writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
|
||||
writes[i].pBufferInfo = &bufInfos[i];
|
||||
}
|
||||
vkUpdateDescriptorSets(device, 3, writes, 0, nullptr);
|
||||
|
||||
// ========== 9. 记录和提交命令缓冲区 ==========
|
||||
VkCommandPoolCreateInfo cmdPoolInfo{};
|
||||
cmdPoolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
|
||||
cmdPoolInfo.queueFamilyIndex = computeFamily;
|
||||
VkCommandPool cmdPool;
|
||||
vkCreateCommandPool(device, &cmdPoolInfo, nullptr, &cmdPool);
|
||||
|
||||
VkCommandBufferAllocateInfo cmdAllocInfo{};
|
||||
cmdAllocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
|
||||
cmdAllocInfo.commandPool = cmdPool;
|
||||
cmdAllocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
|
||||
cmdAllocInfo.commandBufferCount = 1;
|
||||
VkCommandBuffer cmdBuf;
|
||||
vkAllocateCommandBuffers(device, &cmdAllocInfo, &cmdBuf);
|
||||
|
||||
VkCommandBufferBeginInfo beginInfo{};
|
||||
beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
|
||||
vkBeginCommandBuffer(cmdBuf, &beginInfo);
|
||||
|
||||
vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
|
||||
vkCmdBindDescriptorSets(cmdBuf, VK_PIPELINE_BIND_POINT_COMPUTE,
|
||||
pipelineLayout, 0, 1, &descSet, 0, nullptr);
|
||||
vkCmdPushConstants(cmdBuf, pipelineLayout, VK_SHADER_STAGE_COMPUTE_BIT,
|
||||
0, sizeof(uint32_t), &N);
|
||||
vkCmdDispatch(cmdBuf, (N + 255) / 256, 1, 1); // 启动工作组
|
||||
|
||||
vkEndCommandBuffer(cmdBuf);
|
||||
|
||||
// 提交
|
||||
VkFenceCreateInfo fenceInfo{};
|
||||
fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
|
||||
VkFence fence;
|
||||
vkCreateFence(device, &fenceInfo, nullptr, &fence);
|
||||
|
||||
VkSubmitInfo submitInfo{};
|
||||
submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
|
||||
submitInfo.commandBufferCount = 1;
|
||||
submitInfo.pCommandBuffers = &cmdBuf;
|
||||
vkQueueSubmit(computeQueue, 1, &submitInfo, fence);
|
||||
vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX);
|
||||
|
||||
// ========== 10. 读取结果 ==========
|
||||
float* ptrC;
|
||||
vkMapMemory(device, memC, 0, bufferSize, 0, (void**)&ptrC);
|
||||
std::cout << "结果: c[0]=" << ptrC[0] << " c[1]=" << ptrC[1]
|
||||
<< " (期望值 3.0)\n";
|
||||
bool correct = true;
|
||||
for (uint32_t i = 0; i < N; i++) {
|
||||
if (ptrC[i] != 3.0f) { correct = false; break; }
|
||||
}
|
||||
std::cout << (correct ? "全部正确" : "发现错误") << "\n";
|
||||
vkUnmapMemory(device, memC);
|
||||
|
||||
// ========== 清理（简写） ==========
|
||||
vkDestroyFence(device, fence, nullptr);
|
||||
vkDestroyCommandPool(device, cmdPool, nullptr);
|
||||
vkDestroyPipeline(device, pipeline, nullptr);
|
||||
vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
|
||||
vkDestroyDescriptorPool(device, descPool, nullptr);
|
||||
vkDestroyDescriptorSetLayout(device, descLayout, nullptr);
|
||||
vkDestroyShaderModule(device, shaderModule, nullptr);
|
||||
vkDestroyBuffer(device, bufA, nullptr); vkFreeMemory(device, memA, nullptr);
|
||||
vkDestroyBuffer(device, bufB, nullptr); vkFreeMemory(device, memB, nullptr);
|
||||
vkDestroyBuffer(device, bufC, nullptr); vkFreeMemory(device, memC, nullptr);
|
||||
vkDestroyDevice(device, nullptr);
|
||||
vkDestroyInstance(instance, nullptr);
|
||||
|
||||
return 0;
|
||||
}
|
||||
```
|
||||
|
||||
- **是的，向量加法就需要大约 200 行代码。** 相比之下 CUDA 只需要大约 30 行。这就是显式性的代价。但请注意：每一行都有其目的。没有隐藏的驱动决策，没有隐式同步，没有意外的内存分配。你控制一切。
|
||||
|
||||
- 在实践中，你可以将这些样板代码封装到辅助库中（或使用现有的库，如 **vk-bootstrap**、用于内存分配的 **VMA**，或专注于 ML 的 Vulkan 计算库 **kompute**）。
|
||||
|
||||
## Kompute：为 ML 简化的 Vulkan
|
||||
|
||||
- **Kompute** 是一个开源 C++ 库，封装了 Vulkan 用于 GPU 计算的样板代码。同样的向量加法变成：
|
||||
|
||||
```cpp
|
||||
#include <kompute/Kompute.hpp>
|
||||
|
||||
int main() {
|
||||
kp::Manager mgr;
|
||||
|
||||
auto tensorA = mgr.tensor({1, 1, 1, 1, 1});
|
||||
auto tensorB = mgr.tensor({2, 2, 2, 2, 2});
|
||||
auto tensorC = mgr.tensor({0, 0, 0, 0, 0});
|
||||
|
||||
std::string shader = R"(
|
||||
#version 450
|
||||
layout(local_size_x = 1) in;
|
||||
layout(set=0, binding=0) buffer A { float a[]; };
|
||||
layout(set=0, binding=1) buffer B { float b[]; };
|
||||
layout(set=0, binding=2) buffer C { float c[]; };
|
||||
void main() {
|
||||
uint i = gl_GlobalInvocationID.x;
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
)";
|
||||
|
||||
auto algorithm = mgr.algorithm({tensorA, tensorB, tensorC},
|
||||
kompute::Shader::compile_source(shader));
|
||||
|
||||
mgr.sequence()
|
||||
->record<kp::OpTensorSyncDevice>({tensorA, tensorB, tensorC})
|
||||
->record<kp::OpAlgoDispatch>(algorithm)
|
||||
->record<kp::OpTensorSyncLocal>({tensorC})
|
||||
->eval();
|
||||
|
||||
// tensorC 现在包含 [3, 3, 3, 3, 3]
|
||||
}
|
||||
```
|
||||
|
||||
- 可读性强多了。Kompute 处理实例创建、设备选择、内存分配、描述符集和命令缓冲区管理。你只需关注着色器和数据。
|
||||
|
||||
## WebGPU：浏览器中的 GPU 计算
|
||||
|
||||
- **WebGPU** 是 WebGL 的继任者，提供从 JavaScript 访问现代 GPU 的能力。它基于 Vulkan（Linux/Android）、Metal（macOS/iOS）和 DirectX 12（Windows）构建，抽象了平台差异。
|
||||
|
||||
- WebGPU 使用 **WGSL**（WebGPU 着色语言）而非 GLSL：
|
||||
|
||||
```wgsl
|
||||
// add.wgsl — WebGPU 计算着色器
|
||||
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
||||
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
||||
@group(0) @binding(2) var<storage, read_write> c: array<f32>;
|
||||
|
||||
@compute @workgroup_size(256)
|
||||
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
|
||||
let i = id.x;
|
||||
c[i] = a[i] + b[i];
|
||||
}
|
||||
```
|
||||
|
||||
- **JavaScript 设置**（精简版）：
|
||||
|
||||
```javascript
|
||||
const adapter = await navigator.gpu.requestAdapter();
|
||||
const device = await adapter.requestDevice();
|
||||
|
||||
// 创建缓冲区
|
||||
const bufferA = device.createBuffer({ size: N * 4, usage: GPUBufferUsage.STORAGE, mappedAtCreation: true });
|
||||
new Float32Array(bufferA.getMappedRange()).fill(1.0);
|
||||
bufferA.unmap();
|
||||
|
||||
// ...（B 和 C 类似）
|
||||
|
||||
// 从 WGSL 着色器创建管线
|
||||
const pipeline = device.createComputePipeline({
|
||||
layout: 'auto',
|
||||
compute: { module: device.createShaderModule({ code: wgslSource }), entryPoint: 'main' }
|
||||
});
|
||||
|
||||
// 调度
|
||||
const encoder = device.createCommandEncoder();
|
||||
const pass = encoder.beginComputePass();
|
||||
pass.setPipeline(pipeline);
|
||||
pass.setBindGroup(0, bindGroup);
|
||||
pass.dispatchWorkgroups(Math.ceil(N / 256));
|
||||
pass.end();
|
||||
device.queue.submit([encoder.finish()]);
|
||||
```
|
||||
|
||||
- **为什么 WebGPU 对 ML 很重要**：在浏览器中运行推理意味着没有服务器成本、没有延迟，且用户数据永远不会离开设备。像 **ONNX Runtime Web** 和 **Transformers.js** 这样的库使用 WebGPU 完全在客户端运行模型（包括小型 LLM）。
|
||||
|
||||
## 何时使用 Vulkan
|
||||
|
||||
| 场景 | 使用 Vulkan？ | 原因 / 替代方案 |
|
||||
|------|-------------|----------------|
|
||||
| ML 训练 | 否 | CUDA/Triton 在 NVIDIA 上更简单更快速 |
|
||||
| NVIDIA GPU 上的推理 | 否 | TensorRT 或 CUDA 更好 |
|
||||
| AMD/Intel GPU 上的推理 | **是** | 唯一跨厂商的 GPU 计算选项 |
|
||||
| 移动端推理（Android） | **是** | Vulkan 是 Android 上的标准 GPU API |
|
||||
| 移动端推理（iOS） | 否 | 直接使用 Metal（MoltenVK 增加开销） |
|
||||
| 浏览器推理 | **WebGPU** | 基于 Vulkan/Metal/DX12 |
|
||||
| 游戏引擎 + ML | **是** | 引擎已使用 Vulkan 进行渲染 |
|
||||
| 跨平台库 | **是** | 一套代码支持所有 GPU 厂商 |
|
||||
| 学习 GPU 编程 | 视情况而定 | CUDA 更容易上手；Vulkan 能学到更多 |
|
||||
|
||||
## 编码任务（使用 g++ -lvulkan 编译，需要 Vulkan SDK）
|
||||
|
||||
1. 编译并运行上面的向量加法示例。修改着色器以计算 `c[i] = a[i] * b[i] + a[i]`（融合乘加）并验证结果。
|
||||
|
||||
2. 编写一个计算着色器，使用共享内存对一行数据应用 softmax（包括最大值和求和归约步骤）。用已知值进行测试。
|
||||
|
||||
```glsl
|
||||
// softmax.comp — 编译命令: glslangValidator -V softmax.comp -o softmax.spv
|
||||
#version 450
|
||||
|
||||
#define WG_SIZE 256
|
||||
|
||||
layout(local_size_x = WG_SIZE) in;
|
||||
|
||||
layout(set = 0, binding = 0) buffer Input { float input_data[]; };
|
||||
layout(set = 0, binding = 1) buffer Output { float output_data[]; };
|
||||
|
||||
layout(push_constant) uniform PC { uint n; };
|
||||
|
||||
shared float sdata[WG_SIZE];
|
||||
|
||||
void main() {
|
||||
uint gid = gl_GlobalInvocationID.x;
|
||||
uint lid = gl_LocalInvocationID.x;
|
||||
|
||||
// 步骤 1：找最大值（数值稳定性）
|
||||
sdata[lid] = (gid < n) ? input_data[gid] : -1e30;
|
||||
barrier();
|
||||
for (uint s = WG_SIZE / 2; s > 0; s >>= 1) {
|
||||
if (lid < s) sdata[lid] = max(sdata[lid], sdata[lid + s]);
|
||||
barrier();
|
||||
}
|
||||
float maxVal = sdata[0];
|
||||
barrier();
|
||||
|
||||
// 步骤 2：计算 exp(x - max)
|
||||
float expVal = (gid < n) ? exp(input_data[gid] - maxVal) : 0.0;
|
||||
sdata[lid] = expVal;
|
||||
barrier();
|
||||
|
||||
// 步骤 3：exp 值求和
|
||||
for (uint s = WG_SIZE / 2; s > 0; s >>= 1) {
|
||||
if (lid < s) sdata[lid] += sdata[lid + s];
|
||||
barrier();
|
||||
}
|
||||
float sumExp = sdata[0];
|
||||
|
||||
// 步骤 4：归一化
|
||||
if (gid < n) {
|
||||
output_data[gid] = expVal / sumExp;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
3. 修改 C++ 宿主代码以对计算着色器进行基准测试：使用 Vulkan 时间戳查询或 CPU 端栅栏对调度（排除设置阶段）计时，并计算以 GB/s 为单位的实际带宽。
|
||||
Reference in New Issue
Block a user