Rust in Everything
基于 burn 0.22.1

CPU 上的 matmul 慢了 27 倍

主线位置：训练主线的 forward 和 backward，在 CPU 上跑：每一步的时间几乎都花在 Linear 的 matmul 上。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  dep["default-features false"] --> nf["burn-flex 只剩 std"]
  nf --> g1["gemm 单线程"]
  g1 --> slow["1024 方阵 21.4 ms"]
  add["补 rayon simd"] --> g8["gemm 用全部核"]
  g8 --> mid["3.5 ms"]
  amx["再加 apple-amx"] --> fast["1.4 到 1.6 ms"]

实验代码：PyTorch 侧 cpu_bench.py；burn 侧四个 crate，代码是同一个文件，只有 Cargo.toml 里的 burn feature 不同：

crateburn 依赖
pitb-exec-cpudefault-features = false，std、flex、autodiff、optim（本卷模板）
pitb-exec-cpu-simd同上，再加 simd、rayon
pitb-exec-cpu-default默认 feature，加 flex、autodiff、optim
pitb-exec-cpu-amx模板，再加 simd、rayon、apple-amx

全部 release 构建。每项先热身，再计时 20 次，报最小值和中位数；整组跑两轮，下面贴的是两轮的原始输出。测量时 uptime 的 1 分钟负载在 4 到 7 之间。平台 M2 Max（12 核：8 个性能核、4 个能效核）、macOS 27、rustc 1.99.0。

PyTorch 的写法

torch.set_num_threads(int(sys.argv[2]))
a, b = torch.rand(n, n), torch.rand(n, n)
a @ b
ts = []
for _ in range(reps):
  t = time.perf_counter()
  a @ b
  ts.append(time.perf_counter() - t)

torch threads 1: 1024x1024 matmul, 20 runs: min 0.855ms median 0.922ms
torch threads 8: 1024x1024 matmul, 20 runs: min 0.799ms median 0.832ms
...
torch threads 1: 1024x1024 matmul, 20 runs: min 0.798ms median 0.826ms
torch threads 8: 1024x1024 matmul, 20 runs: min 0.805ms median 0.833ms

pip install torch 装上的 wheel 不用任何配置，CPU 上的 matmul 已经是快的。线程数从 8 降到 1 也不变：这个 wheel 的 torch.__config__.show() 写着 BLAS_INFO=accelerate，matmul 交给苹果的 Accelerate，set_num_threads 管不到它（计时一页查过）。

再看一个完整的训练步：MLP 784→512→10，batch 256，交叉熵，SGD，从 forward 到 step：

torch threads 1: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 0.574ms median 0.770ms
torch threads 8: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 0.773ms median 0.829ms
...
torch threads 1: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 0.653ms median 0.861ms
torch threads 8: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 0.698ms median 0.790ms

搬到 burn 的现象

按本卷模板的依赖（default-features = false，只开需要的几个 feature）写同样的 matmul：

let device = Device::flex();
let a = Tensor::<2>::random([n, n], Distribution::Default, &device);
let b = Tensor::<2>::random([n, n], Distribution::Default, &device);
let _ = a.clone().matmul(b.clone()).into_data(); // warm-up
let mut ts = Vec::new();
for _ in 0..reps {
  let t = Instant::now();
  let c = a.clone().matmul(b.clone());
  device.sync().unwrap();
  ts.push(t.elapsed());
  drop(c);
}

训练步是同一个 MLP，建在 Device::flex().autodiff() 上，每步 forward → backward → GradientsParams::from_grads → optim.step，再 sync。四种构建两轮：

