核心架构#

• 多核设计：现代CPU通常拥有2-64个处理核心，每个核心可以独立执行指令
• 缓存层次结构：
  • L1缓存：最快、容量最小(几十KB)，通常分为指令缓存和数据缓存
  • L2缓存：中等速度和容量(几百KB到几MB)
  • L3缓存：较大容量(几MB到几十MB)，在多核间共享
• 指令流水线：将指令执行拆分为多个阶段(取指、译码、执行、访存、写回等)

数据处理能力#

• SIMD指令集：单指令多数据并行处理能力
  • x86架构：SSE、AVX、AVX2、AVX-512等
  • ARM架构：NEON、SVE等
• 内存访问：支持乱序执行和预取机制，减少内存访问延迟

CPU与矩阵计算#

CPU执行矩阵计算时面临以下挑战：

• 内存带宽限制：大型矩阵数据传输可能成为瓶颈
• 缓存局部性：矩阵计算需要优化内存访问模式以利用缓存
• 并行度有限：与GPU相比，CPU的SIMD单元和核心数量较少

CPU核心优势#

• 顺序执行效率：

  • 高单线程性能和复杂的分支预测
  • 大缓存设计优化顺序代码执行
  • 较低的指令延迟，适合串行任务

• 条件控制流处理：

  • 高效的分支预测器，减少条件判断开销
  • 支持复杂的控制逻辑和不规则代码路径
  • 适合具有大量条件判断和异常处理的算法

• 混合工作负载处理：

  • 平衡的设计适合同时处理计算、IO和系统任务
  • 灵活性强，可处理多样化任务而非单一计算模式

这些特性使CPU在需要复杂决策逻辑的AI推理任务、迭代算法开发和混合工作负载中保持优势。

# 利用NumPy在CPU上高效执行矩阵乘法
import numpy as np
import time

# 创建两个大矩阵
A = np.random.rand(1000, 1000)
B = np.random.rand(1000, 1000)

start = time.time()
C = np.matmul(A, B)  # NumPy自动使用优化的BLAS库
end = time.time()

print(f"CPU矩阵乘法耗时: {end - start:.4f}秒")

优化的CPU库(如Intel MKL、OpenBLAS)通过充分利用SIMD指令和多核心，仍能高效处理中小规模矩阵计算。

CPU 的执行流程如下图所示 其中访存步骤花费的时间相对是最多的

案例：如何在 CPU 上高效计算矩阵乘法?#

矩阵乘法是张量计算的基础操作，对于两个矩阵 $A \in \mathbb{R}^{M \times K}$ 和 $B \in \mathbb{R}^{K \times N}$，其乘积 $C = A \times B$ 的计算公式为

$C_{i,j} = \sum_{k=0}^{K-1} A_{i,k} \times B_{k,j}$

直接实现与性能瓶颈#

最简单的实现方式是使用三重循环

for (int i = 0; i < M; i++)
    for (int j = 0; j < N; j++) {
        C[i][j] = 0;
        for (int k = 0; k < K; k++)
            C[i][j] += A[i][k] * B[k][j];
    }

但是这种实现存在严重的性能问题

• 缓存不友好：对矩阵B的访问是按列进行的，而在内存中数据是按行存储的
• 频繁访存：每次计算都需要从内存加载数据，导致CPU大部分时间在等待数据
• 访存次数：对于每个元素C[i][j]的计算，需要读取K个A的元素和K个B的元素

分块矩阵乘法优化#

分块矩阵乘法的核心思想是减少访存操作，提高缓存利用率：

// 分块大小根据缓存大小选择
int BM = 64, BN = 64, BK = 64;

// 初始化C矩阵为0
for (int i = 0; i < M; i++)
    for (int j = 0; j < N; j++) 
        C[i][j] = 0;

// 分块矩阵乘法
for (int i = 0; i < M; i += BM)
    for (int j = 0; j < N; j += BN)
        for (int k = 0; k < K; k += BK) {
            // 计算子矩阵乘法并累加结果
            for (int i1 = i; i1 < min(i+BM, M); i1++)
                for (int j1 = j; j1 < min(j+BN, N); j1++) {
                    float temp = 0;  // 局部累加器
                    for (int k1 = k; k1 < min(k+BK, K); k1++)
                        temp += A[i1][k1] * B[k1][j1];
                    C[i1][j1] += temp;  // 累加到C
                }
        }

