当我把 Numba 带入道场：与 Rust 和 CUDA 的大乱斗

Source: Dev.to

在深入探讨之前，我想先感谢 Shreyan Ghosh (@zenoguy) 以及他那篇精彩的文章 《当时间成为变量——我与 Numba 的旅程笔记》（dev.to 链接）。

他的文章捕捉到了计算的美好之处：实验的乐趣、代码飞速运行的激动，以及对 “如果这样会怎样？” 的好奇心。

“在算法与硬件之间的某个地方，Numba 不仅让我的代码更快，还让探索变得更轻松。”

阅读他的基准测试后，我不禁想问：当我们把 Rust 加入进来会怎样？原生 CUDA 又会怎样？硬件到底在什么地方会捧场？

于是我搭建了一个道场。来一场对决吧。

🎯 挑战

与 Shreyan 原始实验相同的挑战：

f(x) = sqrt(x² + 1) × sin(x) + cos(x/2)

计算 2000 万 个元素。
简单的数学。最大化优化。谁会赢？

🥊 竞争者

Python 队 🐍

Rust 队 🦀

GPU 队 🎮

🏗️ 设置

我希望它是 可复现且公平：

• 所有实现使用相同的计算。
• 相同的数组大小（2000 万 float64 元素）。
• 相同的随机种子（42）。
• 多次热身运行以消除 JIT/缓存影响。
• 取多次运行的 最小值（噪声最小）。

完整的基准套件是开源的：github.com/copyleftdev/numba-dojo

# Run everything yourself
git clone https://github.com/copyleftdev/numba-dojo.git
cd numba-dojo
make all

Source: …

📊 结果

完整排行榜

加速可视化

类别冠军

Source: …

🔬 我学到的

1. GPU ≫ CPU（在适用的情况下）

• RTX 3080 Ti: 0.21 ms → 3,255× 比 NumPy 快。
• 对于极度并行、逐元素的工作负载，GPU 完全处于另一个层次。
• 巨大的并行度（80 个 SM，数千个核心）彻底压制了顺序执行。

2. FP32 ≈ 在消费级 GPU 上比 FP64 快约 20 倍

CUDA FP64: 4.11 ms
CUDA FP32: 0.21 ms   ← 快 20 倍！

• 消费级 GeForce GPU 的 FP64 单元非常少（≈ FP32 吞吐量的 1/32）。
• 如果你的算法能够容忍单精度，使用 FP32。

3. Rust ≈ Numba JIT（单线程）

Rust (single‑threaded): 555.62 ms
Numba JIT:             558.30 ms

• 两者都编译为 LLVM IR，生成的代码几乎相同。
• 微小的差异仅是噪声，验证了 Numba 的宣称：“感觉像 Python，表现却像 C。”

4. Rust 在并行上胜过 Numba（约快 20 %）

Rust Parallel (Rayon): 12.39 ms
Numba Parallel:       15.55 ms

• Rayon 的工作窃取调度器比 Numba 的线程开销更低。
• 对于生产环境中的 CPU 并行工作负载，Rust 显然更有优势。

5. 我们遇到了内存带宽瓶颈

对 FP32 CUDA 核心的分析得到：

Time:       0.21 ms
Bandwidth:  ~777 GB/s achieved
Theoretical: 912 GB/s (RTX 3080 Ti)
Efficiency: 85 %

• GPU 正在以 85 % 的峰值内存带宽 运行。
• 核心大多处于空闲状态；瓶颈在于数据的进出内存。

结论

• GPU 在纯数据并行工作中占据主导。
• 在消费级硬件上，FP32 是最佳选择。
• Rust 在 CPU 上表现不俗（甚至快于 Numba）。
• 内存带宽是两者的最终瓶颈。

欢迎克隆仓库并自行运行基准测试！

等待数据* – 没有算法能超越物理

这就是 Roofline Model 的实际应用：

                    Peak Compute
                         /
                        /
Performance            /
                      /  ←  We're here (memory‑bound)
                     /
                    /
            ──────────────────────
                Memory Bandwidth

对于此工作负载，算术强度很低（每字节的操作很少），因此我们已经达到了内存带宽的上限。

🧪 代码

以下是三个独立实现的相同 kernel。

1️⃣ Numba（原文的英雄）

from numba import njit, prange
import numpy as np

@njit(parallel=True, fastmath=True, cache=True)
def compute_numba_parallel(arr, out):
    """逐元素计算 sqrt(x²+1)·sin(x) + cos(0.5·x)。"""
    n = len(arr)
    for i in prange(n):
        x = arr[i]
        out[i] = np.sqrt(x * x + 1.0) * np.sin(x) + np.cos(0.5 * x)

只需添加 @njit；其余部分完全是 NumPy 风格的 Python。

2️⃣ Rust（挑战者）

use rayon::prelude::*;

/// 逐元素计算 sqrt(x²+1)·sin(x) + cos(0.5·x)。
pub fn compute_parallel(arr: &[f64], out: &mut [f64]) {
    out.par_iter_mut()
        .zip(arr.par_iter())
        .for_each(|(o, &x)| {
            *o = (x * x + 1.0).sqrt() * x.sin() + (0.5 * x).cos();
        });
}

rayon 让数据并行像普通迭代器一样自然。

3️⃣ CUDA C++（冠军）

#include <cuda_runtime.h>
#include <cmath>

__global__ void compute_fp32(const float *arr, float *out, size_t n) {
    // 每个元素对应一个线程
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        float x = arr[idx];
        out[idx] = sqrtf(x * x + 1.0f) * sinf(x) + cosf(0.5f * x);
    }
}

/* 启动内核的辅助函数 */
void launch_compute(const float *d_arr, float *d_out, size_t n) {
    const int threads_per_block = 256;
    const int blocks = (n + threads_per_block - 1) / threads_per_block;
    compute_fp32<<<blocks, threads_per_block>>>(d_arr, d_out, n);
    cudaDeviceSynchronize();   // 为简洁起见省略错误检查
}

一个直接的网格步幅内核，将每个线程映射到数组的每个元素。

参考文献

• Original inspiration – 引发此比较的文章。
• Numba documentation – https://numba.pydata.org/
• Rayon (Rust) – https://github.com/rayon-rs/rayon
• Roofline Model – https://en.wikipedia.org/wiki/Roofline_model

保持实验。保持玩耍。这正是计算的意义所在。 ✨

你最喜欢的性能优化故事是什么？在评论中分享吧！