## round 1
pitb-exec-cpu: 1024x1024 matmul, 20 runs: min 21.424ms median 22.005ms
pitb-exec-cpu-simd: 1024x1024 matmul, 20 runs: min 3.508ms median 3.632ms
pitb-exec-cpu-default: 1024x1024 matmul, 20 runs: min 3.599ms median 4.116ms
pitb-exec-cpu-amx: 1024x1024 matmul, 20 runs: min 1.563ms median 1.721ms
...
pitb-exec-cpu: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 4.410ms median 5.018ms
pitb-exec-cpu-simd: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.747ms median 2.083ms
pitb-exec-cpu-default: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.769ms median 2.249ms
pitb-exec-cpu-amx: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.442ms median 1.847ms
...
## round 2
pitb-exec-cpu: 1024x1024 matmul, 20 runs: min 21.634ms median 22.048ms
pitb-exec-cpu-simd: 1024x1024 matmul, 20 runs: min 3.496ms median 3.802ms
pitb-exec-cpu-default: 1024x1024 matmul, 20 runs: min 3.510ms median 3.751ms
pitb-exec-cpu-amx: 1024x1024 matmul, 20 runs: min 1.609ms median 1.824ms
...
pitb-exec-cpu: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 4.451ms median 4.932ms
pitb-exec-cpu-simd: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.624ms median 2.204ms
pitb-exec-cpu-default: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.874ms median 2.176ms
pitb-exec-cpu-amx: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.235ms median 1.559ms

取两轮里较小的最小值：

构建1024 方阵 matmul和 PyTorch 比训练一步和 PyTorch 比
PyTorch（Accelerate）0.80 ms10.57 ms1
burn 模板21.4 ms27 倍4.41 ms7.7 倍
模板加 simd、rayon3.50 ms4.4 倍1.62 ms2.8 倍
默认 feature3.51 ms4.4 倍1.77 ms3.1 倍
模板加 simd、rayon、apple-amx1.56 ms2.0 倍1.24 ms2.2 倍

代码一行没改，只是依赖写法不同，matmul 差了 27 倍。训练一步差得少一些，因为这个 MLP 的 matmul 小（256×784 乘 784×512），其余的 op 和调度开销占了更大的比例。

原因

Flex 的 matmul 交给 gemm crate（Flex：CPU 后端），用不用多线程由 get_parallelism 决定：

fn get_parallelism(m: usize, n: usize, k: usize) -> gemm::Parallelism {
    let ops = m.saturating_mul(n).saturating_mul(k);
    if ops >= PARALLEL_THRESHOLD {
        #[cfg(feature = "rayon")]
        {
            gemm::Parallelism::Rayon(0) // 0 = use all available threads
        }
        #[cfg(not(feature = "rayon"))]
        {
            gemm::Parallelism::None
        }
    } else {
        gemm::Parallelism::None
    }
}

没有 rayon feature，再大的矩阵也是 Parallelism::None，单线程。burn-flex 的 default 是 ["std", "simd", "rayon"]，burn 的默认 feature 沿 burn-core、burn-dispatch 一路把它打开（这条链和它对 GPU 后端的影响见 burn 卷的关了默认 feature，GPU 变慢）。写了 default-features = false，这条链就断了，Flex 只剩 std。用 cargo tree 看四个 crate 里 burn-flex 实际打开的 feature：

pitb-exec-cpu: burn-flex v0.22.1 [std]
pitb-exec-cpu-simd: burn-flex v0.22.1 [rayon,simd,std]
pitb-exec-cpu-default: burn-flex v0.22.1 [default,rayon,simd,std]
pitb-exec-cpu-amx: burn-flex v0.22.1 [apple-amx,rayon,simd,std]

matmul 上起作用的是 rayon，不是 simd。把 pitb-exec-cpu-simd 的线程数用 RAYON_NUM_THREADS 固定成 1，它和模板一样慢；apple-amx 在 1 个线程时也已经快很多：

### RAYON_NUM_THREADS=1
pitb-exec-cpu-simd: 1024x1024 matmul, 20 runs: min 21.599ms median 22.114ms
pitb-exec-cpu-simd: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 4.520ms median 5.142ms
### RAYON_NUM_THREADS=8
pitb-exec-cpu-simd: 1024x1024 matmul, 20 runs: min 3.237ms median 3.746ms
pitb-exec-cpu-simd: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.650ms median 2.079ms
### RAYON_NUM_THREADS=1
pitb-exec-cpu-amx: 1024x1024 matmul, 20 runs: min 3.488ms median 3.647ms
pitb-exec-cpu-amx: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.038ms median 1.108ms
### RAYON_NUM_THREADS=8
pitb-exec-cpu-amx: 1024x1024 matmul, 20 runs: min 1.423ms median 1.540ms
pitb-exec-cpu-amx: MLP 784-512-10 batch 256 train step, 20 runs after 3 warm-up: min 1.067ms median 1.429ms

