第7板块·第1节：通用算子分类与设计模式

学习目标

学完本节你将能够：

• 对常见 GPU 算子进行分类，理解每类算子的计算特征和优化重点
• 掌握通用 Kernel 设计模式：逐元素、Grid‑Stride Loop、分块、共享内存协作、Warp 合作
• 根据算子类型快速判断最优的并行策略和内存访问模式
• 了解 CUTLASS 模板库的定位与核心抽象，为后续深入学习做铺垫
• 能够为新算子选择合适的实现模板

---

1. 通用算子的分类

在深度学习和科学计算中，GPU 算子种类繁多，但可以根据计算模式归纳为几大类：

算子类别	典型例子	计算特征	优化重点
逐元素（Element‑wise）	ReLU、加法、缩放	每个输出独立，无数据依赖	内存带宽利用、向量化加载
归约（Reduction）	求和、最大值、范数	多对一，需要线程间通信	共享内存、Warp 洗牌、多级归约
扫描（Scan/Prefix Sum）	前缀和、累积分布	存在顺序依赖，但可并行化	Blelloch 算法、Warp 扫描
矩阵乘/卷积（GEMM/Conv）	矩阵乘法、卷积	高计算密度，数据复用率高	分块、共享内存、Tensor Core
排序/搜索（Sort/Search）	排序、TopK、哈希	数据依赖复杂，通信密集	原子操作、协作组、专用库
稀疏算子（Sparse）	SpMV、稀疏注意力	不规则数据访问，负载不均衡	数据格式优化、负载均衡

---

2. 通用 Kernel 设计模式

2.1 逐元素模式（One Thread per Element）

最简单直观的模式，每个线程处理一个输出元素。

__global__ void elementwise(float *out, const float *in, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        out[idx] = in[idx] * 2.0f;
    }
}

适用：计算独立、N 不太大、操作简单。

优化：使用 float4 向量化、Grid‑Stride Loop 处理大数组。

2.2 Grid‑Stride Loop 模式

每个线程处理多个元素，以固定步长遍历数据。

__global__ void gridStride(float *out, const float *in, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < N; i += stride) {
        out[i] = in[i] * 2.0f;
    }
}

优点 ：减少线程启动数量，提高单线程计算密度，适应不同规模 GPU。

适用：大规模逐元素、内存密集型算子。

2.3 分块与共享内存协作模式

将数据分块加载到共享内存，在块内复用数据，减少全局内存访问。

#define BLOCK_SIZE 256
__global__ void tiledKernel(float *out, const float *in, int N) {
    __shared__ float tile[BLOCK_SIZE];
    int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x;
    tile[threadIdx.x] = (idx < N) ? in[idx] : 0.0f;
    __syncthreads();

    // 在共享内存上计算
    out[idx] = tile[threadIdx.x] * 2.0f;
}

适用：需要数据复用的算子，如矩阵转置、GEMM、卷积。

2.4 Warp 合作模式

利用 Warp 洗牌或 Warp 级归约，避免共享内存和同步开销。

__global__ void warpReduce(float *in, float *out, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float val = (idx < N) ? in[idx] : 0.0f;
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    if ((threadIdx.x % 32) == 0) out[blockIdx.x] = val;
}

适用：归约、扫描、小规模数据交换。

2.5 多级归约模式

结合共享内存和 Warp 洗牌，先 Warp 内归约，再跨 Warp 归约，最后跨 Block 归约。适合大规模归约。

__global__ void fullReduction(float *in, float *out, int N) {
    __shared__ float warpResults[32];
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float val = (idx < N) ? in[idx] : 0.0f;

    // Warp 内归约
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }

    int warpId = threadIdx.x / 32;
    int laneId = threadIdx.x % 32;
    if (laneId == 0) warpResults[warpId] = val;
    __syncthreads();

    // 第一个 Warp 归约 warpResults
    if (warpId == 0) {
        val = (laneId < blockDim.x / 32) ? warpResults[laneId] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_down_sync(0xffffffff, val, offset);
        }
        if (laneId == 0) out[blockIdx.x] = val;
    }
}