性能提升原理：

1. 减少访存次数：一次将一个块加载到缓存中，多次复用这些数据
2. 改善局部性：小块矩阵能完全放入L1/L2缓存，大幅减少缓存未命中
3. 计算访存比优化：每个缓存块加载一次后可执行多次计算操作

通过优化，分块矩阵乘法可将访存次数从 $O(MNK)$ 降低到 $O(MNK/B)$，其中B是块大小。当B较大时，访存开销显著降低，性能可提升5-10倍。

除分块外，进一步优化还包括:

• 循环重排序
• 向量化指令(SIMD)
• 多线程并行化
• 内存对齐和预取技术

GPU 执行模型#

单指令多线程 (SIMT) 是GPU并行处理的核心执行模型

• 原理：同一指令被多个线程同时执行，但每个线程拥有独立的程序计数器和寄存器状态
• 组织结构：
  • 线程被组织为线程束（warp/wavefront），典型大小为32/64线程
  • 同一线程束 (warp) 内的线程执行相同指令路径时性能最优
  • 多个线程束组成线程块（block），多个线程块组成网格（grid）
• 分支处理：
  • 当线程束内出现分支时，不同分支路径被序列化执行
  • 执行某分支时，不走该路径的线程被暂时屏蔽
  • 分支散布（divergence）会显著降低性能

• 线程层次结构：
  • 线程束(Warp/Wavefront)：GPU执行的基本单位，通常为32线程(NVIDIA)或64线程(AMD)
  • 线程块(Block)：共享本地内存的线程集合，可包含多个warp
  • 网格(Grid)：整个GPU计算任务，由多个块组成

$\text{Thread} \subset \text{Warp} \subset \text{Block} \subset \text{Grid}$

这种执行模型使GPU能高效处理规整的并行计算任务，如矩阵运算，但在分支条件复杂时性能下降。

GPU 访存延迟#

如何在 GPU 上高效计算矩阵乘法？#

在GPU上高效实现矩阵乘法涉及多层次的优化策略，主要围绕着如何充分利用GPU的并行架构和内存层次结构。

