SIMD 与向量化

SIMD 原理深度解析

现代 CPU 的向量计算能力

现代 x86-64 处理器包含专用的向量处理单元（SIMD 执行单元），可以在一个时钟周期内对多个数据元素执行相同操作。这是 NumPy/PyTorch 等库能够快速处理大型数组的根本原因——它们的 C 底层充分利用了这些指令集。

主流 SIMD 指令集对比

SIMD[DType, size] 类型

Mojo 的 SIMD 一等公民

Mojo 将 SIMD 向量作为内置类型 SIMD[dtype, size]，而不是像 C++ 那样依赖编译器内置函数（intrinsics）或 Python 那样通过 NumPy 间接调用。这让 SIMD 编程变得直观。

from sys.info import simdwidthof

fn simd_basics():
    # 创建 SIMD 向量（8个 Float32）
    var a = SIMD[DType.float32, 8](1.0, 2.0, 3.0, 4.0,
                                     5.0, 6.0, 7.0, 8.0)
    var b = SIMD[DType.float32, 8](2.0)  # 广播：全部设为 2.0

    # 向量运算（一条指令处理全部 8 个元素）
    var c = a + b   # [3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
    var d = a * b   # [2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0]
    var e = a / b   # [0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]

    print(c)

    # 访问单个元素
    print(c[0])    # 3.0
    print(c[7])    # 10.0

    # reduce_add: 水平求和
    print(c.reduce_add())  # 63.0 = 3+4+5+6+7+8+9+10

DType 枚举：支持的数据类型

# DType 支持的元素类型
alias F16Vec = SIMD[DType.float16, 16]  # 16× FP16 = 256 bit
alias F32Vec = SIMD[DType.float32, 8]   # 8× FP32 = 256 bit
alias F64Vec = SIMD[DType.float64, 4]   # 4× FP64 = 256 bit
alias I32Vec = SIMD[DType.int32, 8]     # 8× Int32 = 256 bit
alias U8Vec  = SIMD[DType.uint8, 32]    # 32× UInt8 = 256 bit（图像处理）
alias BF16Vec = SIMD[DType.bfloat16, 16] # 16× BF16（AI训练）

simd_width：自动检测最优向量宽度

跨平台自适应代码

不同 CPU 的 SIMD 宽度不同（128/256/512 bit）。手动写死向量宽度会导致代码无法最优利用当前硬件。Mojo 提供 simdwidthof[DType]() 在编译时获取当前硬件对指定类型最优的 SIMD 宽度。

from sys.info import simdwidthof
from algorithm import vectorize

fn adaptive_add(
    a: List[Float32],
    b: List[Float32],
    inout result: List[Float32]
):
    # 编译期自动确定最优向量宽度
    alias SIMD_WIDTH = simdwidthof[DType.float32]()
    # SSE:  SIMD_WIDTH = 4
    # AVX2: SIMD_WIDTH = 8
    # AVX512: SIMD_WIDTH = 16

    let n = len(a)

    @parameter
    fn vectorized_add[simd_width: Int](i: Int):
        # 加载 simd_width 个 Float32 到向量寄存器
        var va = SIMD[DType.float32, simd_width].load(a, i)
        var vb = SIMD[DType.float32, simd_width].load(b, i)
        (va + vb).store(result, i)

    vectorize[vectorized_add, SIMD_WIDTH](n)

fma：融合乘加指令

FMA 为什么重要？

FMA（Fused Multiply-Add，融合乘加） 是现代 CPU 的关键指令：result = a * b + c 在一个时钟周期内完成，精度优于分两步计算（中间结果不舍入）。神经网络推理中 95% 以上的计算是矩阵乘法，而矩阵乘法的核心就是大量的乘加操作——FMA 是其中的性能基础。

from math import fma

fn fma_demo():
    # 标量 FMA
    var result = fma(3.0, 4.0, 5.0)  # 3.0 * 4.0 + 5.0 = 17.0
    print(result)

    # 向量 FMA（SIMD）
    var va = SIMD[DType.float32, 8](1.0, 2.0, 3.0, 4.0,
                                     5.0, 6.0, 7.0, 8.0)
    var vb = SIMD[DType.float32, 8](2.0)
    var vc = SIMD[DType.float32, 8](10.0)

    # a[i] * b[i] + c[i]，一条 VFMADD 指令处理 8 个元素
    var vd = fma(va, vb, vc)
    print(vd)  # [12.0, 14.0, 16.0, 18.0, 20.0, 22.0, 24.0, 26.0]

实战：向量化矩阵乘法

from sys.info import simdwidthof
from algorithm import vectorize
from math import fma
from time import now

alias F32 = DType.float32
alias SIMD_W = simdwidthof[F32]()

fn matmul_vectorized(
    C: UnsafePointer[Float32],
    A: UnsafePointer[Float32],
    B: UnsafePointer[Float32],
    M: Int, N: Int, K: Int
):
    """C[M×N] = A[M×K] @ B[K×N]"""
    for m in range(M):
        for k in range(K):
            var a_val = SIMD[F32, SIMD_W](A[m * K + k])

            @parameter
            fn inner[sw: Int](n: Int):
                var b_vec = SIMD[F32, sw].load(B + k * N + n)
                var c_vec = SIMD[F32, sw].load(C + m * N + n)
                # FMA: C[m,n:n+sw] += A[m,k] * B[k,n:n+sw]
                fma(a_val, b_vec, c_vec).store(C + m * N + n)

            vectorize[inner, SIMD_W](N)

fn benchmark():
    alias M = 512; alias N = 512; alias K = 512
    var A = UnsafePointer[Float32].alloc(M * K)
    var B = UnsafePointer[Float32].alloc(K * N)
    var C = UnsafePointer[Float32].alloc(M * N)

    let start = now()
    matmul_vectorized(C, A, B, M, N, K)
    let elapsed = (now() - start) / 1_000_000

    let flops = 2 * M * N * K
    let gflops = Float64(flops) / (Float64(elapsed) * 1e6)
    print("耗时：", elapsed, "ms")
    print("GFLOPS：", gflops)

    A.free(); B.free(); C.free()

def main():
    benchmark()