【深度学习】Pytorch gpu加速原理探究

文章目录

Pytoch实现gpu加速的实质是什么

本质上，PyTorch 的"GPU加速"就是：把张量计算从 CPU 指令切到 CUDA 内核，在 GPU 上并行执行。

关键点只有几个：

tensor/model.to('cuda') 后，算子会走 CUDA 后端（不是 CPU 后端）。
底层调用的是高性能 GPU 库/内核（如 cuBLAS、cuDNN，或自定义 CUDA kernel）。
前向和反向传播都在 GPU 上跑（autograd 的 backward 也会调用 GPU 内核）。
真正速度提升来自"大规模并行 + 高带宽"，不是框架本身"变快"。

所以加速效果取决于：

计算量是否足够大（小张量反而可能更慢）
数据是否长期留在 GPU（频繁 CPU↔GPU 拷贝会抵消收益）
算子是否有 CUDA 实现（否则会回退到 CPU）。

什么是算子

例如在 PyTorch 里这些都是算子：

矩阵乘法 matmul
卷积 conv2d
激活函数 relu
加法 add
归一化 batch_norm

模型前向/反向本质上就是大量算子按计算图顺序执行。

"某算子支持 CUDA" 的意思是：这个算子有对应的 GPU 实现内核，可以在显卡上算。

matmul的cuda版本实现是什么样的

在 PyTorch 里，matmul 的 CUDA 实现通常不是"手写一个简单 kernel"，而是走 cuBLAS/cuBLASLt 的

GEMM（矩阵乘法）：

C = A @ B -> 调用类似 cublasGemmEx / cublasLtMatmul

你可以把它理解成三层：

框架层（PyTorch）
把 torch.matmul 按维度规则展开（2D、batched、broadcast），再分发到 CUDA 后端。
库层（cuBLAS/cuBLASLt）
选择最优算法（tile 大小、是否用 Tensor Core、流水线策略等）。
内核层（CUDA kernel）
核心思路是分块（tiling）：每个 thread block 算 C 的一个小块，循环读取 A/B 的子块到 shared memory，累加后写回。

一个教学版的简化 kernel 大概这样：

// C $M,N$ = A $M,K$ * B $K,N$

#define TILE 16

global void matmul_kernel(const float* A, const float* B, float* C,

int M, int N, int K) {
shared float As $TILE$ $TILE$ ;
shared float Bs $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 < (K + TILE - 1) / TILE; ++t) {
      int a_col = t * TILE + threadIdx.x;
      int b_row = t * TILE + threadIdx.y;

      As[threadIdx.y][threadIdx.x] =
          (row < M && a_col < K) ? A[row * K + a_col] : 0.0f;
      Bs[threadIdx.y][threadIdx.x] =
          (b_row < K && col < N) ? B[b_row * N + col] : 0.0f;

      __syncthreads();

      #pragma unroll
      for (int i = 0; i < TILE; ++i) sum += As[threadIdx.y][i] * Bs[i][threadIdx.x];

      __syncthreads();
  }

  if (row < M && col < N) C[row * N + col] = sum;

}

真实的 cuBLAS 实现会比这个复杂很多：会用寄存器分块、warp 级 MMA、Tensor Core（如

mma.sync）、双缓冲流水线等来榨干吞吐。

cublas是什么

cuBLAS 是 NVIDIA 提供的 CUDA 版 BLAS 库（Basic Linear Algebra Subprograms）。

它的作用是给 GPU 提供高性能线性代数算子，比如向量运算、矩阵向量乘、矩阵乘（GEMM）。

PyTorch、TensorFlow 这类框架做 matmul 时，很多情况下就是在调用 cuBLAS（或 cuBLASLt）而不是自

己手写矩阵乘 kernel。

可以把它类比为 CPU 里的 MKL/OpenBLAS，只是目标是 NVIDIA GPU，并且会利用 Tensor Core 等硬件特

性做加速。