__global__ void matrixMul(float *A, float *B, float *C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if(row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

__global__ void matrixMulShared(float *A, float *B, float *C, int M, int N, int K) {
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
    
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    
    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;
    
    float sum = 0.0f;
    
    for (int t = 0; t < (K-1)/TILE_SIZE + 1; t++) {
        // 协作加载A和B的tile到共享内存
        if(row < M && t*TILE_SIZE+tx < K)
            As[ty][tx] = A[row*K + t*TILE_SIZE+tx];
        else
            As[ty][tx] = 0.0f;
            
        if(t*TILE_SIZE+ty < K && col < N)
            Bs[ty][tx] = B[(t*TILE_SIZE+ty)*N + col];
        else
            Bs[ty][tx] = 0.0f;
            
        __syncthreads();
        
        // 计算当前tile的乘积
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += As[ty][k] * Bs[k][tx];
        }
        __syncthreads();
    }
    
    if(row < M && col < N)
        C[row*N + col] = sum;
}

# 使用cuBLAS的矩阵乘法示例
import torch

# 创建GPU上的张量
A = torch.rand(1024, 1024, device="cuda")
B = torch.rand(1024, 1024, device="cuda")

# PyTorch会自动调用cuBLAS
C = torch.matmul(A, B)

算术表达式化简#

1. 通过代数运算等价变换计算图，如

   • $a * 0 \rightarrow 0$
   • $a * 1 \rightarrow a$
   • $a + 0 \rightarrow a$
   • $\log(\exp(x) / y) \rightarrow x - \log y$

2. 利用交换律、结合律等规律调整图中的算子执行顺序，或者删除不必要的算子

   • $(a \cdot b) \odot d + (a \cdot b) \odot c \rightarrow (a \cdot b) \odot (c + d)$
   • e.g. $AB \odot C + AB \odot D \rightarrow AB \odot (C + D)$，其中 $\odot$ 是逐元素乘法，$AB$ 是矩阵乘法

公共子表达式消除#

目的是通过找到程序中等价的计算表达式，然后通过复用结果的方式消除其他冗余表达式的计算

这里的数字只是算子和参数的标号，不是加7减6之类的含义

死代码消除#

移除对程序执行结果没有任何影响的代码，例如一些在训练阶段会使用但在推理时不会执行的计算，可以被消除 (Dropout 操作)

常量折叠#

折叠多余的常量计算，例如 a = 320 * 200 * 32 直接变成 a = 2_048_000。这个道理非常的显然，但是常量折叠具有较好的优化作用，因为在模型的修改调试阶段，用户可能会因为调试而编写了很多非折叠的常量，如果不折叠，这在后续的训练优化中会很消耗性能。此外，模型训练阶段时，模型的参数是变化的，而当训练结束后，在推理阶段，模型的参数可以视作一个个常量，此时可以使用常量折叠进行优化。

一个具体的例子是 BatchNorm。BatchNorm操作在训练阶段的计算公式为： $y = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

其中:

• $x$ 是输入特征
• $\mu$ 是批次均值
• $\sigma^2$ 是批次方差
• $\gamma, \beta$ 是可学习的缩放和偏移参数
• $\epsilon$ 是防止除零的小常数

在推理阶段，$\mu$ 和 $\sigma^2$ 通常使用训练过程中累积的移动平均值，且这些值与 $\gamma$ 和 $\beta$ 都是固定的。此时可以通过常量折叠优化：

# 原始BatchNorm推理计算
def batch_norm_inference(x, gamma, beta, mean, var, epsilon=1e-5):
    return gamma * (x - mean) / torch.sqrt(var + epsilon) + beta

# 常量折叠后的计算
def batch_norm_folded(x, scale, shift):
    return x * scale + shift

# 参数折叠过程
scale = gamma / torch.sqrt(var + epsilon)
shift = beta - scale * mean

# 将原始BatchNorm替换为简单的缩放和平移操作

这种优化将BatchNorm从包含除法的复杂操作简化为单一的缩放和平移（乘加操作），减少了计算量并提高了推理效率。在实际部署中，常量折叠可以将BatchNorm层与相邻的卷积层进一步融合，完全消除BatchNorm的额外计算开销。

在高效的推理引擎实现中，这种折叠通常是自动完成的，显著提升了模型的执行速度，尤其在边缘设备上更为明显。

算子融合#

将向量化的多个算子操作合并成一个算子操作。这种优化的想法是，CPU 内核的启动是有时间开销的，这样做可以减少内存的读取，提高计算密度。例如，可以将 c = mul(a, b) 和 e = add(c, d) 合并成 e = mla(a, b, d)，从而减少一次 c 写入的访存和一次 c 读取的访存。

另外一个具体的例子是实现矩阵乘法的自动融合，例如同时有两个张量 x 和 y 都需要有与张量 z 相乘计算后的结果，可以把 x 和 y 进行拼接再与 z 相乘，从而减少访存的操作和内核启动时间

布局转换#

Tensor 是一个多维阵列，但是在存储中的排列存储方式确是“向量化”的形式，因此如何布局排布对计算 cache 的读取是具有影响的。假设有一个形状为 (2, 3) (行编号, 列编号) 的矩阵 $X$，通常排布有

• 行优先：排布成 $[X_{1, 1:3}, X_{2, 1:3}]$ 存储
• 列优先：排布成 $[X_{1:2, 1}, X_{1:2, 2}, X_{1:2, 3}]$ 存储

在深度学习中，图像张量的形状有两种常用的排布方式

• (B, C, H, W) 型：各个张量维度依次表示 (批次大小, 通道数, 高度, 宽度)，这是输入卷积层的特征图的形状。同一个通道的数据值连续排布，更适合需要对每个通道单独运算的操作，如 MaxPooling，适合 GPU、大带宽的计算模式
• (B, H, W, C) 型：各个张量维度依次表示 (批次大小, 高度, 宽度, 通道数)。不同通道中同一位置元素顺序存储，更适合需要跨通道计算的操作如 $1\times 1$卷积（注意$1\times 1$卷积可以用 NHWC 形式下存储的张量和线性层实现，这样比 NCHW 加卷积层的模式更高效）。这种计算模式更适合多核 CPU 执行，计算控制灵活且复杂

内存分配#

在深度学习训练和推理过程中，内存分配与管理是影响性能和效率的关键因素。内存复用技术可以显著降低内存开销，尤其是在处理大型模型时。

# 使用torch.cuda.memory_stats()查看当前GPU内存使用情况
import torch

# 在不优化的情况下执行操作
def unoptimized():
    x = torch.randn(10000, 10000, device='cuda')
    y = torch.randn(10000, 10000, device='cuda')
    z = x @ y  # 矩阵乘法
    return z

# 使用内存前后对比
torch.cuda.empty_cache()
before = torch.cuda.memory_allocated()
result = unoptimized()
after = torch.cuda.memory_allocated()
print(f"内存增长: {(after - before) / 1024**2:.2f} MB")

这里没有进行优化，导致内存增长为 390.12 MB

内存复用的基本原理

在计算图执行过程中，不同操作产生的中间结果通常只在特定阶段有用。一旦这些中间结果被消费后，其占用的内存空间可以被释放并重新用于存储其他中间结果。

• 原地操作 (Inplace Operation): 如果一块内存不再需要，且下一个操作是 element-wise 的，就可以做原地操作覆盖内存，而无需新内存。例如，PyTorch 的激活函数可以设置采用 torch.nn.ReLU(inplace=True)

• 内存共享 (Memory Sharing): 两个数据共享同一块内存空间，其中有一个数据参与计算后不再需要，后一个数据可以覆盖前一个数据，只要这两个数据的生命周期无重叠即可（不过这对反向传播是有影响的）。

PyTorch中提供了多种内存优化机制

1. 梯度检查点(Checkpoint)

对于大型模型，可以使用梯度检查点技术在反向传播时重新计算中间激活值，而不是存储所有激活值

import torch.utils.checkpoint as checkpoint

class LargeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 定义一个包含多个层的大型模型
        self.layers = torch.nn.ModuleList([
            torch.nn.Linear(1024, 1024) for _ in range(10)
        ])
        
    def forward(self, x):
        for layer in self.layers:
            # 使用checkpoint包装层的前向传播
            x = checkpoint.checkpoint(layer, x)
        return x
        
# 这种方式会在反向传播时重新计算中间结果，降低内存使用但增加计算量

2. 释放未使用的Tensor

# 明确释放不再需要的张量
x = torch.randn(10000, 10000, device='cuda')
y = torch.randn(10000, 10000, device='cuda')
z = x @ y
del x, y  # 显式释放不再需要的大张量
torch.cuda.empty_cache()  # 释放缓存的未使用内存

3. 使用inplace操作

当可能时，使用原地操作可以避免额外的内存分配：

# 非原地操作
x = torch.ones(1000, 1000)
x = x + 1  # 创建新张量

# 原地操作
x = torch.ones(1000, 1000)
x.add_(1)  # 直接修改x，不创建新张量

4. 自动垃圾回收优化

PyTorch 2.0引入了更智能的内存管理机制，包括通过编译优化内存使用：

# 使用torch.compile优化内存使用
@torch.compile
def optimized_function(x, y):
    intermediate1 = x * 2
    intermediate2 = y * 3
    result = intermediate1 @ intermediate2
    return result

# 编译后的函数会优化内存使用模式

5. 流水线并行和模型分片

对于特别大的模型，可以使用PyTorch的分布式训练工具实现模型并行化：

# 使用torch.distributed进行模型并行示例
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# 初始化分布式环境
dist.init_process_group(backend="nccl")

# 将模型的不同部分分配到不同设备
model_part1 = Model1().to(f'cuda:{dist.get_rank() % 2}')
model_part2 = Model2().to(f'cuda:{(dist.get_rank()+1) % 2}')

# 每个设备只存储部分模型参数

通过这些技术，深度学习框架可以在有限内存条件下处理更大的模型和数据集，提高训练和推理效率。

循环展开#

循环展开是一种通过减少循环控制开销和提高指令级并行性来优化循环执行的技术。

基本原理：将循环体内的代码复制多次，减少循环迭代次数，同时增加每次迭代中的计算量。

// 原始循环
for (int i = 0; i < N; i++) {
    sum += array[i];
}

// 循环展开(展开因子为4)
for (int i = 0; i < N; i += 4) {
    sum += array[i];
    sum += array[i+1];
    sum += array[i+2];
    sum += array[i+3];
}

优化效果：

• 减少循环控制开销(分支预测、条件检查)
• 提高指令级并行性
• 增加寄存器利用率
• 减少内存操作的延迟影响

在深度学习计算如卷积或矩阵乘法中，循环展开通常与向量化指令(如SIMD)结合使用，显著提高计算密度。

循环分块#

循环分块（Loop Tiling）通过改变数据访问模式来提高缓存效率，对于矩阵和张量运算尤为重要。

基本原理：将大循环划分为小块，使每块数据能够完全放入CPU缓存中。

// 原始矩阵乘法
for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++) {
        float sum = 0;
        for (int k = 0; k < N; k++)
            sum += A[i][k] * B[k][j];
        C[i][j] = sum;
    }