burn-flex 的 simd 打开的是 macerator、aligned-vec 和 gemm/wasm-simd128-enable，后者从名字看只针对 wasm；f32 matmul 交给 gemm，单线程时有没有 simd 都是 21.4 到 21.6 ms。simd 管的是 Flex 自己写的内核：在 0.22.1 的 burn-flex 里，#[cfg(feature = "simd")] 出现在归约、掩码、激活、比较等算子里，matmul.rs 里用到它的只有 Int 的 matmul。它会改变 f32 求和的结果（怎么读这一卷实测过）。

apple-amx 打开 gemm/experimental-apple-amx，用 M 系列芯片的 AMX 协处理器，名字里写着 experimental，不在默认 feature 里。即使打开它，8 线程的 1024 方阵 matmul 也还是 PyTorch 的 1.8 到 2 倍。PyTorch 那一侧的 0.80 ms 是 Accelerate 给的，它在 burn 这边没有对应：burn 的 accelerate feature 接的是 ndarray 后端的 blas-accelerate，离线构建失败，本页没有测（计时一页试过）。

改法

在 CPU 上训练或推理，依赖里把 rayon 补回来，simd 一起加上；Apple Silicon 上再加 apple-amx：

[dependencies]
burn = { version = "=0.22.1", default-features = false, features = ["std", "flex", "autodiff", "optim", "simd", "rayon", "apple-amx"] }

或者干脆不关默认 feature（burn = { version = "=0.22.1", features = ["flex", "autodiff", "optim"] }），代价是多出 GPU 相关的依赖和更大的二进制（为什么要转量过二进制大小）。上表里两种写法的 matmul 一样快。

几点要一起考虑：

  • 换构建不改变 matmul 的结果：计时一页用确定的输入比过，模板、默认 feature、apple-amx 三种构建的 matmul 逐位相同。逐元素和归约在 simd 下的结果会变，和 PyTorch 逐位对照的实验要说明用的是哪种构建。
  • apple-amx 的训练步在 1 个线程时（1.04 ms）比 8 个线程（1.07 ms）还快一点，矩阵小的时候多线程不一定划算。模型尺寸定了以后，用 RAYON_NUM_THREADS 试几个值。
  • 和 PyTorch 比 CPU 速度时，torch.set_num_threads 管不到 Accelerate 的 matmul，「两边都固定成 1 个线程」不等于两边用了同样的算力。

怎么检查

  • 看实际打开的 feature：在自己的 crate 里跑 cargo tree -f '{p} [{f}]' -i burn-flex --depth 0，方括号里没有 rayon，Flex 的 matmul 就是单线程。
  • 量一次 1024 方阵 matmul：这台机器上单线程约 21 ms，多线程约 3.5 ms，加 apple-amx 约 1.5 ms，PyTorch 约 0.8 ms。

小结

  • PyTorch 的 CPU wheel 用 Accelerate 做 matmul，开箱就快（1024 方阵 0.80 ms）；burn 用 default-features = false 时 Flex 没有 rayon，gemm 单线程，21.4 ms，慢 27 倍，训练一步慢 7.7 倍。
  • 起作用的是 rayon：补上之后 3.5 ms；再加 apple-amx 到 1.4 到 1.6 ms，仍是 PyTorch 的 2 倍左右。simd 不改变单线程 matmul 的时间。
  • 用 cargo tree -i burn-flex 确认打开了哪些 feature；比速度时写清楚构建，PyTorch 的 set_num_threads 管不到 Accelerate。

💬 关联讨论

发起讨论
加载中...