---

3. 算子类型与设计模式对应关系

算子类别	推荐设计模式	关键优化点
逐元素	Grid‑Stride Loop + 向量化	合并访问、float4
归约	Warp 洗牌 + 共享内存多级归约	减少同步、多累加器
扫描	Blelloch 或 Warp 扫描 + 共享内存	平衡同步次数
GEMM	分块共享内存 + 寄存器分块 + Tensor Core	数据复用、Bank Conflict
卷积	im2col + GEMM 或共享内存直接卷积	数据复用、内存布局
稀疏算子	原子操作 + 负载均衡	减少竞争、数据格式

---

4. CUTLASS 模板库简介

CUTLASS 是 NVIDIA 开源的 CUDA C++ 模板库，专门用于实现高性能 GEMM 及相关算子。它的核心思想是将 GEMM 分解为可配置的 tile 和 warp 级别操作，利用模板参数在编译时确定分块大小、数据类型、布局等，从而生成高度优化的代码。

4.1 CUTLASS 的抽象层次

CUTLASS 将 GEMM 分为几个层次：

1. Tile 级别：将大矩阵划分为多个小 tile，每个 threadblock 处理一个 tile。
2. Warp 级别：每个 warp 处理 tile 内的一个子块。
3. 线程级别：每个线程计算一个小的矩阵片段（通常 4×4 或 8×8），使用寄存器分块。

通过模板参数可以指定：

• 数据类型（FP32、FP16、BF16、INT8 等）
• 矩阵布局（行主序、列主序）
• Tile 大小
• 是否使用 Tensor Core
• 流水线阶段数

4.2 CUTLASS 的优势

• 性能接近手写汇编：通过模板和编译时优化，达到 cuBLAS 级别的性能。
• 高度可定制：用户可以轻松修改 tile 大小、数据类型、混合精度策略。
• 可读性好：相比手写 SASS 或复杂宏，CUTLASS 代码更结构化，易于理解和维护。
• 支持多种算子：除了 GEMM，还支持卷积、注意力等。

4.3 简单示例（CUTLASS GEMM）

#include <cutlass/cutlass.h>
#include <cutlass/gemm/device/gemm.h>

// 定义 GEMM 类型
using Gemm = cutlass::gemm::device::Gemm<
    float,              // A 元素类型
    cutlass::layout::RowMajor,   // A 布局
    float,              // B 元素类型
    cutlass::layout::RowMajor,   // B 布局
    float,              // C 元素类型
    cutlass::layout::RowMajor    // C 布局
>;

int main() {
    // 创建 GEMM 对象
    Gemm gemm;
    // 设置参数并运行
    // ...
    return 0;
}

在本板块后续章节中，我们将深入 CUTLASS 的内部结构和优化策略。

---

5. 课后练习

练习1：算子分类

列举你在深度学习框架中见过的算子，并按照本节分类标准归类，说明每类算子的优化重点。

练习2：实现 Grid‑Stride 逐元素算子

使用 Grid‑Stride Loop 实现一个向量平方和算子，与 one‑thread‑per‑element 版本对比性能。

练习3：多级归约

实现一个完整的多级归约（Warp → Block → Grid），处理长度为 1 << 24 的数组，使用 ncu 分析性能。

练习4：了解 CUTLASS

访问 CUTLASS GitHub 仓库，浏览 examples 目录，编译一个简单的 GEMM 示例并运行，记录性能。

练习5：设计模式选择

给定一个新算子（例如 batch normalization），分析其计算特征，选择合适的设计模式，并写出 Kernel 框架。

---

6. 下一步

下一节将深入 CUTLASS 的核心抽象与数据布局，你将学习：

• CUTLASS 的 tile 和 fragment 概念
• 矩阵布局（RowMajor/ColumnMajor/TensorOp 布局）对性能的影响
• CUTLASS 的模板参数如何影响代码生成
• 如何根据硬件能力选择合适的 GEMM 配置