// 分块优化后的矩阵乘法
for (int i0 = 0; i0 < N; i0 += B)
    for (int j0 = 0; j0 < N; j0 += B)
        for (int k0 = 0; k0 < N; k0 += B)
            for (int i = i0; i < min(i0+B, N); i++)
                for (int j = j0; j < min(j0+B, N); j++) {
                    float sum = (k0 == 0) ? 0 : C[i][j];
                    for (int k = k0; k < min(k0+B, N); k++)
                        sum += A[i][k] * B[k][j];
                    C[i][j] = sum;
                }

优化效果：

• 提高缓存命中率
• 减少内存带宽需求
• 增加数据局部性
• 降低TLB缺失率

分块策略在深度学习编译器中广泛应用，特别是在GPU计算中，将数据分块到不同层次的内存(共享内存、寄存器)可以显著提升性能。

循环融合#

循环融合是将多个循环合并成一个循环的优化技术，可减少循环开销并提高数据局部性。

基本原理：当多个循环具有相同的迭代范围且无数据依赖时，可以将它们融合为一个循环。

// 融合前
for (int i = 0; i < N; i++)
    A[i] = B[i] * scalar1;
    
for (int i = 0; i < N; i++)
    C[i] = A[i] + D[i];

// 融合后
for (int i = 0; i < N; i++) {
    A[i] = B[i] * scalar1;
    C[i] = A[i] + D[i];
}

优化效果：

• 减少循环控制开销
• 提高指令缓存效率
• 减少数据在不同循环间的移动
• 提高数据局部性和缓存命中率

在深度学习计算中，循环融合通常用于合并element-wise操作，如将激活函数计算与前一层的线性操作融合，减少中间结果的内存访问。许多现代AI编译器（如TVM、XLA）会自动执行这类融合优化。

循环融合与算子融合密切相关，但更专注于底层循环实现而非图层面的优化。在实际应用中，两者通常结合使用，先在计算图层面融合算子，再在代码生成阶段优化具体循环实现。

循环拆分#

循环拆分#

循环拆分（Loop Splitting）是将单个循环拆分为多个连续执行的循环的优化技术，有助于提高执行效率和并行性。主要是针对含有条件分支的语句进行的优化方法。

基本原理：将循环按照某种条件分成多个部分，每部分处理原循环的不同情况或执行不同的优化策略。

// 原始循环（包含条件判断）
for (int i = 0; i < N; i++) {
    if (i < M)
        A[i] = B[i] * 2;
    else
        A[i] = C[i] + D[i];
}

// 拆分后
for (int i = 0; i < M; i++) {
    A[i] = B[i] * 2;  // 第一部分
}

for (int i = M; i < N; i++) {
    A[i] = C[i] + D[i];  // 第二部分
}

优化效果：

• 消除循环内的条件分支，减少分支预测失误
• 简化循环体，使编译器能应用更多优化
• 为不同循环部分应用不同的优化策略
• 提高向量化和并行化机会

在深度学习计算中，循环拆分常用于处理边界条件、对齐问题或特殊情况，例如将卷积操作中填充(padding)区域和非填充区域的计算分开处理，使主计算循环能够被更好地优化。这种技术也可以与循环展开、向量化等其他优化相结合，进一步提升性能。

参考资料：

• 微软亚洲研究院的 AI System 课程
• TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
• CUDA C++ Programming Guide
• The LLVM Compiler Infrastructure
• Deep Learning Compiler: A Comprehensive Survey
• Optimizing CNN Model Inference on CPUs
• PyTorch Internals