当前位置：首页 > news >正文

51c~CUDA~合集1

news 2025/5/30 12:59:56

我自己的原文哦~ https://blog.51cto.com/whaosoft/13949692

一、CUDA性能简易优化

这里是关于CUDA性能优化的基础教程，主要介绍了GPU的基础知识，帮助理解如何设计和优化深度学习中的层，使其能充分利用GPU性能。

CUDA性能优化简单教程，本篇介绍性能优化背景。

想知道实际中如何优化特定的层，或者某一层怎么设计才可能充分利用GPU，我们需要了解一些GPU的基础知识。

以下教程主要来源自NVIDIA官方：GPU Performance Background User's Guide[1]，主要是讲解深度学习中，网络中一些operation在GPU中是如何执行的，已经一些和性能相关的细节注意点。

具体内容如下：

GPU的基本结构 (GPU Architecture Fundamentals[2])
OP（operation）是如何被拆分设计为并行计算的 (GPU Execution Model[3])
如何通过算术强度（arithmetic intensity）估算性能限制 (Understanding Performance[4])
深度学习操作的大概分类及其各自的性能限制(DNN Operation Categories[5])

1. GPU的基础结构

GPU是一种高度并行的处理器架构，由processing elements and a memory hierarchy组成，和我们经常提到的SM和显存有直接关系。

在高层次上，NVIDIA® GPU由多个流多处理器（SM）组成，具有片上的L2缓存和高带宽DRAM。SM执行算术和其他指令；通过L2缓存从DRAM访问数据和代码。例如，NVIDIA A100 GPU包含108个SM，40 MB L2缓存以及来自80 GB HBM2内存的高达2039 GB/s的带宽。

GPU架构

每个SM都有自己的instruction schedulers和各种instruction execution pipelines。在现代神经网络中，Multiply-add是最常见的操作，作为全连接和卷积层的核心块操作，两者都可以视为向量点积（vector dot-products.）的集合。

下表显示了NVIDIA Ampere GPU架构上单个SM的每个时钟的乘加操作的各种数据类型。每个乘加操作包括两个操作，因此可以将表中的吞吐量乘以2，以获得FLOP counts per clock。

要获取GPU的FLOPS具体值，我们可以将这些乘以SM的数量和SM时钟速率。例如，带有108个SM和1.41 GHz时钟速率的A100 GPU具有156 TF32 TFLOPS和312 FP16 TFLOPS的峰值密集吞吐量（实际能达到的吞吐量取决于后续讨论的多个因素，这只是理论值）。

Multiply-add operations per clock per SM

另外，FP16操作可以在Tensor Cores或NVIDIA CUDA®核心中执行。NVIDIA Turing™架构可以在Tensor Cores或CUDA核心中执行INT8操作。Tensor Cores在NVIDIA Volta™ GPU架构中引入，用于加速机器学习和科学应用的矩阵乘法和累加操作。这些指令在小的矩阵块（例如4x4块）上操作。

Tensor Cores还可以在比输入更高的精度下进行乘积计算和累加。例如，在使用FP16（半精度浮点数）输入进行训练时，张量核心可以在不损失精度的情况下进行乘积计算，并以FP32（单精度浮点数）进行累加。

满足特定操作的可以在Tensor Core中执行，不满足的就在普通Cuda Core中执行（比如一个简单的FP16精度的element-wise的加法操作）。

2. GPU Execution Model

为了利用其并行资源，GPU通常会同时执行许多线程。

了解线程数与GPU性能之间的关系，有两个关键概念：

GPU使用2级线程层次结构来执行函数。给定函数的线程被分组成相同大小的线程块，并启动一组线程块来执行函数。
GPU通过切换到其他线程的执行来隐藏依赖指令的延迟。因此，为了有效利用GPU，所需的线程数量远远高于核心数或指令流水线数。

两级线程层次结构是因为GPU拥有许多流式多处理器（SM），每个SM又具有执行许多线程的流水线，并使其线程通过共享内存和同步进行通信。在运行时，一个线程块会被放置在一个SM上执行，使线程块中的所有线程能够高效地进行通信和同步。

使用单个线程块启动一个函数只会使一个SM工作，因此为了充分利用拥有多个SM的GPU，需要启动许多线程块。由于一个SM可以同时执行多个线程块，通常希望线程块的数量是SM数量的几倍。这样做的原因是为了尽量减少“尾部”效应，即在函数执行结束时只剩下少数活跃的线程块，从而在那段时间内GPU利用率不足，如下图所示。

GPU Execution Model

Utilization of an 8-SM GPU when 12 thread blocks with an occupancy of 1 block/SM at a time are launched for execution. Here, the blocks execute in 2 waves, the first wave utilizes 100% of the GPU, while the 2nd wave utilizes only 50%.

我们使用波（wave）这个术语来指代一组可以并行运行的线程块。以多个wave的线程块来启动函数执行是最高效的方式——在尾波（最后一个波次）中花费的时间占比较小，从而最小化尾效应，因此也减少了处理尾效应的需求（这种效应会导致总体执行效率降低，因为GPU的一部分核心在等待少数几个线程块完成时处于空闲状态）。对于高端GPU来说，通常只有当启动少于300个线程块的情况下，才需要考虑尾效应。

3. Understanding Performance

在给定的处理器上执行函数的性能受以下三个因素之一的限制：

内存带宽
数学带宽
延迟

以下称呼的数学，意思就是某个op的算术强度，就是这个操作需要多少flops才能实现

考虑一个简化模型，其中一个函数从内存中读取其输入，执行数学运算，然后将其输出写入内存。假设：

是显存访问和读取的时间, 是在数学算子上计算的时间。如果我们假设这两个操作可以在不同线程中执行, 可以重叠（overlap），那么总时间就是 )。两个当中谁大谁就是限制当前op性能的原因。如果是数学那么就是math limited, 如果是显存那就是memory limited。

在显存或数学运算中花费多少时间取决于算法及其实现方式，以及处理器的带宽。内存时间等于在内存中访问的字节数除以处理器的内存带宽。数学时间等于操作次数除以处理器的数学带宽。因此，在给定处理器上，如果一个给定算法是受到数学限制，则：

通过简单的代数，不等式可以重新排列为：

左侧是算法实现操作次数和访问字节数的比值，也就是俗称的arithmetic intensity。右侧是处理器计算能力和内存带宽之比，有时被称为_ops:byte比率_。因此，在给定处理器上，如果算法的算术强度高于处理器的ops:byte比率，则该算法受到数学限制。相反地，如果其算术强度低于处理器的ops:byte比率，则该算法受到内存限制。

让我们考虑一些深度神经网络的具体例子，列在下面的表格中。对于这些示例，我们将比较算法的算术强度与在NVIDIA Volta V100 GPU上的ops:byte比率。V100具有125 FP16 Tensor TFLOPS的峰值数学速率，约为900 GB/s 的芯片外存储器带宽和3.1 TB/s 的芯片内L2带宽，因此其ops:byte比率介于40到139之间，取决于操作数据（片内或片外存储器）的来源。

Operation	Arithmetic Intensity	Usually limited by...
Linear layer (4096 outputs, 1024 inputs, batch size 512)	315 FLOPS/B	arithmetic
Linear layer (4096 outputs, 1024 inputs, batch size 1)	1 FLOPS/B	memory
Max pooling with 3x3 window and unit stride	2.25 FLOPS/B	memory
ReLU activation	0.25 FLOPS/B	memory
Layer normalization	< 10 FLOPS/B	memory

如表所示，许多常见操作的算术强度很低——有时每从内存读取并写入两字节元素只执行一次操作。请注意，这种分析是一种简化，因为我们只计算了_算法_中使用的操作。实际上，函数还包含算法中未明确表达的操作指令，如内存访问指令、地址计算指令、控制流指令等。

算术强度和每字节操作数分析假设工作负载足够大，能够充分利用给定处理器的数学和内存管道。然而，如果工作负载不够大，或没有足够的并行性，处理器将被低效利用，性能将受到延迟的限制。例如，考虑启动一个单线程，它将访问16字节并执行16000次数学操作。尽管算术强度为1000 FLOPS/B，并且在V100 GPU上的执行应该是数学限制的，但只创建一个线程会严重低效利用GPU，使其几乎所有的数学管道和执行资源处于空闲状态。

此外，算术强度计算假设输入和输出仅从内存中访问一次。对于算法实现多次读取输入元素的情况并不少见，这会有效降低算术强度。因此，算术强度是一种初级近似；如果需要更准确的分析，需要看profiler。

4. DNN Operation Categories

我们平常的深度学习网络中有很多基础的op层，大概分为三类：

4.1. Elementwise Operations

Elementwise operations 可以是一元或二元操作；关键在于此类层对张量中的每个元素执行数学运算，独立于张量中所有其他元素。

例如，ReLU层对输入张量中的每个x返回max(0, x)。同样，两个张量的逐元素加法计算每个输出和值与其他和值无关。这类层包括大多数非线性函数（sigmoid、tanh等）、比例、偏置、加法等。

这些层往往受到内存限制，因为它们每访问一个字节只执行少量操作。有关激活函数的更多详细信息，请参见《优化内存限制层用户指南》中的激活[6]部分。

4.2. Reduction Operations

reduce operations也就是常见的规约操作。

例如，池化层在输入张量的某些邻域内计算值。Batch normalization在每个输出元素的操作中使用之前，计算张量的均值和标准差。除了池化和归一化层外，SoftMax也属于归约类别。典型的归约操作算术强度较低，因此受到内存限制。有关池化层的更多详细信息，可以在池化[7]中找到。

4.3. Dot-Product Operations

该类操作可以表示为两个tensor的点积，通常是一个权重（学习参数）张量和一个激活张量。也就是常说的矩阵乘法。

包括全连接层（fully-connected layers），可以很自然地表示为矩阵向量和矩阵矩阵乘积。

卷积也可以理解为就是做 matrix-vector and matrix-matrix multiplies，常见的卷积实现im2col就是典型的矩阵乘。不过矩阵乘处于算术瓶颈还是显存瓶颈也分情况，如果batch比较小或者输入比较小，那么还是memory bound，只有当batch大的时候才会是compute bound

5. 总结

如何大概预估我们某个op在某个GPU上的性能受限在哪，有以下几点方法：

弄清楚这个GPU上的 SM 数量，并确定 GPU 的 ops:bytes 比例
预估当前op算法的运算强度（arithmetic intensity）
通过估算线程块的数量和大小，判断是否有足够的并行度来充分利用 GPU。如果线程块的数量至少是 SM 数量的 4 倍，并且每个线程块由几百个线程组成，那么可能就有足够的并行度
最可能的性能瓶颈是：

如果没有足够的并行度，则是GPU常见的延迟问题，延迟无法掩盖
如果有足够的并行度且算法运算强度高于 GPU ops:bytes 比例，则是数学运算瓶颈
如果有足够的并行度且算法运算强度低于 GPU ops:bytes 比例，则是内存瓶颈

说白也就是看Memory-bound还是Compute-bound

参考资料

[1]

GPU Performance Background User's Guide: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html

[2]GPU Architecture Fundamentals: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#gpu-arch

[3]GPU Execution Model: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#gpu-execution

[4]Understanding Performance: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#understand-perf

[5]DNN Operation Categories: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#dnn-op-cat

[6]激活: https://docs.nvidia.com/deeplearning/performance/dl-performance-memory-bound/index.html#activations

[7]池化: https://docs.nvidia.com/deeplearning/performance/dl-performance-memory-bound/index.html#pooling

二、CUDA性能检查清单

文中详细讨论了如何合并全局内存访问、最大化占用率、理解内存与计算限制、最小化线程分化、通过Tiling重用数据、私有化内存使用、线程粗化、以及通过算法重写来提升性能。此外，文章还提到了使用ncu工具来分析和优化kernel性能，并通过实际的CUDA代码示例来阐释这些概念。最后，文章提到了Flash Attention算法作为通过数学角度重写算法以提升性能的一个例子。

大佬的课程笔记，欢迎关注：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/cuda-mode

CUDA-MODE课程笔记第8课: CUDA性能检查清单

课程笔记

这节课实际上算是CUDA-MODE 课程笔记第一课: 如何在 PyTorch 中 profile CUDA kernels 这节课更细节的讲解。另外关于nsight compute相关指标细节解释可以参考 CUDA-MODE 第一课课后实战（上），CUDA-MODE 第一课课后实战（下）这两篇笔记。

将GPU用于计算，我们最关心的肯定是性能。比较幸运的是，当我们掌握一些性能优化技巧之后它往往会经常被用到。这节课将会系统的介绍这些性能优化技巧。

本节课的课件和代码都在 https://github.com/cuda-mode/lectures 开源，我们可以用nvcc编译lecture8下面的cu文件，然后使用ncu进行profile。此外，这里的方法遵循了 https://arxiv.org/pdf/1804.06826.pdf 这篇论文的风格。up主也非常推荐大家阅读下这篇论文，用claude 3.5问了一下paper的主要内容，如下截图：

可以看到这篇论文主要是对Volta架构的GPU的架构细节进行分析，对性能优化是很重要的。

这张Slides从物理学的角度分析了下SRAM和DRAM的区别：

DRAM由1个晶体管和1个电容器构成；SRAM: 由6个晶体管构成
SRAM 比 DRAM 更快，但也更贵；SRAM 占用更多空间且发热更多；实际上SRAM就对应了GPU的Shared Memory，而DRAM对应的则是Shared Memory。

这里的youtube链接作者Bill是NVIDIA的首席科学家，他解释了很多为什么GPU设计成现在这个样子，并且由浅入深，基础细节讲的非常清楚。

这里的"性能检查清单"（Performance checklist），列出了一系列优化GPU程序性能的策略和技巧：

合并全局内存访问（Coalesced Global Memory Access）
最大化占用率（Maximize occupancy）
理解是内存受限还是计算受限（Understand if memory or compute bound）
最小化线程分化（Minimize control divergence）
Tiling以更好的重用数据（Tiling of reused data）
私有化（Privatization）
Thread Coarsening
使用更好的数学方法重写算法（Rewrite your algorithm using better math）

这里的Privatization指的应该就是Shared Memory/寄存器优化全局内存读取，而Coarsening大概指的就是一个线程应该完成多少任务，一般情况下我们让一个线程完成的任务尽量少，但是在Compute Bound情况下，让一个线程执行更多的工作可以让程序运行得更快。最后一点更好的数学方法重写算法的经典例子就是Flash Attention。

这张Slides讲述了GPU内存访问延迟的相关内容，下面的Figure3和表格都来自 https://arxiv.org/pdf/2208.11174 ，这个表格（Table IV），列出了不同类型内存的访问延迟（以时钟周期为单位）：

全局内存（Global memory）: 290 cycles
L2 缓存: 200 cycles
L1 缓存: 33 cycles
共享内存（Shared Memory）: 读取23 cycles，写入19 cycles

我后面也找到了这个paper里面做micro benchmark的代码：https://www.stuffedcow.net/research/cudabmk?q=research/cudabmk ，后面如果有空继续阅读下这篇 paper 以及测试代码。

这张Slides讲述了延迟（latency）在计算机系统中的重要性和一些相关概念。

标题 "It's the latency stupid" 强调了延迟的重要性。
吞吐量（Throughput）和延迟（Latency）的对比：

吞吐量容易提高，但延迟却很难降低。
举例说明：即使你可以并行使用80条电话线，每条线传输一个比特，但100毫秒的延迟仍然存在。

量化（Quantization）技术：
用于减少数据包大小的一种方法。
例如，Bolo（可能是某个系统或协议）尽可能使用字节（byte）而不是16位或32位字来减少数据包大小。
底部提供了一个网址链接，包含更多关于这个话题的详细讨论。

这张Slides开始介绍内存合并（Memory Coalescing）的概念。我们无法减少延迟，但可以通过读取连续的内存元素来隐藏延迟。Slides建议在进行案例研究时要关注以下三个方面：

DRAM Throughput（DRAM吞吐量）
Duration（持续时间）
L1 cache throughput（L1缓存吞吐量）

这里说的内存合并的案例就是 https://github.com/cuda-mode/lectures/blob/main/lecture_008/coalesce.cu 这里所展示的。代码如下：

#include <iostream>  
#include <cuda_runtime.h>  __global__ void copyDataNonCoalesced(float *in, float *out, int n) {  int index = blockIdx.x * blockDim.x + threadIdx.x;  if (index < n) {  out[index] = in[(index * 2) % n];  }  
}  __global__ void copyDataCoalesced(float *in, float *out, int n) {  int index = blockIdx.x * blockDim.x + threadIdx.x;  if (index < n) {  out[index] = in[index];  }  
}  void initializeArray(float *arr, int n) {  for(int i = 0; i < n; ++i) {  arr[i] = static_cast<float>(i);  }  
}  int main() {  const int n = 1 << 24; // Increase n to have a larger workload  float *in, *out;  cudaMallocManaged(&in, n * sizeof(float));  cudaMallocManaged(&out, n * sizeof(float));  initializeArray(in, n);  int blockSize = 128; // Define block size  // int blockSize = 1024; // change this when talking about occupancy  int numBlocks = (n + blockSize - 1) / blockSize; // Ensure there are enough blocks to cover all elements  // Launch non-coalesced kernel  copyDataNonCoalesced<<<numBlocks, blockSize>>>(in, out, n);  cudaDeviceSynchronize();  initializeArray(out, n); // Reset output array  // Launch coalesced kernel  copyDataCoalesced<<<numBlocks, blockSize>>>(in, out, n);  cudaDeviceSynchronize();  cudaFree(in);  cudaFree(out);  return 0;  
}

这里段程序比较简单，用于演示内存合并（Memory Coalescing）的概念和其对性能的影响。它主要做了以下事情：

定义了两个CUDA kernel：

copyDataNonCoalesced kernel：非合并内存访问模式，以非连续的方式读取输入数组（使用 (index * 2) % n 作为索引），这种访问模式会导致非合并的内存访问，可能降低性能。
copyDataCoalesced kernel：合并内存访问模式，以连续的方式读取输入数组（直接使用 index 作为索引），这种访问模式允许合并内存访问，可以提高性能。

主函数：
分配统一内存（Unified Memory）用于输入和输出数组，初始化输入数组。
设置CUDA网格和块的大小，分别运行非合并和合并的kernel，在每次kernel执行后使用 cudaDeviceSynchronize() 确保GPU操作完成。

接着使用nvcc -o benchmark coalesce.cu来编译程序，然后执行ncu benchmark来Profile程序。

对于copyDataNonCoalesced kernel来说，DRAM内存吞吐量大约是89%，L1 Cache的吞吐量是30%，kernel的执行时间是764us。

对于copyDataCoalesced kernel来说，L1 Cache的吞吐量大约是37%，DRAM内存吞吐量是82%，执行时间是558us。

我们可以看到合并内存访问的kernel是有明显的性能提升的。可以预见，随着输入数据量的增大合并内存访问的优势会更明显。ncu的结果里面还提示计算的理论occupancy（100.0%）和实测的实际occupancy占用（77%）之间的差异可能是由于 kernel 执行期间的warp调度开销或工作负载不平衡导致的。在同一kernel 的不同块之间以及块内的不同 warps 之间都可能发生负载不平衡。把上面程序中的int blockSize = 128改成int blockSize = 1024再次用ncu profile，可以发现occupancy提升到了85.94%。

这张Slides讨论了GPU中的占用率（Occupancy）问题，主要内容如下：

两种quantization问题：

a) Tile quantization：矩阵维度不能被线程块Tile大小整除。
b) Wave quantization：Tile总数不能被GPU上的SM（流多处理器）数量整除。

性能图表比较和分析：
左图(a)：cuBLAS v10 上 NN GEMM 的性能
右图(b)：cuBLAS v11 上 NN GEMM 的性能
两图都是在 M = 1024, N = 1024 的矩阵维度下进行的测试
左图(a)显示性能呈现明显的阶梯状，有大幅波动。
右图(b)显示性能波动较小，整体更加平滑。我们可以看到cuBLAS v11 可能采用了更好的调度策略或优化技术，减少了由于Tile和Wave Quantization 导致的性能波动。

这张Slides讲解了在PyTorch中使用padding（填充）来解决Tensor Core矩阵乘法维度要求的问题。具体内容如下：

在PyTorch环境中，使用padding是解决某些问题的方法。
表格展示了不同cuBLAS和cuDNN版本下，使用Tensor Core的数据精度要求。这些要求适用于矩阵维度M、N和K。
版本区分：

左列：cuBLAS < 11.0 和 cuDNN < 7.6.3 的旧版本
右列：cuBLAS ≥ 11.0 和 cuDNN ≥ 7.6.3 的新版本

数据类型的要求：
INT8：旧版本要求16的倍数；新版本总是可用，但16的倍数最高效，在A100上128的倍数最佳。
FP16：旧版本要求8的倍数；新版本总是可用，但8的倍数最高效，在A100上64的倍数最佳。
TF32：旧版本不适用；新版本总是可用，但4的倍数最高效，在A100上32的倍数最佳。
FP64：旧版本不适用；新版本总是可用，但2的倍数最高效，在A100上16的倍数最佳。

新版本的cuBLAS和cuDNN提供了更灵活的Tensor Core使用条件。而A100 GPU可能需要更大的倍数来获得最佳性能。Padding可以用来将矩阵维度调整为这些推荐的倍数，以提高性能。

在CUDA中提升Occupancy的一个方法是修改kernel。

CUDA Occupancy calculator工具可以帮我们自动计算达到更好Occupancy的kernel启动参数，在上一节合并访存的.cu中调用这个Api结果显示，对于T4 GPU，最优的配置是网格大小为40，块大小为1024。代码见：https://github.com/cuda-mode/lectures/blob/main/lecture_008/occupancy.cu

在对这个程序进行ncu的时候有新的问题，那就是下面所展示的：

警告（WRN）：内存利用率高于计算利用率：请查看内存工作负载分析部分以识别DRAM瓶颈。检查内存重放（合并）指标，以确保您正在有效利用传输的字节。同时考虑是否可以通过每次内存访问执行更多工作（kernel融合）或是否有可以（重新）计算的值。

接下来开始讨论这个问题

讨论之前需要先了解一下这张Slides展示的Roofline模型，它决定了一个cuda kernel是compute bound还是memory bound。

这张Slides讲解了算术强度（Arithmetic intensity）的概念及其在处理器性能分析中的应用。这个slides来自gtc2019的一个讲解。

左侧指标是数学运算和内存操作的算法混合，称为算术强度。右侧指标是处理器的ops/byte比率。例如，V100 GPU可以执行125/0.9=139 FLOPS/B。比较算术强度和ops/byte比率可以指出算法受什么因素限制。

下面还给出了操作类型及其算术强度表格：

Residual addition（残差加法）：0.166，受内存限制
ReLU activation（ReLU激活）：0.25，受内存限制
Batch normalization（批量归一化）：O(10)，受内存限制
Convolution（卷积）：1-10000+（假设FP16数据），可能受内存或数学运算限制

链接：https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s9926-tensor-core-performance-the-ultimate-guide.pdf

这张slides讲解了ReLU（Rectified Linear Unit）函数的算术强度分析：

ReLU函数定义：f(x) = max(0, x)，应用于向量的每个元素。
操作描述：对每个元素进行1次读取、1次比较操作，可能还有1次写入。
数据类型：假设使用float32，即每个数占4字节（32位）。
计算分析：

操作数（Ops）：1（每个元素一次比较操作）
字节数（Byte）：2 * 4 = 8（读取和可能的写入，每次4字节）

算术强度计算：
最坏情况：1/8（当每个元素都需要写入时）
最好情况：1/4（当不需要写入时，只有读取操作）结论：1/4 < 1，表明ReLU操作受内存带宽限制（Memory bound）

这张Slides对Float16的ReLU进行了算术强度分析，可以看打这种情况下最坏的算术强度是1/4，而不是Float32时的1/8，因此量化是可以提高计算强度的。

这张Slides讲解了矩阵乘法（Matmul）的算术强度分析。其中：

FLOPS（浮点运算次数）计算：

对C中的每个输出元素，需要A的一行和B的一列做点积
需要N次乘法和N次加法
总FLOPS = M * K * 2N

字节数计算：
加载矩阵A和B：MN + NK
写入输出矩阵C：MK
总字节数 = MN + NK + MK
算术强度（AI）计算：
AI = 2MNK / (MN + NK + MK)
结论：
对于大型矩阵，计算受限（Compute bound）
否则，带宽受限（Bandwidth bound）

这张Slides总结了如何优化不同类型的kernels：

带宽受限的kernel（Bandwidth Bound Kernels）优化策略：

Fuse（融合）：合并多个操作以减少内存访问
Quantize（量化）：使用更小的数据类型来减少内存传输
Compile（编译）：可能指使用特定的编译技术来优化内存访问模式

计算受限的kernel（Compute Bound Kernels）优化策略：
Write a better algorithm（编写更好的算法）：这意味着需要从算法层面进行优化

关于矩阵乘法Tiling减少全局内存访问请查看以前的CUDA-MODE 课程笔记第四课: PMPP 书的第4-5章笔记。

这张Slides对应这里的代码：https://github.com/cuda-mode/lectures/blob/main/lecture_008/divergence.cu ，主要是对下面2个kernel进行分析：

__global__ void processArrayWithDivergence(int *data, int N) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  if (idx < N) {  if (data[idx] % 2 == 0) {  data[idx] = data[idx] * 2; // 注意这个分支比下面的分支要慢，可能一个Warp里执行这个分支的线程会落后，Warp里的其它线程必须等待这些线程计算完成  } else {  data[idx] = data[idx] + 1;  }  }  
}  __global__ void processArrayWithoutDivergence(int *data, int N) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  if (idx < N) {  int isEven = !(data[idx] % 2); // 这里做的事情和上面相同，但是规避了线程分化问题  data[idx] = isEven * (data[idx] * 2) + (!isEven) * (data[idx] + 1);  }  
}

控制分歧（control divergence）与占用率（occupancy）有关，但如果条件语句导致大量线程闲置，这是不好的。
processArrayWithDivergence 耗时 0.074272 毫秒; processArrayWithoutDivergence 耗时 0.024704 毫秒;这表明去除control divergence可以显著提高性能（约3倍）。
"ncu --set full divergence" 用这行命令来设置线程control divergence分析。

对于compute bound的kernel，让线程可以做更多工作，可能会更快。

性能比较：

运行命令：main ~/lecturex ./benchmark
VecAdd 执行时间：0.245600 ms
VecAddCoarsened 执行时间：0.015264 ms

关键观察：
VecAddCoarsened启动了一半的线程数量
尽管线程数减少，但执行速度显著提高（约16倍）

这里的代码在 https://github.com/cuda-mode/lectures/blob/main/lecture_008/coarsening.cu 。

这也许可以解释Lecture 7中为什么对于Int4 Weight Only量化的高效kernel实现比普通的fp16的Kernel跑得更快。

这张Slides讨论了在GPU编程中的"私有化"（Privatization）技术。要点为：

将部分更新应用到数据的私有副本上，然后再写回全局或共享内存。
示例：

滑动窗口算法（Sliding window algorithm）
图示：1 2 [3] [4] [5] 6 7
这表明算法在一个局部窗口内进行操作。

Privatization的优势：
更高的占用率（Higher occupancy）
更高的计算SM吞吐量（Higher compute SM throughput）
更低的DRAM吞吐量（Lower DRAM throughput）

这个滑动窗口算法对应的例子就是Mistral等大模型里面的滑动窗口注意力算法，

解释下这个图：

左侧矩阵：Vanilla Attention

展示了传统注意力机制，每个token都可以关注到所有其他token。
矩阵是下三角形，表示每个token只能关注到它自己和之前的token。

中间矩阵：Sliding Window Attention
展示了滑动窗口注意力机制，每个token只关注固定窗口大小内的相邻token。
这里窗口大小W=3，可以看到每个token只与前后3个token有连接。
右侧图：Effective Context Length
展示了多层滑动窗口注意力如何扩大有效上下文长度。
每一层都可以使信息向前传播W个token。

总结来说，传统注意力的操作数量与序列长度成二次方关系，内存使用随token数线性增长。在推理时，这会导致更高的延迟和更低的吞吐量，因为缓存可用性降低。滑动窗口注意力通过限制每个token最多只关注前一层的W个token来缓解这个问题。虽然窗口外的token不直接参与注意力计算，但它们仍然可以影响下一个词的预测。在每个注意力层，信息可以向前传播W个token。经过k个注意力层后，信息可以向前传播最多k × W个token。

继续讨论Privatization技术，https://github.com/cuda-mode/lectures/blob/main/lecture_008/privatization.cu 这里的核心代码是：

// CUDA kernel for vector addition without privatization  
__global__ void vectorAdd(const float *a, const float *b, float *result, int n) {  int index = threadIdx.x + blockIdx.x * blockDim.x;  if (index < n) {  result[index] = a[index] + b[index];  }  
}  // CUDA kernel for vector addition with privatization  
__global__ void vectorAddPrivatized(const float *a, const float *b, float *result, int n) {  int index = threadIdx.x + blockIdx.x * blockDim.x;  if (index < n) {  float a_private = a[index]; // Load into private memory  float b_private = b[index]; // Load into private memory  result[index] = a_private + b_private;  }  
}

我们把a[index]，b[index]加载到private memory里面避免对全局内存的直接操作，但在这个VectorAdd的例子中没有加速。

但是在下面的滑动窗口求和的例子中通过把global memory加载到shared memory中，然后进行累加时求和操作就是在shared memory中进行操作。代码链接：https://github.com/cuda-mode/lectures/blob/main/lecture_008/privatization2.cu

// Kernel without privatization: Direct global memory access  
__global__ void windowSumDirect(const float *input, float *output, int n, int windowSize) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  int halfWindow = windowSize / 2;  if (idx < n) {  float sum = 0.0f;  for (int i = -halfWindow; i <= halfWindow; ++i) {  int accessIdx = idx + i;  if (accessIdx >= 0 && accessIdx < n) {  sum += input[accessIdx];  }  }  output[idx] = sum;  }  
}  // Kernel with privatization: Preload window elements into registers  
__global__ void windowSumPrivatized(const float *input, float *output, int n, int windowSize) {  int idx = blockIdx.x * blockDim.x + threadIdx.x;  int halfWindow = windowSize / 2;  __shared__ float sharedData[1024]; // Assuming blockDim.x <= 1024  // Load input into shared memory (for demonstration, assuming window fits into shared memory)  if (idx < n) {  sharedData[threadIdx.x] = input[idx];  __syncthreads(); // Ensure all loads are complete  float sum = 0.0f;  for (int i = -halfWindow; i <= halfWindow; ++i) {  int accessIdx = threadIdx.x + i;  // Check bounds within shared memory  if (accessIdx >= 0 && accessIdx < blockDim.x && (idx + i) < n && (idx + i) >= 0) {  sum += sharedData[accessIdx];  }  }  output[idx] = sum;  }  
}

作者最后讲的一点是以Flash Attention为例，如果你可以从数学的角度重写算法，有可能可以让代码的性能大幅度提升。比如Flash Attention利用Safe Softmax的数学形式分块计算Attention。这部分讲解已经非常多了，大家可以参考 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/README.md 里面收集的Flash Attention相关的资料，最后这几张Slides就不再赘述了。

总结

这节课相当于对Lecture 1更系统的补充，重点介绍GPU kernel优化的一些实用技术和分析工具ncu。课程深入探讨了“性能检查清单”，概述了关键的优化策略：

合并全局内存访问：确保线程访问连续的内存位置，以最大限度地提高内存带宽利用率。
最大化占用率：优化kernel启动参数，以充分利用 GPU 的处理能力。
理解内存与计算限制：分析kernel特征以确定限制因素（内存带宽或计算能力），并应用适当的优化技术。
最小化线程的分化：避免导致warp内的线程采用不同执行路径的条件语句，从而导致性能下降。
Tiling以重用数据：组织数据访问模式，以最大限度地提高缓存中的数据局部性和重用率。
私有化：利用私有内存（寄存器或共享内存）来减少全局内存访问并提高占用率。
线程粗化：调整线程粒度以平衡工作负载并最小化线程开销。
算法重写：探索替代的数学公式或算法以提高计算效率。

up主还通过cuda代码示例和使用 ncu 工具的分析结果来说明这些概念。它展示了内存合并、占用率优化和控制分歧最小化的好处。它还介绍了用于分析kernel性能瓶颈的 Roofline 模型，以及算术强度的概念，以确定kernel是受内存限制还是受计算限制。接着up主讨论了私有化技术，重点介绍了在滑动窗口算法中使用共享内存来改善数据局部性和减少全局内存访问。最后，简单描述了一下通过从数学角度重写算法来获得显著性能提升的潜力，并以 Flash Attention 为主要示例，这里就没有截最后几张slides了，毕竟Flash Attention的资料非常多了。

三、手写CUDA卷积算子

这里主要介绍如何利用CUDA实现一个2D卷积算子，实现过程较为简单，最终的实现效果可以在较小的尺寸下取得比cudnn快较大的性能。实测在以下参数配置下可以达到平均

1.2倍cudnn的性能。

CUDA介绍（from chatGPT）

现在深度学习大行其道，作为深度学习的基础软件设施，学习cuda也是很有意义的。本篇文章主要介绍如何利用CUDA实现一个2D卷积算子，实现过程较为简单，最终的实现效果可以在较小的尺寸下取得比cudnn快较大的性能。实测在以下参数配置下可以达到平均1.2倍cudnn的性能（娱乐结果，还与cudnn配置有关，更小更快）。

TIPS: 跳过cudnn初始化的时间，99轮平均时间

const int inC = 6;const int inH = 768;const int inW = 512;const int kernelH = 6;const int kernelW = 6;const int outC = 6;const int outH = inH - kernelH + 1;const int outW = inW - kernelW + 1;

1 卷积操作通俗介绍

1.1 数据布局(data layout)

卷积操作主要针对图像进行运算，我们常见的RGB即为三通道的二维图像，那么就可以通过一个一维数组存储所有的数据，再按照不同的布局去索引对应的数据，现在主要使用nchw和nhwc两种数据布局，其中

n - batch size 也可以理解为"图像"数量
c - channel num 即我们说的通道数量
h - height 图像高度，每个通道的高度宽度是一致的
w - width 图像宽度

那么显然nchw就是逐个通道的读取图像，nhwc即对所有通道的同样位置读取数据后，再切换到下一个为止

一个是优先通道读取，一个是优先位置读取

还有一种CHWN布局，感觉比较奇怪，并未过多了解

详细的可以参考英伟达官方文档Developer Guide : NVIDIA Deep Learning cuDNN Documentation (https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html)

nchw layout

nhwc layout

本文是按照nchw数据格式来进行算子的实现的。

1.2 直接卷积

相信大家都或多或少听过卷积，可以通过gpt的回答来直观地认识卷积操作

最基本的直接卷积操作是十分简单的，你可以想象一个滑动的矩阵窗口在原矩阵上移动，对应位置进行点积，得到结果后求和放到目标矩阵上，可以用以下图像直观地理解这一过程，向老师称为对对碰:)

图源：国科大模式识别课程

你会注意到上述过程中怎么没有什么channel的参与，只有一个矩阵

多输入通道的情况下，就是对每个通道的相同位置分别与卷积核进行对对碰，结果累加作为输出矩阵值；

多输入多输出通道，即对每个输出通道都进行上述操作

对于通道的理解建议参考[@双手插袋]的文章CNN卷积核与通道讲解 (https://zhuanlan.zhihu.com/p/251068800)

那么我们需要知道的是直接卷积操作其实就是原矩阵与卷积核间的对对碰，产生所谓的特征图feature map，十分的简单，这也方便我们对其进行并行任务划分

注意到上述文章中并没有提到padding和stride，本篇文章并没有针对padding和stride的实现

padding

padding是作为对图像的填充，可以发现上面的特征图尺寸缩小了一圈，是因为直接卷积势必会造成这一结果

通过padding可以加强图像边缘特征，避免边缘特征被忽略

stride

stride可以简单的理解为跨步，即上面的小窗口在矩阵上滑动的步长，默认为1

即上述图像中下一次卷积的中心应该是4为中心的3*3子矩阵

如果你设置为2，那么下一次是3为中心的3*3子矩阵了

1.3 其他卷积计算方法

除去直接卷积，也有一些其他方法进行卷积，感兴趣的读者可以自行了解，仅举以下几例参考

Img2col

即把图像展开为一个行向量组，卷积核/滤波器（kernel/filter）展开为一列或多列向量，转化为矩阵乘去计算卷积结果

FFT method

利用傅里叶变换的频域变换去做卷积，这样做的优势是计算量会小很多

Winograd Algorithm

也是一种将图像变换到另外一个空间再去做运算再做变换得到结果，会减少很多乘法运算

2 整体实现思路

2.1 block与thread划分

首先我们需要考虑如何对代表图像的多通道矩阵来进行block与thread的划分，这一部分是有说法的

不同的切分方式会让block在SM上的流转效率有很大的差别

本文仅提供一个十分草率的切分，我们都清楚目前在英伟达的GPU上，任务的调度最小单元是warp

一个warp以32个线程为一组，故通过8*4的block来进行矩阵的切分，每个block里共32个位置

这样可以保证每个block上到SM时不用去与其他的block拼接线程，产生额外开销

注意我这里用的是位置，并不是元素，32个线程，每个线程去负责一个位置的计算

以16*20的矩阵为例，对其进行划分的结果如下图所示，(x,y)是笛卡尔坐标系，与行主序不同

2.2 数据转移

关于位置和规模(size)

那么为什么说一个block有32个位置，而不是32个元素呢，首先注意到卷积操作虽然遍历到了原矩阵的所有元素

但是你按中心点的位序去数的话（以卷积核3*3为例），结果应该是这个样子

注意这里仅示意卷积中心点范围，请与后文作区分

按3*3矩阵的中心来看，中心正好是去掉外面一圈的位置，按照左上角元素来看，恰好应该是（左上角，右下角）

这样一个区间，参数解释如下

row_num 原矩阵中一行元素的数目
inH inW 原矩阵的H W
kernelH kernelW 卷积核的H W
outH outW 输出矩阵的H W

当然你也可以用中心点而不是左上角的元素作为窗口的标识来设计算法

恰巧你上面算出来的这个范围也正是你得到的feature map的下标范围

我们也可以得到输出矩阵的规模为

请注意大小和位置下标的区别，一个从1开始数一个从0开始数

一个block的数据转移

确定了整体的尺寸，那么我们来看一个block需要的数据尺寸是多少呢

显然你可以发现，对于输出矩阵进行block划分是更合理的，这样可以保证一个block

32个位置恰好对应输出矩阵的32个位置，而不用过多的去考虑输出矩阵的排布

那么对于上述提到的划分，可以通过下图来直观感受block划分在原矩阵的效果

22*18的in产生20*16的out

那么一个block用到的元素范围应该是哪些呢，我们要做的是卷积操作，每个中心点应该有对应卷积核大小的矩阵参与运算，那么以(0,0)和(4,1)的block为例，给出他们的涉及原矩阵范围如下图所示

蓝色为一个block需要用到的原矩阵元素

那么我们可以确定一个block，8×4的情况下需要读取10×6的原矩阵的元素，也是+kernelH-1来确定的

那么对应输出矩阵就是一个萝卜一个坑了，不需要额外考虑

这样就确定了一个block需要从GMEM到SMEM的元素范围

至于怎么转移，我们在代码实现中讲述，当然你可以单独指定某几个进程去完成所有的转移任务

2.3 计算逻辑

不考虑channel

不考虑channel的情况下，即单输入通道单输出通道单卷积核这样最简单的情况

我们只需要做三件事

① 将block对应的数据转移到SMEM中

② 利用线程的tid去计算对应输出矩阵位置的结果

③ 将结果写回输出矩阵

只考虑inC

这种情况下我们要做的额外的事儿就多一点

加一层循环，让每个线程计算多个in channel的数据，并累加起来作为结果

需要用到一个寄存器来存储这个中间结果

考虑inC与outC

其实要做的事情也就比上面多一点，就是开大点空间

让线程去存储多个outC的中间结果，分别累加

最后写回的时候也分别写回即可

3 详细实现过程

3.1 整体实现思路

主要从自己的角度出发去还原怎样一步步构造出这样一个初级的算法

首先实现一个最简单的版本，CPU串行版本，并保证CPU串行版本可以获取正确的结果

此后再在其基础上进行并行化的改造，而直接卷积运算的过程其实相对是比较简单的

我们在不考虑padding与stride的情况下，是可以不借助任何参考资料来直接完成第一版代码的

3.1.1 CPU串行版本的卷积算子

#define element_type float
#define OFFSET(row, col, ld) ((row) * (ld) + (col))/*@brief: 串行卷积实现 CPU代码 NCHW@param in inC inH inW: 输入矩阵(数组) channel height width@param out outC outH outW: 输出矩阵 channel height width@param kernel kernelH kernelW: 卷积核 height width
*/
void serial_convolution(element_type *in, element_type *out, element_type *kernel, int batch_size,int inC, int inH, int inW,int outC, int outH, int outW,int kernelH, int kernelW)
{float val;int out_pos, in_pos, kernel_pos;for (int oc = 0; oc < outC; oc++) // 每个输出通道{// 对一个位置的操作 用当前输入channel卷积去对相应的输出channel// 保证每个outChannel都是多inChannel累积的结果for (int i = 0; i < outH; i++){for (int j = 0; j < outW; j++){val = 0; // 避免累积和需要多次读取写入out_pos = oc * outH * outW + OFFSET(i, j, outW);for (int ic = 0; ic < inC; ic++) // 对每个输入通道{for (int ii = 0; ii < kernelH; ii++){for (int jj = 0; jj < kernelW; jj++){in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);val += in[in_pos] * kernel[kernel_pos];}}}out[out_pos] = val; // 与cudnn计算结果为相反数}}}
}

这是我最终完成的CPU串行版本代码，可以发现套了足足有5层循环

在我们传统观念中，这可是 O(n5)O(n^5)O(n^5) 的最笨算法了

不过没有关系，我们关注的并不是他的性能，cuda上也不会去跑这一版代码

我们需要关注的是怎么样能得到正确的结果，且如何设计循环的嵌套关系来使用尽量少的访存次数

使用尽量多的本地中间结果，这样可以尽可能地减少我们的算法在访存方面的消耗

要明白GPU上的线程如果去读GMEM上的数据需要几百个时钟周期，读SMEM需要几十个时钟周期

读取SM上的寄存器需要的时钟周期会更少！

因此我们需要竭力优化的一部分是如何减少访存，多用本地存储来代替

另一方面这也是因为计算本身是十分简单的点积，不太可能去做出更大的优化

3.1.2 循环顺序设计

逐层去观察循环的嵌套顺序，发现是

outC-->H-->W--->inC-->kernelH-->kernelW

这样的计算顺序不一定是最优化的，笔者也没有进行详细的计算论证，但是这样的计算顺序是出于以下角度考虑

① 多通道卷积结果的维度/通道数/feature map数就是我们的outC，是我们最终要写回的out矩阵的维度，将其放在最外层循环，作用是：

一次循环内完成这个out channel中的所有计算，再接着进行下一个outC的计算
由于out数据是在一维数组中存储，且为nchw格式，那么不同outC中的数据跨度其实是很大的，连续的完成一个outC的内容可以更好的利用局部性原理
个人理解逐个outC的计算是很是一种比较直观和自然(方便想象与理解)的角度
串行过程中我们可以使用尽量少的中间变量去维护中间结果，如果你先遍历inC后遍历outC的话，其实你是需要维护outC个中间变量的

这样的顺序也是在后面做并行化改造过程中逐渐发现的一个较为合理的顺序，我们可以在后文中更加直观的感受到这样设计的优势

② 出于nchw布局的涉及，H W的顺序是基本固定的，当然你也可以先W后H，不过一般是行主序存储.. 还是先H比较快一些

③ inC为何出现在H W之后？请回顾多通道卷积的过程，一个feature map的值是由多个inC与kernel分别点击累加形成的，如果你将inC放置在H W之前的话，在下方的代码中，你是不是就需要设置height×width个中间变量来存储这里的val值呢？

in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);
val += in[in_pos] * kernel[kernel_pos];

将inC放置在H W之后，是相当于在一个outC上进行计算，对不同inC同样的位置分别计算得到了val的准确值，最终写回，这样在串行的版本中，我们只需要一个float即可存储好中间结果来避免空间的浪费！

TIPS：注意上方对于下标的计算，我们以两个位序举例说明

in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);

nchw的数据布局格式下，这里是默认n为1的，注意本文所有的实现都是建立在n假设为1的情况，其实n为更大值也不是很有意义，这样的布局下，下一张图像在计算意义上是没有任何差别的，无非是你将数据的起始地址跳过一大部分，切到下一张图像

说回这个式子，其中ic为in channel，inH inW分别是输入矩阵的高度与宽度，后面宏定义的OFFSET其实就是简略写法，你也可以写成(i+ii)*inW + j + jj

in_pos的含义是在当前循环变量下输入矩阵的位置

同理，out_pos的计算是一样的

out_pos = oc * outH * outW + OFFSET(i, j, outW);

ii和jj是相对于卷积核的相对位置循环变量，输出位置是用不到他们的

进行并行化改造

其实当你把串行版本设计明白后，你对于并行化改造的想法也差不多有个七七八八了

主要是出于以下三个角度去设计并优化的

① 尽量减少访存次数（当然不是不访问），尤其是减少访问GMEM的次数，善用SMEM与register

（对于GMEM SMEM和register等访存层次相关知识不熟的读者可以去了解一下CUDA的存储层次）

② 此外要划分明确各个线程要负责的任务区域和他的行为应达到的效果，做好下标计算

③ 计算行为是很快的，我们要尽可能去掩盖访存延迟，让线程去火力全开计算（预取prefetch）

下面的章节都是在并行化改造过程中的一些细节，代码其实是一版版写出来的，这里是对最终版本进行说明

（所谓的一版版就是划分出不同块，分别测试是否与预期一致，再去完成下面的块）

3.2 线程任务均分

这部分其实是源于 @有了琦琦的棍子在GMEM讲解中的数据转移部分，基本算是照抄了

十分感谢前辈，不过还不知道这种方法的确切名字，目前暂时称为均分，其实思想是很朴素的

我们的block设计的是8*4的大小，对应32个线程，但是涉及到in矩阵的数据可不只是32个元素，那么

我们需要尽可能地平均分配任务给线程，保证每个线程承担差不多的任务量来达到更好的平均性能

差不多是因为，不太可能都是整除的情况

这部分主要通过图示讲解，自己设计的过程中大多是通过纸笔演算确定下标的

首先确定一些变量，注意CUDA的笛卡尔坐标系和笔者的行号row和列号col的区别

int block_row = blockIdx.y;
int block_col = blockIdx.x;
int thread_row = threadIdx.y, thread_col = threadIdx.x;
int tid = thread_row * threadW + thread_col;

由于要重复使用inC内的数据，我们肯定是要开一个SMEM去存储这部分数据的，那么就有一个GMEM->SMEM的数据转移过程，以8×4的block和3×3的kernel为例，我们可以得到如下的景象

其中橙色部分是我们的block，一个tid(thread id)是一个线程，也是block中的一个位置，也是outC中的一个位置

那么白色部分就是我们在block范围之外但会用到的数据，这部分数据可以看到像两条网格

那么我们怎么把这些数据从GMEM转移到SMEM呢，首先我们考虑（以下部分为自己笨拙的思考过程）

方案① 边缘线程负责白色区域

橙色为仅负责自己的位置，紫色负责3个位置，红色负责9个

看起来是不是好像也还行，只要我们通过thread_row和thread_col判断一下当前进程是否在边缘

对这些进程进行单独的编码就可以了，不过在写代码前可以先算一笔账

这个网格共有10×6=60个元素，我们有32个线程，那么最好的情况下，是每个线程负责

60/32=1.875个元素，也就是花费1.875个单位时间（这里的单位时间是抽象概念，假定为每个线程处理每个元素的时间）

那么可以看一下这种划分方式下，每个线程平均负责的元素为

后面的项是权重，前面的项如说明这个线程处理9个线程，那么花费的时间应当是9倍，所以性能应当是九分之一（相当于只处理一个元素的线程），且线程是warp调度的，32个线程里面有这么一个拖后腿分子，想必并行情况下整体花费时间是取决于这个31号线程的

这个方案的效率是理想情况的一半都不到，说明这种方案是不太可行的，写出来效果也不一定好呢，换！

方案② 平均划分

其实笔者也想过一些其他奇怪的方法，但是感觉平均思想似乎是最佳的，那么何不一步到胃呢？

我们先来定义一些变量，后面再来逐步解释

// 分块边界 boundary是限制正常范围 edge是需要补的范围
int row_boundary = outH / BLOCK_HEIGHT - 1,col_boundary = outW / BLOCK_WIDTH - 1;
int row_edge = outH % BLOCK_HEIGHT, col_edge = outW % BLOCK_WIDTH;
···
int single_trans_ele_num = 4;                               // 线程一次转移的数据数
int cur_in_block_height = BLOCK_HEIGHT + KERNEL_HEIGHT - 1, // 读入in的block heightcur_in_block_width = BLOCK_WIDTH + KERNEL_WIDTH - 1,    // 读入in的block widthin_tile_thread_per_row,                                 // 以tile为单位转移数据，一行需要的thread数in_tile_row_start,                                      // tile的行起始位置in_tile_col,                                            // tile的列in_tile_row_stride;                                     // tile行跨度// 修正边缘block尺寸
if (block_row == row_boundary)
{cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;

3.2.1 “block”设计与修正

不要急着头大，我们逐个说明，首先看顶头部分的变量，是关于限制范围的

因为我们要首先确定一个block内的线程要负责多少元素呢，因此需要界定这样的范围

我们前面只提到了block涉及到的in范围是扩大了一圈的，其实你的in矩阵相对于out矩阵也是多了一圈的

当多的这么一圈不能构成新的block时，那么注定我们的block网格是不能覆盖到out矩阵的！

我们还是上图比较直观

咱们的block网格只有16×20这么大，out矩阵有18×22这么大，明显可以看到蓝色的两条

是不足以构成新的block的，那么还有红色的部分，就是in矩阵的大小了，可以看到有20×24这么大

而我们的block是建立在out矩阵上的，所以我们起码也要覆盖到蓝色矩阵的所有范围吧

那么在不修改block尺寸的情况下，最简单的方法就是人为地去修正这些特定block的大小啦

修正后的block应该是这个样子的

修正后的block把out全覆盖了~

怎么修正呢？无非就是利用block位序去判断并修改尺寸啦，即这两行代码

// 修正边缘block尺寸
if (block_row == row_boundary)
{cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}

结合图片，是不是这些变量的概念就清晰了起来

注意我们所有变量都是有一个in的标识，这是标注in矩阵的范围

out矩阵的划分自然是有out的标识，且步骤都是一样的，只不过需要补的范围不太一样罢了

3.2.2 线程行为指定

还有一段代码我们没有解释，是这一段（thread_num_per_block本文默认为32，没有修改）

in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;

这段我觉得是最抽象的部分也恰恰是最为精华的设计，首先要明确，是通过行里面的小片/tile作为线程处理的最小单元来进行设计的

其实变量名已经做了一部分的解释，可以大概解释为如下的含义

in_tile_thread_per_row 一行里面会有多少个tile
in_tile_row_start 当前线程负责的tile的起始行号
in_tile_col 当前线程负责的列号
in_tile_row_stride 如果还有元素要处理，那么需要跳过的行数/stride

好像不是那么的直观，我们再上一张图

左面是我们的block与in矩阵的关系，我们要把他都转移过来，且利用了fetch_float4的向量指令（也是single_trans_ele_num设置为4的原因）

以7号线程为例，当前的in_block为10×6大小，那么上面四个变量的值分别为1,7,0,32

这个例子比较简单，可以发现一行其实是有一个半的tile的，那么需要一点点小小的修正来让每个线程

读取4+2个元素，这点小小的修正我们可以看代码

那么再来一个复杂的例子，假设我们在考虑out矩阵的事情，那么一个线程负责一个元素的话

请问这种方式对嘛？

是不是直观上你感觉应该是这样的，他可以丝滑的衔接好每个元素，完成我们的分配~

那么给出我们利用这个均分思想让每个线程负责任务的代码如下，大家再想一想分配后的图像

for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;i += in_tile_row_stride)
{/*do something*/
}

浅浅一个for循环，只不过所有条件都是我们仔细设计的，循环内部就是每个线程根据这些位序

去对应的显存位置上对数据一通操作罢了

那么注意部分，线程在跨过一个stride时，这个单位是不是row？那么意味着0号线程在下次任务会踩到30号的位置！如下图所示

实际上的线程分配

这样才是正确的线程操作顺序，当然由于我们是通过CUDA并行计算的，实际上上半部分是并行的，下半部分是在0-29号线程完成了上面的任务后才进行计算的（注意他们是32个一组/warp调度上来执行的）

这样其实有个小隐患，30号和31号以及0,1号会对这两个位置上重复进行操作，如果他们的行为不一致的话

会导致我们的结果出错，本例中他们的行为是一致的，故无所谓先后

通过这样的机制，我们可以指定每个线程负责的元素位置以及个数（tile大小），灵活地应用于不同的任务！

3.3 预取机制

这部分就是很基本的数据预取，计算的效率远远大于访存，计算时读取数据进来，完成基本的运算

（复杂运算也不是一行代码可以解决的）

再把结果存到对应位置，我们发现是不是即使是计算你也需要访存，节省访存开销是十分重要的

整体的数据传输逻辑是GMEM->SMEM->register->GMEM->MEM

并没有使用到Constant Memory和Texture Memory，那么结合数据预取的机制下

整体的框架如下方伪代码所示

初始化我们所需要的所有变量并修正block规模;
分配好shared memory用于加速访存;// 预读取第一个channel的数据
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;i += in_tile_row_stride)
{把in中的数据从GMEM转到SMEM;
}
// 预读取第一个kernel的数据 可以使用很简单的读取策略 因为数据很少
if (thread_row >= 0 && thread_row < KERNEL_HEIGHT && thread_col == 0)
{把kernel的数据从GMEM转到SMEM;
}__syncthreads();// 这里oc在外ic在内的设计是为了方便写回
for (int oc = 0; oc < outC; oc++)
{for (int ic = 0; ic < inC; ic++){// i,j 是相当于当前block起始位置而言// 用ic的每个block去对oc的kernel进行计算for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height;i += out_tile_row_stride){计算当前ic与oc的结果，存到register;}// 读取下一个in channel数据 3,932,160if (ic + 1 < inC){for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;i += in_tile_row_stride){读取下一个channel的数据;}}__syncthreads();}if (oc + 1 < outC){读取下一个kernel数据;}__syncthreads();// 注意这样的循环顺序下已经完成了一个outC的计算for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height; i += out_tile_row_stride;){写回当前outC的数据;}// 预读取下一个in channel数据 需要注意这时候要从头读了for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;i += in_tile_row_stride){读取第一个channel的数据;}
}

到这里其实我们就完成了大部分内容了，整体骨架就是这样，其余就是一些细节上的下标计算问题了

3.4 一些杂项却又需要细节

3.4.1 中间结果存储设计

可以看到我们的伪代码中循环顺序是先oc再ic

可以想象一下，如果你先ic再oc的话，这样确实是我们只需要遍历一遍ic，oc多次遍历

但是我们也要考虑写回部分，写回你还需要单独再去写，理论上先ic的话会快一些

这里就不给大家放图了，读者可以自己想象一下两种计算顺序的区别

需要注意的是

线程能利用的硬件资源是有限的，一个warp共用一个SM上的寄存器，具体到每个线程大概32-255个寄存器（来源于chatGPT,不严谨，需要核实，后面gpt又说v100一个线程可以用800个..）

总之我们还是能少用就少用几个

当register存不下我们这些中间变量，就会放到local memory中

所谓的local memory是位于GMEM上的，如果发生这种情况，每次读取中间结果

你还得跑到GMEM上去访存，是非常之浪费时间的

两种循环其实需要的register数目都是oc×2（2是因为你一个线程要负责好几个位置的）

出于修正考虑，哥们儿直接开4倍，保证不会越界

3.4.2 下标计算

这部分其实，你串行算的明白，你并行就算的明白，我们举几个例子来说明一下

FETCH_FLOAT4(load_reg[0]) =FETCH_FLOAT4(in[begin_pos + OFFSET(in_tile_row_start + i, in_tile_col, inW)]);
s_in[in_tile_row_start + i][in_tile_col] = load_reg[0];
s_in[in_tile_row_start + i][in_tile_col + 1] = load_reg[1];
s_in[in_tile_row_start + i][in_tile_col + 2] = load_reg[2];
s_in[in_tile_row_start + i][in_tile_col + 3] = load_reg[3];

这里是利用向量指令去一次读取4个32位数据，s_in是开在SMEM上的，in是GMEM上的一位数据

那么可以看这个后面的下标

begin_pos 代表当前block的起始位序
OFFSET 是一个宏定义，代表行×一行元素数目
in[xxx] 下标其实就是当前block位置+block内的位置

再看一个写入中间结果的位置

temp_pos = i / out_tile_row_stride + j +oc * (cur_out_block_height / out_tile_row_stride + 1);

这里要考虑到线程是在计算它负责的第几个元素，那么就要用i / out_tile_row_stride来判断

如果处理多个元素，那你还得用j来控制一下当前是第几个元素

还要考虑到不同的oc，一个oc内负责的元素有cur_out_block_height / out_tile_row_stride +1这么多个

我们再看一个

out_pos = oc * outH * outW +block_row * BLOCK_HEIGHT * outW + block_col * BLOCK_WIDTH +OFFSET(out_tile_row_start + i, out_tile_col + j, outW);

首先略过几个oc的范围，再计算当前block的起始位置，再计算上block内的相对位置

每个下标都要明白其计算的含义，本例中有很多公共表达式没有提取出来提前计算，会影响一定性能

3.5 完整代码

完整代码已上传：

https://github.com/Pegessi/conv2d_direct

3.6 性能测试

虽然是娱乐测试，但是也严谨一点，可以发现这个代码会受channel数目影响很大

代码还有一点小bug，不过不影响你执行，大家可能会发现（亟待修复）

不同数据规模下性能在cudnn的1/10到10倍上下横跳，有空给大家测一下放个完整的图。

四、大模型手撕CUDA

在CUDA编程中实现不同矩阵运算和操作的优化技巧。

前段时间参加了一些大模型的面试，大部分都要手撕CUDA，因此也整体复习了一遍CUDA优化相关的内容，整理了一些高频题的基本写法，保存在这里也便于日后自己复习。当然，有些代码不一定是最优化解，比如GEMM，想要在面试短短的30分钟内写一个好的GEMM Kernel，是有些难度的。印象比较深刻的是，其中有一场面试2个多小时，一个小时问项目，剩下一个小时在写GEMM，说实话，如果不是事先有准备过一些，直接上手写优化版还是会有点慌。就自己的经验而言，命中率还挺高（目前还没有遇到这些题目之外的），考虑深度优化的，一般也不会让你在面试短短几十分钟手撸出来。要是遇到有面试官让你手撸一个FlashAttention，那说明，你们实在是没有缘分，还是尽早好聚好散的好，或者提前结束面试，把时间省下来，出去吃顿烧烤也不错...，附GitHub链接：

https://github.com/DefTruth/cuda-learn-note

TIPS: 文章整理为方便自己复习，不喜欢的请自动跳过哈。

01 高频面试题汇总简介

相关kernel如下。也就是不到1000行代码，建议背下来，我个人是喜欢背记，背的过程中基本就慢慢理解所有细节。当然，每个人的学习方法都不一样哈，自己觉得舒服就行。

sgemm naive, sgemm + block-tile + k-tile + vec4
sgemv k32/k128/k16 kernel
warp/block reduce sum/max, block all reduce + vec4
dot product, dot product + vec4
elementwise, elementwise + vec4
histogram, histogram + vec4
softmax, softmax + vec4 (grid level memory fence)
sigmoid, sigmoid + vec4
relu, relu + vec4
layer_norm, layer_norm + vec4
rms_norm, rms_norm + vec4
....

题内话，大模型相关的岗位，手撕CUDA的概率非常大，leetcode反而写的少，就前段时间个人的经验，基本是4:1的比例，还是建议好好复习下CUDA。当然，这些只是最简单的kernel实现，比如flash_attn，FMHA这些优化手段，就不在这篇文章里写了，面试中基本都会问到。FlashAttention系列原理详解，可以看我写的另一篇文章：

DefTruth：[FlashAttention系列][2w字] 原理详解: 从Online-Softmax到FlashAttention-1/2/FlashDecoding629(https://zhuanlan.zhihu.com/p/668888063)

02 sgemm naive, sgemm + block-tile + k-tile + vec4

#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <vector>
#include <algorithm>
#include <cuda_runtime.h>#define WARP_SIZE 32
#define INT4(value) (reinterpret_cast<int4*>(&(value))[0])
#define FLOAT4(value) (reinterpret_cast<float4*>(&(value))[0])// SGEMM: Block Tile + K Tile, with smem
// Block Tile (BM, BN) + K Tile (BK=32)
// grid((N + BN - 1) / BN, (M + BM - 1) / BM), block(BN, BM)
// a: MxK, b: KxN, c: MxN, compute: c = a * b, all row major  
__global__ void sgemm(float* a, float* b, float* c, int M, int N, int K) {// [1] Block Tile: 32x32的block处理c上一块32x32的元素计算// [2]     K Tile: 使用共享内存，并将K分块为BK大小的块constexpr int BM = 32;constexpr int BN = 32;constexpr int BK = 32;__shared__ float s_a[BM][BK], s_b[BK][BN]; int bx = blockIdx.x;int by = blockIdx.y;int tx = threadIdx.x;int ty = threadIdx.y;int tid = threadIdx.y * blockDim.x + tx; // tid within the block// load values to shared memory, 32x32 threads working together // to fetch data along the row direction of a and b both for s_a // and s_b 32x32x4x2=8KB, we use 32x32 threads within block to // load 32x32 elements from global memory to shared memory, namely, // each thread will load 1 element.int load_smem_a_m = tid / 32; // 0~31, tid / 32, tid / BM, threadIdx.yint load_smem_a_k = tid % 32; // 0~31, tid % 32, tid % BK, threadIdx.xint load_smem_b_k = tid / 32; // 0~31, tid / 32, tid / BK, threadIdx.yint load_smem_b_n = tid % 32; // 0~31, tid % 32, tid % BN, threadIdx.xint load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and cint load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and c// if (load_gmem_a_m >= M || load_gmem_b_n >= N) return;float sum = 0.f;for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {int load_gmem_a_k = bk * BK + load_smem_a_k;int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;s_a[load_smem_a_m][load_smem_a_k] = a[load_gmem_a_addr];int load_gmem_b_k = bk * BK + load_smem_b_k;int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n;s_b[load_smem_b_k][load_smem_b_n] = b[load_gmem_b_addr];__syncthreads();#pragma unrollfor (int k = 0; k < BK; ++k) {int comp_smem_a_m = load_smem_a_m;int comp_smem_b_n = load_smem_b_n;sum += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];}__syncthreads();}int store_gmem_c_m = load_gmem_a_m;int store_gmem_c_n = load_gmem_b_n;int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;c[store_gmem_c_addr] = sum;
}// SGEMM: Block Tile + Thread Tile + K Tile + Vec4, with smem
// BK:TILE_K=8 BM=BN=128
// TM=TN=8 增加计算密度 BM/TM=16 BN/TN=16
// dim3 blockDim(BN/TN, BM/TM);
// dim3 gridDim((N + BN - 1) / BN, (M + BM - 1) / BM)
__global__ void sgemm_thread_tile_vec4(float* a, float* b, float* c, int M, int N, int K) {// [1]  Block Tile: 一个16x16的block处理C上大小为128X128的一个目标块// [2] Thread Tile: 每个thread负责计算TM*TN(8*8)个元素，增加计算密度// [3]      K Tile: 将K分块，每块BK大小，迭代(K+BK-1/BK)次，//                  每次计算TM*TN个元素各自的部分乘累加// [4]   Vectorize: 减少load和store指令，使用float4constexpr int BM = 128;constexpr int BN = 128;constexpr int BK = 8; constexpr int TM = 8;constexpr int TN = 8;int bx = blockIdx.x;int by = blockIdx.y;int tx = threadIdx.x;int ty = threadIdx.y;int tid = threadIdx.y * blockDim.x + tx; // tid within the block__shared__ float s_a[BM][BK], s_b[BK][BN]; // 2*128*8*4=8KB// 0. 先计算shared memory中的索引// tid和需要加载的smem s_a[BM][BK] 之间的索引关系 BM=128 BK=8 按行读取 A行主序// 对于s_a每行8个数据，每个线程读取4个，需要2个线程；总共128行，需要128x2刚好256线程int load_smem_a_m = tid / 2; // tid/2 (128/8)*(128/8)=256 threads per block, tid/2->[0,128), BM=128 0~127int load_smem_a_k = (tid % 2 == 0) ? 0 : 4;  // (tid%2 == 0) ? 0 : 4, col of s_a 0,4// tid和需要加载的smem s_b[BK][BN] 之间的索引关系 BK=8 BN=128 按行读取 B行主序// 对于s_b每行128个数据，每个线程读4个数据，需要32个线程；总共8行，需要32x8=256个线程int load_smem_b_k = tid / 32; // tid/32, row of s_b 256/32=8 行 0~7int load_smem_b_n = (tid % 32) * 4;  // (tid % 32) * 4, col of s_b 0,4,...,124// 1. 再计算全局内存中的索引// 要加载到s_a中的元素对应到A全局内存中的行数 每个block负责出C中大小为BM*BN的块int load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and cint load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and cfloat r_c[TM][TN] = {0.0}; // 8x8// 2. 先对K进行分块，每块BK大小for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {// 加载数据到共享内存smem s_a BM*BK 128*8 vectorize float4int load_gmem_a_k = bk * BK + load_smem_a_k; // global col of aint load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;FLOAT4(s_a[load_smem_a_m][load_smem_a_k]) = FLOAT4(a[load_gmem_a_addr]);// 加载数据到共享内存smem s_b BK*BN 8*128 vectorize float4int load_gmem_b_k = bk * BK + load_smem_b_k; // global row of bint load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n; FLOAT4(s_b[load_smem_b_k][load_smem_b_n]) = FLOAT4(b[load_gmem_b_addr]); __syncthreads();#pragma unrollfor (int k = 0; k < BK; k++) {// 3. 每个线程负责计算BM*BN(12x128)中的TM*TN(8x8)个元素#pragma unrollfor (int m = 0; m < TM; m++) {#pragma unrollfor (int n = 0; n < TN; n++) {// k from 0~7，0 ~ BK, ty and tx range from 0 to 15, 16x8=128int comp_smem_a_m = ty * TM + m;  // 128*8 128/TM(8)=16 M方向 16线程int comp_smem_b_n = tx * TN + n;  // 8*128 128/TN(8)=16 N方向 16线程r_c[m][n] += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];}}}__syncthreads();}#pragma unrollfor (int m = 0; m < TM; ++m) {int store_gmem_c_m = by * BM + ty * TM + m;#pragma unrollfor (int n = 0; n < TN; n += 4) {int store_gmem_c_n = bx * BN + tx * TN + n;int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;FLOAT4(c[store_gmem_c_addr]) = FLOAT4(r_c[m][n]);}}
}

这里gemm的实现比较简单，只使用了CUDA Cores，并且只实现Block Tile + K Tile以及Block Tile + K Tile+Thread Tile+向量化的版本。主要在于如何加载gmem中的数据到smem，也就是把全局内存中的数据索引mapping到共享内存中的。核心思维：把一个block中的线程id按照线性来理解，然后把这个线性的id和全局内存索引以及共享内存索引进行匹配。比如Block Tile + K Tile的实现，block内一共32x32个Threads，需要加载到smem的数据也是32x32，那么，最简单的做法，只需要每个线程加载一个互不重复数据即可。NOTE，本文的gemm kernel修改自：紫气东来：CUDA（三）：通用矩阵乘法：从入门到熟练(https://zhuanlan.zhihu.com/p/657632577)

03 warp/block reduce sum/max

// Warp Reduce Sum
template<const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum(float val) {#pragma unrollfor (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {val += __shfl_xor_sync(0xffffffff, val, mask);}return val;
}// Warp Reduce Max
template<const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_max(float val) {#pragma unrollfor (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, mask));}return val;
}// Block reduce sum/max/min device helper for Layer/RMS Norm/Softmax etc.
// grid 1D block 1D, grid(N/128), block(128)
template<const int NUM_THREADS=128>
__device__ __forceinline__ float block_reduce_sum(float val) {// always <= 32 warps per block (limited by 1024 threads per block)constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;int warp = threadIdx.x / WARP_SIZE;int lane = threadIdx.x % WARP_SIZE;static __shared__ float shared[NUM_WARPS];val = warp_reduce_sum<WARP_SIZE>(val);if (lane == 0) shared[warp] = val;__syncthreads();val = (lane < NUM_WARPS) ? shared[lane] : 0.0f;val = warp_reduce_sum<NUM_WARPS>(val);return val;
}template<const int NUM_THREADS=128>
__device__ __forceinline__ float block_reduce_max(float val) {// always <= 32 warps per block (limited by 1024 threads per block)constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;int warp = threadIdx.x / WARP_SIZE;int lane = threadIdx.x % WARP_SIZE;static __shared__ float shared[NUM_WARPS];val = warp_reduce_max<WARP_SIZE>(val);if (lane == 0) shared[warp] = val;__syncthreads();val = (lane < NUM_WARPS) ? shared[lane] : -FLT_MAX;val = warp_reduce_max<NUM_WARPS>(val);return val;
}

warp reduce几乎已经成为大部分reduce kernel的标准写法了，比如vLLM中，就是这种经典的写法。所以，先搞懂warp reduce（也就是搞懂各种warp functions的用法），再去写其他kernel，思路就会容易很多。需要注意的是，warp函数处理的是寄存器上的数据，也就是说，此时，没必要先加载数据到smem，再进行reduce，直接加载到寄存器即可（以前犯过这个小错误...）。Warp Functions建议参考：jhang：CUDA编程入门之Warp-Level Primitives(https://zhuanlan.zhihu.com/p/572820783)

04 block all reduce + vec4

// Block All Reduce Sum
// grid(N/128), block(128)
// a: Nx1, y=sum(a)
template<const int NUM_THREADS = 128>
__global__ void block_all_reduce_sum(float* a, float* y, int N) {int tid = threadIdx.x;int idx = blockIdx.x * NUM_THREADS + tid;constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;__shared__ float reduce_smem[NUM_WARPS];// keep the data in register is enougth for warp operaion.float sum = (idx < N) ? a[idx] : 0.0f;int warp = tid / WARP_SIZE;int lane = tid % WARP_SIZE;// perform warp sync reduce.sum = warp_reduce_sum<WARP_SIZE>(sum);// warp leaders store the data to shared memory.if (lane == 0) reduce_smem[warp] = sum;__syncthreads(); // make sure the data is in shared memory.// the first warp compute the final sum.sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;if (warp == 0) sum = warp_reduce_sum<NUM_WARPS>(sum);if (tid == 0) atomicAdd(y, sum);
}// Block All Reduce Sum + float4
// grid(N/128), block(128/4)
// a: Nx1, y=sum(a)
template<const int NUM_THREADS = 128/4>
__global__ void block_all_reduce_sum_vec4(float* a, float* y, int N) {int tid = threadIdx.x;int idx = (blockIdx.x * NUM_THREADS + tid) * 4;constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;__shared__ float reduce_smem[NUM_WARPS];float4 reg_a = FLOAT4(a[idx]);// keep the data in register is enougth for warp operaion.float sum = (idx < N) ? (reg_a.x + reg_a.y + reg_a.z + reg_a.w) : 0.0f;int warp = tid / WARP_SIZE;int lane = tid % WARP_SIZE;// perform warp sync reduce.sum = warp_reduce_sum<WARP_SIZE>(sum);// warp leaders store the data to shared memory.if (lane == 0) reduce_smem[warp] = sum;__syncthreads(); // make sure the data is in shared memory.// the first warp compute the final sum.sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;if (warp == 0) sum = warp_reduce_sum<NUM_WARPS>(sum);if (tid == 0) atomicAdd(y, sum);
}

block all reduce是在warp reduce的基础上进行的，reduce_smem这部分的共享内存申请无法避免，这是用来同步每个warp之间得到局部结果。注意，最后，还需要atomicAdd做一个block级别的原子操作，以得到全局的和。float4向量化优化访存，可以减缓WarpScheduler发送指令的压力。

05 sgemv k32/k128/k16 kernel

// SGEMV: Warp SGEMV K32
// 假设K为32的倍数，每个warp负责一行
// grid(M/4), block(32,4) blockDim.x=32=K, blockDim.y=4
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
__global__ void sgemv_k32(float* a, float* x, float* y, int M, int K) {int tx = threadIdx.x;         // 0~31int ty = threadIdx.y;         // 0~4int bx = blockIdx.x;          // 0~M/4int lane = tx % WARP_SIZE;    // 0~31int m = bx * blockDim.y + ty; // (0~M/4) * 4 + (0~3)if (m < M) {float sum = 0.0f;int NUM_WARPS = (K + WARP_SIZE - 1) / WARP_SIZE;#pragma unrollfor (int w = 0; w < NUM_WARPS; ++w) {// 若NUM_WARPS>=2，先将当前行的数据累加到第一个warp中int k = w * WARP_SIZE + lane;sum += a[m * K + k] * x[k];}sum = warp_reduce_sum<WARP_SIZE>(sum);if (lane == 0) y[m] = sum;}
}// SGEMV: Warp SGEMV K128 + Vec4
// 假设K为128的倍数 float4
// grid(M/4), block(32,4) blockDim.x=32=K, blockDim.y=4
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
__global__ void sgemv_k128(float* a, float* x, float* y, int M, int K) {// 每个线程负责4个元素，一个warp覆盖128个元素int tx = threadIdx.x;         // 0~31int ty = threadIdx.y;         // 0~3int bx = blockIdx.x;          // 0~M/4int lane = tx % WARP_SIZE;    // 0~31int m = blockDim.y * bx + ty; // (0~M/4) * 4 + (0~3)if (m < M) {float sum = 0.0f;// process 4*WARP_SIZE elements per warp.int NUM_WARPS = (((K + WARP_SIZE - 1) / WARP_SIZE) + 4 - 1) / 4;#pragma unrollfor (int w = 0; w < NUM_WARPS; ++w) {int k = (w * WARP_SIZE + lane) * 4;float4 reg_x = FLOAT4(x[k]);float4 reg_a = FLOAT4(a[m * K + k]);sum += (reg_a.x * reg_x.x + reg_a.y * reg_x.y + reg_a.z * reg_x.z + reg_a.w * reg_x.w);}sum = warp_reduce_sum<WARP_SIZE>(sum);if(lane == 0) y[m] = sum;}
}// SGEMV: Warp SGEMV K16
// 假设K为16 < 32,每个warp负责2行，每行有16个元素
// NUM_THREADS=128, NUM_WARPS=NUM_THREADS/WARP_SIZE;
// NUM_ROWS=NUM_WARPS * ROW_PER_WARP, grid(M/NUM_ROWS), block(32,NUM_WARPS)
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
template<const int ROW_PER_WARP = 2> 
__global__ void sgemv_k16(float* A, float* x, float* y, int M, int K) {constexpr int K_WARP_SIZE = (WARP_SIZE + ROW_PER_WARP - 1) / ROW_PER_WARP;int tx = threadIdx.x;       // 0~31int ty = threadIdx.y;       // 0~NUM_WARPSint bx = blockIdx.x;        // 0~M/NUM_ROWS (NUM_ROWS=NUM_WARPS * ROW_PER_WARP)int lane = tx % WARP_SIZE;  // 0~31int k = lane % K_WARP_SIZE; // 0~15// gloabl row of a: MxK and y:Mx1, blockDim.y=NUM_WARPSint m = (blockDim.y * bx + ty) * ROW_PER_WARP + lane / K_WARP_SIZE;if (m < M) {float sum = A[m * K + k] * x[k];sum = warp_reduce_sum<K_WARP_SIZE>(sum);// 注意是k == 0，而不是lane == 0if(k == 0) y[m] = sum; }
}

估计有些大佬倒立都能写sgemv的各种优化版了，核心思路其实也是基于warp reduce，考虑K的不同情况进行优化。本文的sgemv kernel修改自：有了琦琦的棍子：深入浅出GPU优化系列：gemv优化(https://zhuanlan.zhihu.com/p/494144694)

06 dot product, dot product + vec4

// Dot Product
// grid(N/128), block(128)
// a: Nx1, b: Nx1, y=sum(elementwise_mul(a,b))
template<const int NUM_THREADS = 128>
__global__ void dot(float* a, float* b, float* y, int N) {int tid = threadIdx.x;int idx = blockIdx.x * NUM_THREADS + tid;constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;__shared__ float reduce_smem[NUM_WARPS];// keep the data in register is enougth for warp operaion.float prod = (idx < N) ? a[idx] * b[idx] : 0.0f;int warp = tid / WARP_SIZE;int lane = tid % WARP_SIZE;// perform warp sync reduce.prod = warp_reduce_sum<WARP_SIZE>(prod);// warp leaders store the data to shared memory.if (lane == 0) reduce_smem[warp] = prod;__syncthreads(); // make sure the data is in shared memory.// the first warp compute the final sum.prod = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;if (warp == 0) prod = warp_reduce_sum<NUM_WARPS>(prod);if (tid == 0) atomicAdd(y, prod);
}// Dot Product + Vec4
// grid(N/128), block(128/4)
// a: Nx1, b: Nx1, y=sum(elementwise_mul(a,b))
template<const int NUM_THREADS = 128/4>
__global__ void dot_vec4(float* a, float* b, float* y, int N) {int tid = threadIdx.x;int idx = (blockIdx.x * NUM_THREADS + tid) * 4;constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;__shared__ float reduce_smem[NUM_WARPS];float4 reg_a = FLOAT4(a[idx]);float4 reg_b = FLOAT4(b[idx]);float prod = (idx < N) ? (reg_a.x * reg_b.x + reg_a.y * reg_b.y + reg_a.z * reg_b.z + reg_a.w * reg_b.w) : 0.0f;int warp = tid / WARP_SIZE;int lane = tid % WARP_SIZE;// perform warp sync reduce.prod = warp_reduce_sum<WARP_SIZE>(prod);// warp leaders store the data to shared memory.if (lane == 0) reduce_smem[warp] = prod;__syncthreads(); // make sure the data is in shared memory.// the first warp compute the final sum.prod = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;if (warp == 0) prod = warp_reduce_sum<NUM_WARPS>(prod);if (tid == 0) atomicAdd(y, prod);
}

dot product kernel的核心就是block reduce，不多说了。

07 elementwise, elementwise + vec4

// ElementWise Add  
// grid(N/128), block(128)
// a: Nx1, b: Nx1, c: Nx1, c = elementwise_add(a, b)
__global__ void elementwise_add(float* a, float* b, float* c, int N) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < N) c[idx] = a[idx] + b[idx];
}// ElementWise Add + Vec4
// grid(N/128), block(128/4)
// a: Nx1, b: Nx1, c: Nx1, c = elementwise_add(a, b)
__global__ void elementwise_add_vec4(float* a, float* b, float* c, int N) {int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);if (idx < N) {float4 reg_a = FLOAT4(a[idx]);float4 reg_b = FLOAT4(b[idx]);float4 reg_c;reg_c.x = reg_a.x + reg_b.x;reg_c.y = reg_a.y + reg_b.y;reg_c.z = reg_a.z + reg_b.z;reg_c.w = reg_a.w + reg_b.w;FLOAT4(c[idx]) = reg_c;}
}

elementwise可以考虑加点向量化进行访存优化。

08 histogram, histogram + vec4

// Histogram
// grid(N/128), block(128)
// a: Nx1, y: count histogram
__global__ void histogram(int* a, int* y, int N) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < N) atomicAdd(&(y[a[idx]]), 1);
}// Histogram + Vec4
// grid(N/128), block(128/4)
// a: Nx1, y: count histogram
__global__ void histogram_vec4(int* a, int* y, int N) {int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);if (idx < N) {int4 reg_a = INT4(a[idx]);atomicAdd(&(y[reg_a.x]), 1);atomicAdd(&(y[reg_a.y]), 1);atomicAdd(&(y[reg_a.z]), 1);atomicAdd(&(y[reg_a.w]), 1);}
}

统计频数直方图，很简单，两行代码搞定。

09 softmax, softmax + vec4 (grid level memory fence)

// Softmax x: N, y: N
// grid(N/128), block(K=128)
template<const int NUM_THREADS = 128>
__global__ void softmax(float* x, float* y, float* total, int N) {const int tid = threadIdx.x;const int idx = blockIdx.x * blockDim.x + tid; constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;__shared__ float reduce_smem[NUM_WARPS];float sum = (idx < N) ? expf(x[idx]) : 0.0f;int warp = tid / WARP_SIZE;int lane = tid % WARP_SIZE;sum = warp_reduce_sum<WARP_SIZE>(sum);if (lane == 0) reduce_smem[warp] = sum;__syncthreads();// compute the final sum in each warpsum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;sum = warp_reduce_sum<NUM_WARPS>(sum); // sum(e^x_0,...,e^x_n-1)// get the total sum of all blocks.if (tid == 0) atomicAdd(total, sum);__threadfence(); // grid level memory fence 注意这里需要网格级别的内存同步// e^x_i/sum(e^x_0,...,e^x_n-1) if (idx < N) y[idx] = block_smem[tid] / (*total); 
}// Softmax x: N, y: N
// grid(N/128), block(K=128)
template<const int NUM_THREADS = 128>
__global__ void softmax_v2(float* x, float* y, float* total, int N) {const int tid = threadIdx.x;const int idx = blockIdx.x * blockDim.x + tid; float exp_val = (idx < N) ? expf(x[idx]) : 0.0f;float sum = block_reduce_sum<NUM_THREADS>(exp_val);// get the total sum of all blocks.if (tid == 0) atomicAdd(total, sum);__threadfence(); // grid level memory fence  注意这里需要网格级别的内存同步// e^x_i/sum(e^x_0,...,e^x_n-1) if (idx < N) y[idx] = exp_val / (*total); 
}// Softmax Vec4 x: N, y: N
// grid(N/128), block(128/4)
template<const int NUM_THREADS = 128/4>
__global__ void softmax_v2_vec4(float* x, float* y, float* total, int N) {const int tid = threadIdx.x;const int idx = (blockIdx.x * blockDim.x + tid) * 4; float4 reg_x = FLOAT4(x[idx]);float4 reg_exp;reg_exp.x = (idx < N) ? expf(reg_x.x) : 0.0f;reg_exp.y = (idx < N) ? expf(reg_x.y) : 0.0f;reg_exp.z = (idx < N) ? expf(reg_x.z) : 0.0f;reg_exp.w = (idx < N) ? expf(reg_x.w) : 0.0f;float exp_val = (reg_exp.x + reg_exp.y + reg_exp.z + reg_exp.w);float sum = block_reduce_sum<NUM_THREADS>(exp_val);// get the total sum of all blocks.if (tid == 0) atomicAdd(total, sum);__threadfence(); // grid level memory fence  注意这里需要网格级别的内存同步// e^x_i/sum(e^x_0,...,e^x_n-1) if (idx < N) {float4 reg_y;reg_y.x = reg_exp.x / (*total);reg_y.y = reg_exp.y / (*total);reg_y.z = reg_exp.z / (*total);reg_y.w = reg_exp.w / (*total);FLOAT4(y[idx]) = reg_y; }
}

softmax稍微要注意的就是内存同步的问题，这里，你需要做一个网格级别的同步，而不能仅仅是block级别，否则拿不到全局的exp sum作为分母项。因此使用 __threadfence 这个网格及内存同步操作。不过效率我还没测过，实在要高效的话，可能得整成FA2那样的 1-pass + online softmax的实现。不过，如果是面试的话，就不要太为难自己了... ，但是FA1/FA2的论文很经典，强烈建议多读几遍。

10 sigmoid, sigmoid + vec4

// Sigmoid x: N, y: N y=1/(1+exp(-x))
// grid(N/128), block(K=128) 
__global__ void sigmoid(float* x, float* y, int N) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < N) y[idx] = 1.0f / (1.0f + expf(-x[idx]));
}// Sigmoid x: N, y: N y=1/(1+exp(-x)) Vec4
// grid(N/128), block(128/4)
__global__ void sigmoid_vec4(float* x, float* y, int N) {int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;if (idx < N) {float4 reg_x = FLOAT4(x[idx]);float4 reg_y;reg_y.x = 1.0f / (1.0f + expf(-reg_x.x));reg_y.y = 1.0f / (1.0f + expf(-reg_x.y));reg_y.z = 1.0f / (1.0f + expf(-reg_x.z));reg_y.w = 1.0f / (1.0f + expf(-reg_x.w));FLOAT4(y[idx]) = reg_y;}
}

11 relu, relu + vec4

// Relu x: N, y: N y=max(0,x)
// grid(N/128), block(K=128) 
__global__ void relu(float* x, float* y, int N) {int idx = blockIdx.x * blockDim.x + threadIdx.x;if (idx < N) y[idx] = fmaxf(0.0f, x[idx]);
}// Relu x: N, y: N y=max(0,x) Vec4
// grid(N/128/4), block(128/4) 
__global__ void relu_vec4(float* x, float* y, int N) {int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;if (idx < N) {float4 reg_x = FLOAT4(x[idx]);float4 reg_y;reg_y.x = fmaxf(0.0f, reg_x.x);reg_y.y = fmaxf(0.0f, reg_x.y);reg_y.z = fmaxf(0.0f, reg_x.z);reg_y.w = fmaxf(0.0f, reg_x.w);FLOAT4(y[idx]) = reg_y;}
}

12 layer_norm, layer_norm + vec4

// Layer Norm: x: NxK(K=128<1024), y': NxK, y'=x-mean(x)/std(x) each row
// mean(x) = sum(x)/K, 1/std(x) = rsqrtf( sum( (x-mean(x))^2 )/K ) each row
// grid(N*K/K), block(K<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g + b (g: scale, b: bias)
template<const int NUM_THREADS=128>
__global__ void layer_norm(float* x, float* y, float g, float b, int N, int K) {int tid = threadIdx.x; // 0..K-1int bid = blockIdx.x; // 0..N-1int idx = bid * blockDim.x + threadIdx.x;const float epsilon = 1e-5f;__shared__ float s_mean; // shared within block__shared__ float s_variance; // shared within blockfloat value = (idx < N * K) ? x[idx] : 0.0f; // load once onlyfloat sum = block_reduce_sum<NUM_THREADS>(value);if (tid == 0) s_mean = sum / (float) K;// wait for s_mean in shared memory to be ready for all threads__syncthreads();float variance = (value - s_mean) * (value - s_mean);variance = block_reduce_sum<NUM_THREADS>(variance);if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);// wait for s_variance in shared memory to be ready for all threads__syncthreads();if (idx < N * K) y[idx] = ((value - s_mean) * s_variance) * g + b;
}// Layer Norm Vec4: x: NxK(K=128<1024), y': NxK, y'=x-mean(x)/std(x) each row
// mean(x) = sum(x)/K, 1/std(x) = rsqrtf( sum( (x-mean(x))^2 )/K ) each row
// grid(N*K/K), block(K/4<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g + b (g: scale, b: bias)
template<const int NUM_THREADS=128/4>
__global__ void layer_norm_vec4(float* x, float* y, float g, float b, int N, int K) {int tid = threadIdx.x; // 0..K-1int bid = blockIdx.x; // 0..N-1int idx = (bid * blockDim.x + threadIdx.x) * 4;const float epsilon = 1e-5f;__shared__ float s_mean; // shared within block__shared__ float s_variance; // shared within blockfloat4 reg_x = FLOAT4(x[idx])float value = (idx < N * K) ? (reg_x.x + reg_x.y + reg_x.z + reg_x.w) : 0.0f;float sum = block_reduce_sum<NUM_THREADS>(value);if (tid == 0) s_mean = sum / (float) K;// wait for s_mean in shared memory to be ready for all threads__syncthreads();float4 reg_x_hat;reg_x_hat.x = reg_x.x - s_mean;reg_x_hat.y = reg_x.y - s_mean;reg_x_hat.z = reg_x.z - s_mean;reg_x_hat.w = reg_x.w - s_mean;float variance = reg_x_hat.x * reg_x_hat.x + reg_x_hat.y * reg_x_hat.y + reg_x_hat.z * reg_x_hat.z + reg_x_hat.w * reg_x_hat.w;variance = block_reduce_sum<NUM_THREADS>(variance);if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);// wait for s_variance in shared memory to be ready for all threads__syncthreads();float4 reg_y;reg_y.x = reg_x_hat.x * s_variance * g + b;reg_y.y = reg_x_hat.y * s_variance * g + b;reg_y.z = reg_x_hat.z * s_variance * g + b;reg_y.w = reg_x_hat.w * s_variance * g + b;if (idx < N * K) FLOAT4(y[idx]) = reg_y;
}

layer norm实现的核心同样也是block reduce和warp reduce，然后再整点向量化...

13 rms_norm, rms_norm + vec4

// RMS Norm: x: NxK(K=128<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template<const int NUM_THREADS=128>
__global__ void rms_norm(float* x, float* y, float g, int N, int K) {int tid = threadIdx.x; // 0..K-1int bid = blockIdx.x; // 0..N-1int idx = bid * blockDim.x + threadIdx.x;const float epsilon = 1e-5f;__shared__ float s_variance; // shared within blockfloat value = (idx < N * K) ? x[idx] : 0.0f; // load once onlyfloat variance = value * value;variance = block_reduce_sum<NUM_THREADS>(variance);if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);// wait for s_variance in shared memory to be ready for all threads__syncthreads(); if (idx < N * K) y[idx] = (value * s_variance) * g;
}// RMS Norm Vec4: x: NxK(K=128<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K/4<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template<const int NUM_THREADS=128/4>
__global__ void rms_norm_vec4(float* x, float* y, float g, int N, int K) {int tid = threadIdx.x; // 0..K-1int bid = blockIdx.x; // 0..N-1int idx = (bid * blockDim.x + threadIdx.x) * 4;const float epsilon = 1e-5f;__shared__ float s_variance; // shared within blockfloat4 reg_x = FLOAT4(x[idx]);float variance = (idx < N * K) ? (reg_x.x * reg_x.x + reg_x.y * reg_x.y + reg_x.z * reg_x.z + reg_x.w * reg_x.w) : 0.0f;variance = block_reduce_sum<NUM_THREADS>(variance);if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);// wait for s_variance in shared memory to be ready for all threads__syncthreads(); float4 reg_y;reg_y.x = reg_x.x * s_variance * g;reg_y.y = reg_x.y * s_variance * g;reg_y.z = reg_x.z * s_variance * g;reg_y.w = reg_x.w * s_variance * g;if (idx < N * K) FLOAT4(y[idx]) = reg_y;
}

rms norm实现的核心同样也是block reduce和warp reduce...，然后再加点float4向量化什么的。

14 NMS

struct Box {float x1, y1, x2, y2, score;float area() const {return (std::abs(x2 - x1 + 1)) * std::abs(y2 - y1 + 1); }float iou_of(const Box& other) const{float inner_x1 = x1 > other.x1 ? x1 : other.x1;float inner_y1 = y1 > other.y1 ? y1 : other.y1;float inner_x2 = x2 < other.x2 ? x2 : other.x2;float inner_y2 = y2 < other.y2 ? y2 : other.y2;float inner_h = inner_y2 - inner_y1 + 1.0f;float inner_w = inner_x2 - inner_x1 + 1.0f;float inner_area = inner_h * inner_w;return (inner_area / (area() + tbox.area() - inner_area));}
}
void hard_nms(std::vector<Box> &input, std::vector<Box> &output, float iou_threshold){if (input.empty()) return;std::sort(input.begin(), input.end(),[](Box& a, Box& b) { return a.score > b.score; });int box_num = input.size();std::vector<int> merged(box_num, 0);for (int i = 0; i < box_num; ++i) {if (merged[i]) continue;merged[i] = 1;for (int j = i + 1; j < box_num; ++j) {if (merged[j]) continue;float iou = input[i].iou_of(input[j]);if (iou > iou_threshold) merged[j] = 1;}output.push_back(input[i]);}
}

CV相关的经常会要手撕NMS，也记录下。

15 总结

可以发现，大部分kernel的基本写法都是依赖warp reduce和block reduce的，基本上只要熟练应用warp functions各种场景的写法，应该问题不大；softmax需要考虑网格级同步的问题，或者online softmax以及FlashAttention；sgemm的优化是个很大的课题，不是案例中写的这么简单，但是入门的话，基本就是tiling的思想以及如何做索引之间的mapping；sgemv的优化则主要考虑K不同的值（因为M为1了），比如K=16,64,128等情况下，如何按照warp来处理；relu、sigmoid等都是elementwise的操作，很好实现，可以再考虑加点向量化优化访存；layer norm和rms norm在数学上其实也是挺清晰简单的，落实到cuda kernel时，只要按照逐个token来处理，headdim没有超过1024的情况下（一个block最多可以放1024个threads），可以放到一个block处理，这样并行化就很好写。当然，核心还是warp reduce和block reduce；NMS是乱入的，没有CUDA版本，别问了...

最后，附GitHub repo:

https://github.com/DefTruth/cuda-learn-note/

五、基于CUDA的高效排序算法（Sort on GPU）

一、前言

排序（sort），是计算机内经常进行的一种操作，其目的是将一组“无序”的记录序列调整为“有序”的记录序列，是各领域中广泛应用的核心算法之一。受限于算法本身的限制（所有基于比较的排序算法复杂度下限为O(n)，对于一组超大量数据单纯的cpu排序所需时间过长，通常1000万的数据排序就需要1秒左右，更大量的数据排序则需要分钟级的运行时间。因此，可以通过并行排序操作来进一步加速排序算法。

二、基于cpu的排序算法

经过数十年的发展，目前的排序算法已有基数排序、冒泡排序、选择排序、插入排序、希尔排序、堆排序、归并排序、快速排序等。显然，我们希望在已有的高效排序算法上进行进一步的优化，因些我们忽略冒泡排序、选择排序、希尔排序这一类指数级复杂度的算法，也忽略基数排序这一类不适于于浮点型变量的算法，而在堆排序、归并排序、快速排序这些O(n)算法中，直观的我们能发现归并排序、快速排序都是基于二分的策略，所以可以自然地同时对左右两半同时进行排序；而归并的过程也可以通过直接查找每个元素在合并数组中的位置来契合并行化，但快速排序中划分的过程无法并行实现（至少我没有想到，欢迎各路大神提出好的思路）；因此归并排序是最适于并行化的排序算法。

然而，这些已有的高效排序算法通常只在数据量较大时才有着更高的效率，在小数据时反而常数更小的选择排序、插入排序等有着更出色的性能；如中的std::sort中在长度小于16或32时不再继续划分，而是调用插入排序来实现；因此在基于cuda的排序中也需要在小数据上选择另一种算法以达到更出色的性能。在此，我们选择并不常见、但复杂度O(nn)相对较低的双调排序来实现，自然是因为它能很好的并行化。

因此我们选择双调排序和归并排序来实现CUDA上的排序操作，如需详细了解这两个算法请查阅相关资料。

三、基于cuda的排序算法的实现

1. 双调排序的并行化

我们暂定一个线程块的线程数为 256（tile_size），因为双调排序中每次需要进行一次比较来判断是否需要交换，因此每一个线程块处理连续的tile_size x 2个数据（以下称为一个数据段），每个线程负责一次比较操作。因此该核函数线程块大小为tile_size，网格大小为总共的数据段数量。

#define tile_size 256u
#define log_tile_size 8uunsigned int b = n + (tile_size << 1) - 1 >> log_tile_size + 1, s = tile_size << 1;
binoticsort_gpu<<<b, s >> 1>>>(data, n);

首先先将需要的数据载入共享内存中，由于每一个线程块处理tile_size x 2个数据，所以一个线程块负责的数据段起始为“线程块索引 x tile_size x 2”；每一个线程将相隔tile_size的两个数据载入共享内存中，超过数据规模的部分直接填充浮点数的最大值，这样一个线程束访问一段连续的地址以达到更高的访存效率；还需添加一个线程块内同步函数，确保所有数据都已载入共享内存后再开始排序。

__shared__  double buffer[tile_size << 1];unsigned int d = blockIdx.x << log_tile_size + 1, i = threadIdx.x;
data += d; n -= d;
buffer[i] = i < n ? data[i] : DBL_MAX;
buffer[i + blockDim.x] = i + blockDim.x < n ? data[i + blockDim.x] : DBL_MAX;
__syncthreads();

然后需要将随机序列转换为双调序列，在这个过程中需要交换操作，因此需要先定义一个设备端的交换函数，该交换函数可定义为模板函数以支持不同类型的数据交换。

// auxiliary functions
template<typename type>__device__ void swap(type& v1, type& v2)
{type t = v1;v1 = v2;v2 = t;
}

注：浅蓝色段为升序，深蓝色段为降序

unsigned int s, t;
for (unsigned int k = 2; k <= blockDim.x; k <<= 1)for (unsigned int p = k >> 1; p; p >>= 1){s = (i & -p) << 1 | i & p - 1;t = s | p;if (s & k ? buffer[s] < buffer[t] : buffer[s] > buffer[t])swap(buffer[s], buffer[t]);__syncthreads();}

注：排序过程中的顺序与需求相关，若为升序排序则全为升序，反之则全为降序

该过程转化为代码的操作与上述合并过程中的内循环完全一致。

for (unsigned int p = blockDim.x; p; p >>= 1)
{s = (i & -p) << 1 | i & p - 1;t = s | p;if (buffer[s] > buffer[t])swap(buffer[s], buffer[t]);__syncthreads();
}

最后将该数据段的有效排序结果写回全局内存。

if (i < n)data[i] = buffer[i];
if (i + blockDim.x < n)data[i + blockDim.x] = buffer[i + blockDim.x];

最终分块双调排序的代码如下：

#include <float.h>// binotic sort
__global__ void binoticsort_gpu(double* data, unsigned int n)
{__shared__  double buffer[tile_size << 1];unsigned int d = blockIdx.x << log_tile_size + 1, i = threadIdx.x;data += d; n -= d;buffer[i] = i < n ? data[i] : DBL_MAX;buffer[i + blockDim.x] = i + blockDim.x < n ? data[i + blockDim.x] : DBL_MAX;__syncthreads();unsigned int s, t;for (unsigned int k = 2; k <= blockDim.x; k <<= 1)for (unsigned int p = k >> 1; p; p >>= 1){s = (i & -p) << 1 | i & p - 1;t = s | p;if (s & k ? buffer[s] > buffer[t] : buffer[s] < buffer[t])swap(buffer[s], buffer[t]);__syncthreads();}for (unsigned int p = blockDim.x; p; p >>= 1){s = (i & -p) << 1 | i & p - 1;t = s | p;if (buffer[s] > buffer[t])swap(buffer[s], buffer[t]);__syncthreads();}if (i < n)data[i] = buffer[i];if (i + blockDim.x < n)data[i + blockDim.x] = buffer[i + blockDim.x];
}

经过该分块双调排序后，规模为n的数据被划分为多个长为 tile_size x 2的数据段，每个数据段内的 tile_size x 2个数据已经被成功排序，下面只需要逐步将这些有序数据段合并即可。

2. 归并操作的并行化

有序数据段合并操作即像归并排序中的归并一样，但是需要对其做一定的修改以更好地适用于并行化。通常的合并操作需从头逐个比较两个数组中的元素，取较小者置于结果数组中——但是这个循环是具有数据依赖性的，故这种相互之间不独立的操作无法进行并行化处理。因此我们需要用另一种方法来进行合并操作，我们直接计算两个待合并数组中每个元素在结果数组中应当存在的位置，然后把该元素置于该位置即可。

对于第一个数据段中的每个元素，计算第二个数据段中小于该元素的个数，再加上该元素在第一个数据段中的索引即为该元素在结果数据段中的索引；对于第二个数据段类似，但要计算第一个数据段中小于等于其每个元素的个数。

显然对于其中一待合并数组中比某元素小的元素个数必然是该元素的索引，因此只需找到另一待合并数组中小于（或小于等于）该元素的元素个数即可，注意这儿必须要一个数组用小于而另一个数组用小于等于，否则两个相同的数字就会落在同一位置。在此直接引用中的std::lower_bound和std::upper_bound中的算法来实现，只是需要将其编译为设备端函数。

// auxiliary functions
template<typename type>__device__ unsigned int lower_bound(const type* data, unsigned int n, const type& value)
{unsigned int len = n, half;const type* start = data, * mid;while (len){half = len >> 1;mid = data + half;if (*mid < value){data = mid + 1;len = len - half - 1;}elselen = half;}return data - start;
}
template<typename type>__device__ unsigned int upper_bound(const type* data, unsigned int n, const type& value)
{unsigned int len = n, half;const type* start = data, *mid;while (len){half = len >> 1;mid = data + half;if (value < *mid)len = half;else{data = mid + 1;len = len - half - 1;}}return data - start;
}

mergedirect_gpu<<<dim3(s >> log_tile_size, std::min(b, 32768u), b + 32767 >> 15), tile_size>>>(data, n, s, tmp);

// merge directly
__global__ void mergedirect_gpu(const double* data, unsigned int n, unsigned int s, double* result)
{unsigned int d1 = (blockIdx.z * gridDim.y + blockIdx.y) * s << 1, d2 = d1 + s, i = d1 + blockIdx.x * blockDim.x + threadIdx.x;if (i < n)result[d2 < n ? i + lower_bound(data + d2, min(s, n - d2), data[i]) : i] = data[i];if (s + i < n)result[i + upper_bound(data + d1, s, data[s + i])] = data[s + i];
}

然后将这两个函数整合成一个排序函数即可，该函数接受两个参数：指向数据首地址的指针和数据规模，返回指向已排序数据首地址的指针，其中的指针均指向设备端。该函数首先需要申请一块与输入数据等长的临时数组用于存放合并的结果，然后先进行一次双调排序将输入数据分块排序为多个长为 tile_size x 2的有序数据段，再来一个循环将这些有序数据段逐轮合并直至数据段长度大于数据规模，最后释放临时数组空间并返回结果。

double* sortdirect_gpu(double* data, unsigned int n)
{double* tmp;cudaMalloc(&tmp, sizeof(double) * n);unsigned int b = n + (tile_size << 1) - 1 >> log_tile_size + 1, s = tile_size << 1;binoticsort_gpu<<<b, s >> 1>>>(data, n);for (b = b + 1 >> 1; s < n; s <<= 1, b = b + 1 >> 1){mergedirect_gpu<<<dim3(s >> log_tile_size, std::min(b, 32768u), b + 32767 >> 15), tile_size>>>(data, n, s, tmp);std::swap(data, tmp);}cudaDeviceSynchronize();cudaFree(tmp);return data;
}

至此，我们已初步实现了基于CUDA的排序算法，我们选择了不同的数据规模在RTX A6000上对它的性能进行测试，测试代码如下：

// test.cu
#include <cuda.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <iostream>
#include <chrono>
#include <algorithm>int main(int argc, char** argv)
{cudaSetDevice(0);cudaFree(0);unsigned int test[] = { 1025u, 65537u, 1048577u, 16777217u, 134217729u, 1073741825u };const unsigned int runtime = 3;std::chrono::high_resolution_clock::time_point start, end;unsigned long long elasped[3] = { 0ull };for (unsigned int t = 0; t < sizeof(test) / sizeof(unsigned int); ++t){unsigned int n = test[t];double* original = new double[n];double* data = new double[n];elasped[0] = elasped[1] = elasped[2] = 0ull;srand(time(nullptr));for (int i = 0; i < runtime; ++i){for (unsigned int i = 0; i < n; ++i)original[i] = data[i] = (rand() * 2.0 / RAND_MAX - 1.0) * rand();// cpu sort 1start = std::chrono::high_resolution_clock::now();qsort(data, n, sizeof(double), [](const void* left, const void* right)->int {double tmp = *(double*)left - *(double*)right; if (tmp > 0) return 1; else if (tmp < 0) return -1; else return 0; });end = std::chrono::high_resolution_clock::now();elasped[0] += std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();// cpu sort 2memcpy(data, original, sizeof(double) * n);start = std::chrono::high_resolution_clock::now();std::sort(data, data + n);end = std::chrono::high_resolution_clock::now();elasped[1] += std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();// gpu sortdouble* pd_data;cudaMalloc(&pd_data, sizeof(double) * n);cudaMemcpy(pd_data, original, sizeof(double) * n, cudaMemcpyHostToDevice);start = std::chrono::high_resolution_clock::now();pd_data = sortdirect_gpu(pd_data, n);end = std::chrono::high_resolution_clock::now();elasped[2] += std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();cudaMemcpy(data, pd_data, sizeof(double) * n, cudaMemcpyDeviceToHost);cudaFree(pd_data);}printf("data number: %u\n", n);printf("cpu merge sort cost %.5f ms\n", elasped[0] / 1e6 / runtime);printf("cpu quick sort cost %.5f ms\n", elasped[1] / 1e6 / runtime);printf("gpu sort cost %.5f ms\n", elasped[2] / 1e6 / runtime);printf("-------------------------------------------\n");delete[] data;delete[] original;}return 0;
}

我们先不加任何优化进行编译

nvcc -o ./benchmark -g ./test.cu && ./benchmark

其运行结果如下（只统计排序所消耗的时间，忽略数据传输的耗时）：

然后我们开启O2优化进行编译

nvcc -O2 -o ./benchmark_O2 -g ./test.cu && ./benchmark_O2

其运行结果如下（只统计排序所消耗的时间，忽略数据传输的耗时）：

我们再加上数据传输耗时，其耗时如下：

可以发现在不计数据传输开销时，从数万到10亿的数据规模下，该基于CUDA的排序算法相较于单线程cpu排序算法而言都能获得数十乃至百倍的加速比；若算上数据传输的耗时，数据规模较小时其性能不如cpu排序算法，但在数据规模达100万及以上仍有着20倍左右的加速比（相较于更快的std::sort）。

四. 基于cuda的排序算法的改进

忽略数据传输的耗时，我们已经实现了约60倍的加速比（相较于更快的std::sort），但是仍能注意到该排序算法中合并操作是主要的耗时步骤，而其主要的耗时步骤则在查找每个元素在另一数据段中的位置，因为当数据段长度较大时，在很大的地址区间内进行二分查找是一个很费时的操作，这既不利于缓存（cache）的命中又不利于线程束（warp）的合并访问（coalesced access）。因此我们试图对需要合并的数据段进行进一步的分块，然后分别将两个数据段中对应的块进行合并，以进一步提升该排序算法的性能。

1. 分块合并

首先我们需要考虑如何对数据段进行合理的分块，因为必须要保证待合并的两个数据段第一块中的所有元素均不大于第二块中的所有元素，这样才能在第一对块合并前确认第二对块合并后的位置以保证并行性。最直观的方法为将第一个数据段直接分块，然后在第二个数据段中找到对应的块，可通过在第二个数据段中查找第一个数据段各块首元素的位置，按这些位置便可将第二个数据段分块。

第一个数据段等分为长为4的块，再根据每块首元素将第二个数据段分块，然后将对应块合并

这样的分块策略对于随机数据来说能取得不错的效果，但如果两个数据段中的元素分布不均匀，甚至第一个数据段中的所有元素都小于（或大于）第二个数据段，该分块策略就不能进行有效的切分，也不能达到良好的性能。因此除了第一个数据段各块首元素外，还用第二个数据段等分后每块首元素来进行切分，且同时对两个数据段进行切分以得到对应的块。

两个数据段先都等分为长为4的块，再根据每块首元素将对两个数据段重新分块。由于每个数据段都会被自己块的首元素所切分，因此该分块策略下每块最大长度为4，保证了分块的有效性。

接下来仍然是最令人头疼（甚至是头秃）的时刻，就是把上述流程转换为核函数的代码。最后一步必然是将对应块进行合并，而对于切分操作有两种策略：一是先取出所有块首元素，排序后再用来切分两个数据段；二是先用所有块首元素来切分两个数据段，再将切分位点进行排序。显然前者需要3步而后者只需2步，鉴于核函数启动的耗时我们选择后者。

第一步是用所有块首元素来切分两个数据段，为此需要有一个rank数组来保存每个首元素的切分位点。我们让一个线程负责计算一个块首元素在两个数据段中的切分点，在当前数据段中的切分位点可直接由块索引乘以块大小得到，在另一数据段中的切分位点由二分查找得到。rank数组中依次保存第一、二个数据段的块首元素在第一、二个数据段中的切分位点，

一半的线程负责计算所有第一个数据段中的块首元素分别在对应的两个数据段中的切分位点

另一半的线程负责计算所有第二个数据段中的块首元素分别在对应的两个数据段中的切分位点

__global__ void rank_gpu(const double* data, unsigned int n, unsigned int s, unsigned int* rank)
{unsigned int pl = s >> log_tile_size, t = blockIdx.x * blockDim.x + threadIdx.x, k = t / (pl << 1), i = t & (pl << 1) - 1;data += k * s << 1; n -= k * s << 1; rank += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (n <= s || i >= p)return;rank[i] = i & pl ? lower_bound(data, s, data[i << log_tile_size]) : i << log_tile_size;rank[p + i] = i & pl ? (i ^ pl) << log_tile_size : lower_bound(data + s, min(n - s, s), data[i << log_tile_size]);
}

第二步将所有数据段的两半切分位点进行排序，实际也是两半有序数组的合并，故其原理与直接合并操作相同，由于每个数据段可能会有超过 tile_size 个块，所以需要一个线程块组来负责一次合并；这儿我们选择让一个线程块组来负责一对数据段中切分位点的合并（两次合并），因为切分位点的总个数与这对数据段的长度相关（最后一对数据段中第二个数据段可能是不满的）。所有线程块组仍然排列为二维，因为网格大小第二、三维上限是-1。

每一组线程块负责一对数据段中切分位点的合并，合并操作仍然由二分查找实现

首先仍然是数据段包含的块数量、、p及 data、rank、split数组偏移量的计算，除数据段组别k和当前线程负责的数据索引i分别由线程块索引的第二、三维和线程块索引的第一维、线程索引得到外，其余计算与第一步完全一致，若偏移量大于总数据规模则直接返回。最后在判断数据索引是否有效（i< 块数量）后，利用二分查找将切分位点进行排序，结果写入split数组中。

__global__ void ranksort_gpu(const unsigned int* rank, unsigned int n, unsigned int s, unsigned int* split)
{unsigned int pl = s >> log_tile_size, k = blockIdx.z * gridDim.y + blockIdx.y, i = blockIdx.x * blockDim.x + threadIdx.x;if (k * s << 1 >= n)return;n -= k * s << 1; rank += k * pl << 2; split += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (i < pl){split[i + lower_bound(rank + pl, pr, rank[i])] = rank[i];split[p + i + lower_bound(rank + p + pl, pr, rank[p + i])] = rank[p + i];}if (i < pr){split[i + upper_bound(rank, pl, rank[pl + i])] = rank[pl + i];split[p + i + upper_bound(rank + p, pl, rank[p + pl + i])] = rank[p + pl + i];}
}

第三步则是根据每组数据段的切分位点进行分块合并，每个线程块负责一对对应块的合并，因此每对数据段的合并需要一组线程块来完成。

一组线程块负责一对数据段的合并，每个线程块根据切分位点获取待合并的数据，合并后置于结果数组中

__global__ void mergesegment_gpu(const double* data, const unsigned int* split, unsigned int n, unsigned int s, double* result)
{__shared__  double databuffer[tile_size << 1], resultbuffer[tile_size << 1];unsigned int pl = s >> log_tile_size, k = blockIdx.z * gridDim.y + blockIdx.y;if (k * s << 1 >= n)return;data += k * s << 1; result += k * s << 1; n -= k * s << 1; split += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (n <= s){unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;if (i < n)result[i] = data[i];}else if (blockIdx.x < p){unsigned int fstart = split[blockIdx.x], lstart = split[p + blockIdx.x];unsigned int fend = blockIdx.x + 1 < p ? split[blockIdx.x + 1] : s, lend = blockIdx.x + 1 < p ? split[p + blockIdx.x + 1] : min(n - s, s);unsigned int flen = fend - fstart, llen = lend - lstart;if (threadIdx.x < flen)databuffer[threadIdx.x] = data[fstart + threadIdx.x];if (threadIdx.x < llen)databuffer[flen + threadIdx.x] = data[s + lstart + threadIdx.x];__syncthreads();if (threadIdx.x < flen)resultbuffer[threadIdx.x + lower_bound(databuffer + flen, llen, databuffer[threadIdx.x])] = databuffer[threadIdx.x];if (threadIdx.x < llen)resultbuffer[threadIdx.x + upper_bound(databuffer, flen, databuffer[flen + threadIdx.x])] = databuffer[flen + threadIdx.x];__syncthreads();if (threadIdx.x < flen)result[fstart + lstart + threadIdx.x] = resultbuffer[threadIdx.x];if (threadIdx.x < llen)result[fstart + lstart + flen + threadIdx.x] = resultbuffer[flen + threadIdx.x];}
}

最后将上面三个函数及之前的分块双调排序函数组合成最终的排序函数，这一改进版的排序函数中除了用于存放临时合并结果的数组外，还额外需要两个索引数组用于存放切分位点和排序后的切分位点，每个索引数组的长度应是所有数据的总分块数 x 2，因为每个数据段还要被同组另一数据段中的块首元素切分。仍先进行一次双调排序进行分块排序，然后逐轮分块合并，最后释放临时数组空间并返回结果。注意分块合并中第一步的总线程数是固定的，即为总分块数，因为一个线程负责索引数组中两个索引的计算；第二步中一组线程块负责一对数据段中切分位点的排序，因此线程块组的总数即为长为2s的数据段对数（即为网格的第二、三维大小），而每个线程负责四个索引的排序，所以一对数据段中切分位点的排序总共需要的线程数恰等于每个数据段的初始分块数（即为网格的第一维大小和线程块大小）；第三步中一组线程块负责一对数据段的分块合并，所以其线程块组的总数与第二步相同，但若存在额外的不足一个数据段的剩余数据（所有数据的总分块数是偶数时）则需要一个额外的线程块组来进行额外数据的拷贝，而每组线程块需要的线程块数为每个数据段的分块数，恰等于其初始分块数 x 2（即为网格的第一维大小）。

double* sortsegment_gpu(double* data, unsigned int n)
{double* tmp;cudaMalloc(&tmp, sizeof(double) * n);unsigned int* rank, index_len = (n + tile_size - 1 >> log_tile_size) << 1;cudaMalloc(&rank, sizeof(unsigned int) * index_len << 1);unsigned int* split = rank + index_len;unsigned int b = n + (tile_size << 1) - 1 >> log_tile_size + 1, s = tile_size << 1, m = n + tile_size - 1 >> log_tile_size;binoticsort_gpu<<<b, s >> 1>>>(data, n);for (b = b + 1 >> 1; s < n; s <<= 1, b = b + 1 >> 1){unsigned int p = s + tile_size - 1 >> log_tile_size, k = (n - 1) / s + 1 >> 1;rank_gpu<<<m + tile_size - 1 >> log_tile_size, tile_size>>>(data, n, s, rank);ranksort_gpu<<<dim3(p + tile_size - 1 >> log_tile_size, std::min(k, 32768u), k + 32767 >> 15), tile_size>>>(rank, n, s, split);if (!((n - 1) / s & 1))++k;mergesegment_gpu<<<dim3(p << 1, std::min(k, 32768u), k + 32767 >> 15), tile_size>>>(data, split, n, s, tmp);std::swap(data, tmp);}cudaDeviceSynchronize();cudaFree(tmp);cudaFree(rank);return data;
}

至此我们已经实现了基于分块合并的改进版 CUDAsort，下面我们对其性能进行评估。

2. 性能测试

我们用相同的数据规模对该改进版进行性能测试，测试结果如下：

我们发现数据规模在百万及以下时该改进版的性能反而有所下降，这是因为在数据规模较小时GPU的缓存能较好地缓解对全局内存频繁访问的时延，而此时分块合并法由于有更多的核函数启动会消耗更多的时间，故其性能反而不如直接合并法；但当数据规模在千万及以上时该改进版展现了更优异的性能，10亿双精度浮点型数据的排序仅需约1.1S，相较于初始版本有40%的性能提升。

五. 总结

我们基于双调排序和归并排序实现了GPU上的排序算法，若忽略数据传输的耗时，改进版 CUDAsort在数据规模较大时能达到百倍左右的加速比，能在1s左右完成10亿数据的排序。但是作者发现在tile_size设为512或1024时，该排序算法似乎会出错，思索半天未果，在此欢迎各路大神发表看法，包括但不限于出错的原因、更高性能的排序算法等，本人由于并不是计算机专业的，水平有限，因此只能想到这一种GPU上的高效排序算法了。

最后贴上完整的代码：

#include <cuda.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <float.h>#define tile_size 256u
#define log_tile_size 8u// auxiliary functions
template<typename type>__device__ void swap(type& v1, type& v2)
{type t = v1;v1 = v2;v2 = t;
}
template<typename type>__device__ unsigned int lower_bound(const type* data, unsigned int n, const type& value)
{unsigned int len = n, half;const type* start = data, * mid;while (len){half = len >> 1;mid = data + half;if (*mid < value){data = mid + 1;len = len - half - 1;}elselen = half;}return data - start;
}
template<typename type>__device__ unsigned int upper_bound(const type* data, unsigned int n, const type& value)
{unsigned int len = n, half;const type* start = data, *mid;while (len){half = len >> 1;mid = data + half;if (value < *mid)len = half;else{data = mid + 1;len = len - half - 1;}}return data - start;
}
// binotic sort
__global__ void binoticsort_gpu(double* data, unsigned int n)
{__shared__  double buffer[tile_size << 1];unsigned int d = blockIdx.x << log_tile_size + 1, i = threadIdx.x;data += d; n -= d;buffer[i] = i < n ? data[i] : DBL_MAX;buffer[i + blockDim.x] = i + blockDim.x < n ? data[i + blockDim.x] : DBL_MAX;__syncthreads();unsigned int s, t;for (unsigned int k = 2; k <= blockDim.x; k <<= 1)for (unsigned int p = k >> 1; p; p >>= 1){s = (i & -p) << 1 | i & p - 1;t = s | p;if (s & k ? buffer[s] < buffer[t] : buffer[s] > buffer[t])swap(buffer[s], buffer[t]);__syncthreads();}for (unsigned int p = blockDim.x; p; p >>= 1){s = (i & -p) << 1 | i & p - 1;t = s | p;if (buffer[s] > buffer[t])swap(buffer[s], buffer[t]);__syncthreads();}if (i < n)data[i] = buffer[i];if (i + blockDim.x < n)data[i + blockDim.x] = buffer[i + blockDim.x];
}
// merge directly
__global__ void mergedirect_gpu(const double* data, unsigned int n, unsigned int s, double* result)
{unsigned int d1 = (blockIdx.z * gridDim.y + blockIdx.y) * s << 1, d2 = d1 + s, i = d1 + blockIdx.x * blockDim.x + threadIdx.x;if (i < n)result[d2 < n ? i + lower_bound(data + d2, min(s, n - d2), data[i]) : i] = data[i];if (s + i < n)result[i + upper_bound(data + d1, s, data[s + i])] = data[s + i];
}
double* sortdirect_gpu(double* data, unsigned int n)
{double* tmp;cudaMalloc(&tmp, sizeof(double) * n);unsigned int b = n + (tile_size << 1) - 1 >> log_tile_size + 1, s = tile_size << 1;binoticsort_gpu<<<b, s >> 1>>>(data, n);for (b = b + 1 >> 1; s < n; s <<= 1, b = b + 1 >> 1){mergedirect_gpu<<<dim3(s >> log_tile_size, std::min(b, 32768u), b + 32767 >> 15), tile_size>>>(data, n, s, tmp);std::swap(data, tmp);}cudaDeviceSynchronize();cudaFree(tmp);return data;
}
// merge in segments
__global__ void rank_gpu(const double* data, unsigned int n, unsigned int s, unsigned int* rank)
{unsigned int pl = s >> log_tile_size, t = blockIdx.x * blockDim.x + threadIdx.x, k = t / (pl << 1), i = t & (pl << 1) - 1;data += k * s << 1; n -= k * s << 1; rank += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (n <= s || i >= p)return;rank[i] = i & pl ? lower_bound(data, s, data[i << log_tile_size]) : i << log_tile_size;rank[p + i] = i & pl ? (i ^ pl) << log_tile_size : lower_bound(data + s, min(n - s, s), data[i << log_tile_size]);
}
__global__ void ranksort_gpu(const unsigned int* rank, unsigned int n, unsigned int s, unsigned int* split)
{unsigned int pl = s >> log_tile_size, k = blockIdx.z * gridDim.y + blockIdx.y, i = blockIdx.x * blockDim.x + threadIdx.x;if (k * s << 1 >= n)return;n -= k * s << 1; rank += k * pl << 2; split += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (i < pl){split[i + lower_bound(rank + pl, pr, rank[i])] = rank[i];split[p + i + lower_bound(rank + p + pl, pr, rank[p + i])] = rank[p + i];}if (i < pr){split[i + upper_bound(rank, pl, rank[pl + i])] = rank[pl + i];split[p + i + upper_bound(rank + p, pl, rank[p + pl + i])] = rank[p + pl + i];}
}
__global__ void mergesegment_gpu(const double* data, const unsigned int* split, unsigned int n, unsigned int s, double* result)
{__shared__  double databuffer[tile_size << 1], resultbuffer[tile_size << 1];unsigned int pl = s >> log_tile_size, k = blockIdx.z * gridDim.y + blockIdx.y;if (k * s << 1 >= n)return;data += k * s << 1; result += k * s << 1; n -= k * s << 1; split += k * pl << 2;unsigned int pr = min(n - s, s) + tile_size - 1 >> log_tile_size, p = pl + pr;if (n <= s){unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;if (i < n)result[i] = data[i];}else if (blockIdx.x < p){unsigned int fstart = split[blockIdx.x], lstart = split[p + blockIdx.x];unsigned int fend = blockIdx.x + 1 < p ? split[blockIdx.x + 1] : s, lend = blockIdx.x + 1 < p ? split[p + blockIdx.x + 1] : min(n - s, s);unsigned int flen = fend - fstart, llen = lend - lstart;if (threadIdx.x < flen)databuffer[threadIdx.x] = data[fstart + threadIdx.x];if (threadIdx.x < llen)databuffer[flen + threadIdx.x] = data[s + lstart + threadIdx.x];__syncthreads();if (threadIdx.x < flen)resultbuffer[threadIdx.x + lower_bound(databuffer + flen, llen, databuffer[threadIdx.x])] = databuffer[threadIdx.x];if (threadIdx.x < llen)resultbuffer[threadIdx.x + upper_bound(databuffer, flen, databuffer[flen + threadIdx.x])] = databuffer[flen + threadIdx.x];__syncthreads();if (threadIdx.x < flen)result[fstart + lstart + threadIdx.x] = resultbuffer[threadIdx.x];if (threadIdx.x < llen)result[fstart + lstart + flen + threadIdx.x] = resultbuffer[flen + threadIdx.x];}
}
double* sortsegment_gpu(double* data, unsigned int n)
{double* tmp;cudaMalloc(&tmp, sizeof(double) * n);unsigned int* rank, index_len = (n + tile_size - 1 >> log_tile_size) << 1;cudaMalloc(&rank, sizeof(unsigned int) * index_len << 1);unsigned int* split = rank + index_len;unsigned int b = n + (tile_size << 1) - 1 >> log_tile_size + 1, s = tile_size << 1, m = n + tile_size - 1 >> log_tile_size;binoticsort_gpu<<<b, s >> 1>>>(data, n);for (b = b + 1 >> 1; s < n; s <<= 1, b = b + 1 >> 1){unsigned int p = s + tile_size - 1 >> log_tile_size, k = (n - 1) / s + 1 >> 1;rank_gpu<<<m + tile_size - 1 >> log_tile_size, tile_size>>>(data, n, s, rank);ranksort_gpu<<<dim3(p + tile_size - 1 >> log_tile_size, std::min(k, 32768u), k + 32767 >> 15), tile_size>>>(rank, n, s, split);if (!((n - 1) / s & 1))++k;mergesegment_gpu<<<dim3(p << 1, std::min(k, 32768u), k + 32767 >> 15), tile_size>>>(data, split, n, s, tmp);std::swap(data, tmp);}cudaDeviceSynchronize();cudaFree(tmp);cudaFree(rank);return data;
}

六、gpu和CUDA编程的经典入门书籍

如果让我现在推荐一个项目，我会建议去看Karpathy的llm.c，写得巨好：https://github.com/karpathy/llm.c

CUDA入门

WHY CUDA？

GPU 底层的硬件设计导致它特别擅长于执行逻辑简单但是并行程度极高的计算任务，而 CUDA 的主要目的是“向显卡提交一个又一个任务，每一个任务都形如“给定一个函数，与调用它的参数，请在显卡上运行这个函数”。

关于CPU和GPU的工作原理，Nvidia给出了一个比较生动的例子：

NVIDIA现场形象展示CPU和GPU工作原理上的区别：https://www.bilibili.com/video/BV1ry4y1y7KZ/?

How to Start CUDA Learning？

知乎找到的一个官方学习路线，同样也可以参考：https://zhuanlan.zhihu.com/p/273607744

首先，至少要有一块GPU，以及配置相应的CUDA环境（安装CUDA-Toolkit等）

Configure Software Environment

教程：https://zhuanlan.zhihu.com/p/432092988
推荐OS：任意Linux系统（推荐Ubuntu）
推荐IDE：VSCode（Remote SSH）/Clion（Jetbrain配置CUDA环境便捷）
推荐CUDA版本：11.x（12.x有点问题）
推荐GPU：实验室的V100/Titan/P100

配好后，验证是否已经安装好，只需建立hello_world.cu执行下面的例程：

// hello_world.cu
// CUDA: Hello,World
#include <stdio.h>__global__ void hello_world(void)
{printf("GPU: Hello world!\\n");
}int main(int argc, char** argv)
{printf("CPU: Hello world!\\n");hello_world<<<1, 10>>>();cudaDeviceReset(); // 如果没有这一行，无法从GPU输出 "hello world"return 0;
}
nvcc hello_world.cu -o hello_world
./hello_world
//Output:
// CPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!
// GPU: Hello world!

How to Understand CUDA?

如何理解包括Block、Thread、wrap等概念和操作

推荐的相关资料：

谭升的博客：https://face2ai.com/program-blog/#GPU编程（CUDA）
我见过的最好的中文入门CUDA教程
相关代码：https://github.com/Tony-Tan/CUDA_Freshman
CUDA - From Correctness to Performance：https://wiki.lcpu.dev/zh/hpc/from-scratch/cuda
尝试在纠错中理解CUDA相关概念和注意事项
相关代码：https://github.com/interestingLSY/CUDA-From-Correctness-To-Performance-Code
Xiao's CUDA Optimization Guide：https://github.com/XiaoSong9905/CUDA-Optimization-Guide
同样是一份不错的中文入门CUDA教程，更偏向优化方面
高性能并行编程与优化：https://github.com/parallel101/course
尝试理解并行计算的本质，和写CUDA本身无关但能够明白并行化思维
相关视频：https://space.bilibili.com/263032155
Efficient Distributed GPU Programming for Exascale, an SC/ISC Tutorial
尝试理解如何利用MPI在单机多卡/多机多卡实现GPU Programming，世界最顶级超算会议ISC23/SC23教程
相关代码：https://github.com/FZJ-JSC/tutorial-multi-gpu/tree/main

推荐的书籍：

Programming Massively Parallel Processors（好像链接失效了，如果有需要可以问我要）：https://github.com/h3ct0rjs/HighPerformanceComputing/blob/master/BookRef/Programming Massively Parallel Processors.pdf

How to Understand CUDA in a Real Project

在项目中尝试理解CUDA的用处，目前来说，它最热门的用处是大模型的训练、推理以及CG方向的渲染，当然它也可以迁移到图计算、图卷积神经网络（GCN）方面

LightSeq：A High Performance Library for Sequence Processing and Generation
字节跳动开源的生成模型推理加速引擎，BERT、GPT、VAE等等全都支持，速度很快
代码地址：https://github.com/bytedance/lightseq
FasterTransformer：Transformer related optimization, including BERT, GPT
英伟达开源的Transformer推理加速引擎
代码地址：https://github.com/NVIDIA/FasterTransformer
Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity
当前推理的SOTA，达摩院9月的最新研究
代码地址：https://github.com/AlibabaResearch/flash-llm
HierarchicalKV：A part of NVIDIA Merlin and provides hierarchical key-value storage to meet RecSys requirements
将键值特征嵌入存储在GPU的高带宽内存和主机内存中，还可以用作通用键值存储
代码地址：https://github.com/NVIDIA-Merlin/HierarchicalKV
GPTQ inference TVM kernel
TVM推理，能完整理解一个小模型的训练、推理，亮点是用了TVM（深度学习模型编译框架）编译，这是一个未来的好方向
代码地址：https://github.com/LeiWang1999/AutoGPTQ.tvm
3D Gaussian Splatting for Real-Time Radiance Field Rendering（CG方向）
代码地址：https://github.com/graphdeco-inria/gaussian-splatting
Instant Neural Graphics Primitives with a Multiresolution Hash Encoding（CG方向）
代码地址：https://github.com/NVlabs/instant-ngp

课程推荐

课程推荐：CS344: Intro to Parallel Programming：https://developer.nvidia.com/udacity-cs344-intro-parallel-programming

七、CUDA~编程接口与硬件实现

整理CUDA手册中的CUDA编程接口和GPU硬件实现部分。C++ 实现

CUDA的编程接口

CUDA C++为熟悉C++编程语言的用户提供了一个简单的路径，使他们能够轻松地编写供设备执行的程序。

它由C++语言的最小扩展集和一个运行时库组成。

核心语言扩展已在编程模型中介绍。它们允许程序员将内核定义为C++函数，并使用一些新的语法每次调用函数时指定grid and block的维度。所有扩展的完整描述可以在C++语言扩展中[1]找到。包含这些扩展的任何源文件都必须按照NVCC编译中的说明使用nvcc进行编译[2]。

运行时在CUDA运行时[3]中被介绍。它提供在主机上执行的C和C++函数，用于分配和释放设备内存，传输主机内存和设备内存之间的数据，管理具有多个设备的系统等。运行时的完整描述可以在CUDA参考手册中找到。

运行时是基于一个更低级别的C API构建的，即CUDA驱动API，应用程序也可以访问。驱动API通过公开更低级别的概念（如CUDA上下文 - 设备的主机进程的类比，以及CUDA模块 - 设备的动态加载库的类比）提供了额外的控制级别。大多数应用程序不使用驱动API，因为它们不需要这个额外的控制级别，而当使用运行时时，上下文和模块管理是隐式的，从而产生更简洁的代码。由于运行时与驱动API是互操作的，大多数需要一些驱动API功能的应用程序可以默认使用运行时API，并仅在需要时使用驱动API。驱动程序API在Driver API[4]中介绍。

使用NVCC进行编译

内核可以使用CUDA指令集架构PTX编写，该架构在PTX参考手册中有描述。然而，通常更有效的方法是使用高级编程语言如C++。在两种情况下，内核必须由nvcc编译成二进制代码才能在设备上执行。

nvcc是一个编译器驱动程序，简化了编译C++或PTX代码的过程：它提供简单和熟悉的命令行选项，并通过调用实现不同编译阶段的工具集合来执行它们。本节概述了nvcc工作流程和命令选项。完整说明可在nvcc用户手册中找到。

Compilation Workflow

Offline Compilation 离线编译

使用nvcc编译的源文件可以包括host代码（即在主机上执行的代码）和设备代码（即在设备上执行的代码）的混合。nvcc的基本工作流程包括将设备代码与主机代码分开，然后:

将设备代码编译成汇编形式（PTX代码）和/或二进制形式（cubin对象）,
并通过替换在内核中引入的<<<...>>>语法（并在执行配置中更详细地描述[1]）来修改主机代码，以必要的CUDA运行时函数调用从PTX代码和/或cubin对象加载和启动每个编译的内核

修改后的主机代码输出为C++代码，可以使用另一个工具编译，或者直接通过让nvcc在最后的编译阶段调用主机编译器输出为对象代码。

Applications can then:

要么链接到编译后的主机代码（这是最常见的情况），
要么忽略修改后的主机代码（如果有的话），并使用CUDA驱动API（参见驱动API）加载和执行PTX代码或cubin对象。

Just-in-Time Compilation 即时编译

任何在运行时由应用程序加载的PTX代码都会被设备驱动程序进一步编译成二进制代码。这被称为即时编译。即时编译增加了应用程序的加载时间，但允许应用程序受益于每个新设备驱动程序带来的任何新的编译器改进。这也是应用程序在被编译时尚不存在的设备上运行的唯一方式，如在应用程序兼容性[6]中详细描述。

当设备驱动程序为某个应用程序即时编译一些PTX代码时，它会自动缓存生成的二进制代码的副本，以避免在后续调用应用程序时重复编译。该缓存 - 称为计算缓存 - 在升级设备驱动程序时会自动失效，这样应用程序就可以从新的即时编译器中受益，该编译器内置于设备驱动程序中。

环境变量可用于控制即时编译，如在CUDA环境变量中所述[7]。

作为使用nvcc编译CUDA C++设备代码的替代方法，可以在运行时使用NVRTC编译CUDA C++设备代码到PTX。NVRTC是一个用于CUDA C++的运行时编译库；更多信息可以在NVRTC用户指南中找到。

Binary Compatibility

二进制代码是特定于架构的。使用编译器选项-code 生成_cubin_ 对象，该选项指定目标架构：例如，使用-code=sm_80 编译会为计算能力[8] 8.0的设备生成二进制代码。从一个次要版本到下一个次要版本保证二进制兼容性，但从一个次要版本到前一个次要版本或跨主要版本则不保证。换句话说，为计算能力_X.y_ 生成的_cubin_ 对象只会在计算能力为_X.z_ 的设备上执行，其中_z≥y_ 。

Binary compatibility is supported only for the desktop. It is not supported for Tegra. Also, the binary compatibility between desktop and Tegra is not supported.

PTX Compatibility

某些PTX指令仅在具有更高计算能力的设备上受支持。例如，Warp Shuffle函数仅在计算能力为5.0及以上的设备上受支持。-arch编译器选项指定编译C++到PTX代码时假定的计算能力。因此，包含warp shuffle的代码，例如，必须使用-arch=compute_50（或更高）进行编译。

为某个特定计算能力生成的PTX代码始终可以编译为更大或相等的计算能力的二进制代码。请注意，从早期PTX版本编译的二进制文件可能不会使用某些硬件功能。例如，针对计算能力7.0（Volta）的设备编译的二进制文件，从为计算能力6.0（Pascal）生成的PTX编译，将不会使用Tensor Core指令，因为Pascal上没有这些指令。因此，最终的二进制文件的性能可能比使用PTX的最新版本生成的二进制文件差。

Application Compatibility

为了在具有特定计算能力的设备上执行代码，应用程序必须加载与此计算能力兼容的二进制或PTX代码，如在二进制兼容性[9]和PTX兼容性[10]中所述。特别地，为了能够在具有更高计算能力的未来架构上执行代码（为其尚未生成二进制代码），应用程序必须加载将为这些设备即时编译的PTX代码（参见即时编译[11]）。

CUDA C++应用程序中嵌入的PTX和二进制代码是由-arch和-code编译器选项或-gencode编译器选项控制的，如nvcc用户手册中详细描述的那样。例如，

nvcc x.cu -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code="compute_70,sm_70"

嵌入了与计算能力5.0和6.0兼容的二进制代码（第一和第二个-gencode选项）以及与计算能力7.0兼容的PTX和二进制代码（第三个-gencode选项）。

主机代码在运行时自动生成，以自动选择最适合的代码进行加载和执行，上述示例中将是：

5.0的二进制代码用于计算能力为5.0和5.2的设备，
6.0的二进制代码用于计算能力为6.0和6.1的设备，
7.0的二进制代码用于计算能力为7.0和7.5的设备，
PTX代码在运行时编译为计算能力为8.0和8.6的设备的二进制代码。

x.cu可以有一个优化的代码路径，例如使用仅在计算能力为8.0及更高版本的设备上受支持的warp reduction操作。__CUDA_ARCH__宏可用于基于计算能力区分各种代码路径。它仅为设备代码定义。例如，当使用-arch=compute_80编译时，__CUDA_ARCH__等于800。

使用驱动API的应用程序必须将代码编译为单独的文件，并在运行时明确加载和执行最合适的文件。

Volta架构引入了独立线程调度，这改变了GPU上线程的调度方式。对于依赖于先前架构中SIMT调度的特定行为的代码，独立线程调度可能会更改参与线程的集合，导致结果不正确。为了在实施独立线程调度中详细描述的纠正措施时帮助迁移，Volta开发人员可以选择使用编译器选项组合-arch=compute_60 -code=sm_70来选择Pascal的线程调度。

nvcc用户手册列出了-arch、-code和-gencode编译器选项的各种简写。例如，-arch=sm_70是-arch=compute_70 -code=compute_70,sm_70的简写（与-gencode arch=compute_70,code="compute_70,sm_70"相同）。

C++ Compatibility

编译器的前端根据C++语法规则处理CUDA源文件。主机代码支持完整的C++。但是，只有C++语言支持中描述的部分C++代码在device代码中得到完全支持。

64-位兼容

nvcc的64位版本以64位模式编译设备代码（即指针为64位）。以64位模式编译的设备代码只能与以64位模式编译的主机代码一起使用。

CUDA Runtime

运行时在cudart库中实现，该库链接到应用程序，可以通过cudart.lib或libcudart.a静态链接，或通过cudart.dll或libcudart.so动态链接。需要cudart.dll和/或cudart.so进行动态链接的应用程序通常将它们作为应用程序安装包的一部分包含在内。只有在链接到CUDA运行时的同一实例的组件之间传递CUDA运行时符号的地址才是安全的。

它的所有入口点都带有cuda前缀。

如异构编程[12]中所述，CUDA编程模型假定一个由主机和设备组成的系统，每个系统都有自己的独立内存。设备内存[13]概述了用于管理设备内存的运行时函数。
共享内存[14]说明了如何使用在线程层次结构[15]中引入的共享内存来最大化性能。
页锁定的主机内存[16]介绍了与设备内存之间的数据传输重叠的内核执行所需的页锁定的主机内存。
异步并发执行[17]描述了用于在系统的各个级别启用异步并发执行的概念和API。
多设备系统[18]展示了编程模型如何扩展到与同一主机连接的多个设备的系统。
错误检查[19]描述了如何正确检查运行时生成的错误。
调用堆栈[20]提到了用于管理CUDA C++调用堆栈的运行时函数。
纹理和表面内存[21]介绍了纹理和表面内存空间，这些空间提供了另一种访问设备内存的方法；它们还公开了GPU纹理硬件的一个子集。
图形互操作性[22]介绍了运行时提供的与两个主要图形API（OpenGL和Direct3D）互操作的各种函数。

Initialization

从CUDA 12.0开始，cudaInitDevice()和cudaSetDevice()调用会初始化与指定设备关联的运行时和主上下文。如果没有这些调用，运行时将隐式使用设备0，并根据需要自我初始化以处理其他运行时API请求。在计时运行时函数调用和解释第一次调用运行时的错误代码时，需要记住这一点。在12.0之前，cudaSetDevice()不会初始化运行时，应用程序通常会使用无操作运行时调用cudaFree(0)来隔离运行时初始化和其他api活动（为了计时和错误处理）。

运行时为系统中的每个设备创建一个CUDA上下文（有关CUDA上下文的更多详细信息，请参见上下文[23]）。这个上下文是该设备的主上下文，并在第一个需要在此设备上有活动上下文的运行时函数上初始化。它在应用程序的所有主机线程之间共享。作为此上下文创建的一部分，如果需要，设备代码会进行即时编译（参见即时编译[24]）并加载到设备内存中。这一切都是透明的。如果需要，例如，为了驱动API互操作性，可以从驱动API访问设备的主上下文，如运行时和驱动API之间的互操作性[25]中所述。

当主机线程调用cudaDeviceReset()时，这将销毁主机线程当前操作的设备的主上下文（即，在设备选择[26]中定义的当前设备）。任何将此设备作为当前设备的主机线程进行的下一个运行时函数调用都将为此设备创建一个新的主上下文。

CUDA接口使用在主机程序初始化期间初始化并在主机程序终止期间销毁的全局状态。CUDA运行时和驱动程序无法检测此状态是否无效，因此在程序初始化或终止期间（在main之后）使用这些接口（隐式或显式）将导致未定义的行为。

从CUDA 12.0开始，cudaSetDevice()现在将在为主机线程更改当前设备后明确初始化运行时。CUDA的先前版本在cudaSetDevice()之后延迟了新设备上的运行时初始化，直到进行了第一个运行时调用。这一变化意味着现在检查cudaSetDevice()的返回值以查找初始化错误非常重要。

参考手册的错误处理和版本管理部分的运行时函数不会初始化运行时。

Device Memory

如在异构编程[27]中所提及的，CUDA编程模型假设一个由主机和设备组成的系统，每个都有自己的独立内存。内核在设备内存中运行，因此运行时提供了分配、释放和复制设备内存的函数，以及在主机内存和设备内存之间传输数据。

设备内存可以分配为线性内存或CUDA数组。

CUDA数组是为纹理提取优化的不透明内存布局。它们在 Texture and Surface Memory[28]中有描述。

线性内存在单一统一的地址空间中分配，这意味着单独分配的实体可以通过指针相互引用，例如在二叉树或链表中。地址空间的大小取决于主机系统（CPU）和所使用的GPU的计算能力：

在计算能力为5.3（Maxwell）及更早版本的设备上，CUDA驱动程序会创建一个未提交的40位虚拟地址保留，以确保内存分配（指针）位于支持的范围内。这个保留会显示为保留的虚拟内存，但直到程序实际分配内存之前，不会占用任何物理内存。

线性内存通常使用cudaMalloc()进行分配，并使用cudaFree()进行释放，主机内存和设备内存之间的数据传输通常使用cudaMemcpy()完成。在内核的向量加法代码示例中，需要将向量从主机内存复制到设备内存：

// Device code  
__global__ void VecAdd(float* A, float* B, float* C, int N)  
{  int i = blockDim.x * blockIdx.x + threadIdx.x;  if (i < N)  C[i] = A[i] + B[i];  
}  // Host code  
int main()  
{  int N = ...;  size_t size = N * sizeof(float);  // Allocate input vectors h_A and h_B in host memory  float* h_A = (float*)malloc(size);  float* h_B = (float*)malloc(size);  float* h_C = (float*)malloc(size);  // Initialize input vectors  ...  // Allocate vectors in device memory  float* d_A;  cudaMalloc(&d_A, size);  float* d_B;  cudaMalloc(&d_B, size);  float* d_C;  cudaMalloc(&d_C, size);  // Copy vectors from host memory to device memory  cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);  cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);  // Invoke kernel  int threadsPerBlock = 256;  int blocksPerGrid =  (N + threadsPerBlock - 1) / threadsPerBlock;  VecAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);  // Copy result from device memory to host memory  // h_C contains the result in host memory  cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);  // Free device memory  cudaFree(d_A);  cudaFree(d_B);  cudaFree(d_C);  // Free host memory  ...  
}

GPU的硬件实现

NVIDIA GPU架构是围绕可扩展的多线程流式多处理器（SMs）阵列构建的。当主机CPU上的CUDA程序调用一个内核网格时，网格的块会被枚举并分配给具有可用执行能力的多处理器。一个线程块的线程在一个多处理器上并发执行，一个多处理器上可以并发执行多个线程块。当线程块终止时，新的块在空闲的多处理器上启动。

一个多处理器被设计为并发执行数百个线程。为了管理如此大量的线程，它采用了一种称为SIMT（单指令，多线程）的独特架构，该架构在SIMT架构[29]中有描述。指令是流水线式的，利用单个线程内的指令级并行性，以及通过同时的硬件多线程实现广泛的线程级并行性，如硬件多线程[30]中详细描述的那样。与CPU核心不同，它们是按顺序发出的，没有分支预测或推测执行。

SIMT架构[31]和硬件多线程[32]描述了流处理器的架构特性，这些特性对所有设备都是通用的。计算能力5.x[33]、计算能力6.x[34]和计算能力7.x[35]分别为计算能力5.x、6.x和7.x的设备提供了具体细节。

NVIDIA GPU架构使用小端表示法。

SIMT Architecture

多处理器创建、管理、调度并执行称为_warps_的32个并行线程组。组成warp的单个线程在同一程序地址上同时开始，但它们有自己的指令地址计数器和寄存器状态，因此可以自由地分支和独立执行。_warp_一词来源于编织（weaving），编织（weaving）是一种古老的技术，其中多个线程交织在一起创建一个布料。在这里，这个词被用作一个比喻，表示多个线程并行工作，就像在编织中的线程一样。_半warp_是warp的第一半或第二半。_四分之一warp_是warp的第一、第二、第三或第四部分（ A half-warp is either the first or second half of a warp. A quarter-warp is either the first, second, third, or fourth quarter of a warp）。

详细解释下上述这段内容：

在CUDA架构中，多处理器是执行线程的硬件单元。这里提到的“warp”是一个包含32个线程的组，这些线程在多处理器上并行执行。
尽管所有线程都从同一个程序地址开始，但每个线程都有自己的指令计数器和寄存器状态。这意味着，尽管它们开始时执行相同的指令，但随后可以根据各自的数据和条件独立地分支和执行。
一个warp包含32个线程，所以一个“半warp”就是其中的16个线程，可以是前16个或后16个线程。一个warp被分为四个部分，每部分包含8个线程。所以，“四分之一warp”可以是warp的前8个、第9到16个、第17到24个或最后8个线程。

当一个多处理器被赋予一个或多个要执行的线程块时，它将它们划分为warps，每个warp由warp调度器调度执行。一个块被划分为warps的方式总是相同的；每个warp包含连续的、递增的线程ID，第一个warp包含线程0。线程层次结构[36]描述了线程ID如何与块中的线程索引相关。

一个warp一次执行一个公共指令，所以当一个warp的所有32个线程执行路径都一致（没有分支），就会实现完全效率。如果warp的线程通过数据依赖的条件分支发散，warp会执行每个分支路径，禁用不在该路径上的线程。分支发散只在warp内部发生；不同的warps无论它们是否执行公共或不同的代码路径都是独立执行的。

SIMT架构类似于SIMD（单指令，多数据）向量组织，其中一个单独的指令控制多个处理元素。一个关键的区别是SIMD向量组织将SIMD宽度暴露给软件，而SIMT指令指定单个线程的执行和分支行为。与SIMD向量机器相比，SIMT允许程序员为独立的、标量的线程编写线程级并行代码，以及为协调的线程编写数据并行代码。为了正确性的目的，程序员基本上可以忽略SIMT行为；然而，通过确保代码很少需要warp中的线程发散，可以实现大量的性能提升（however, substantial performance improvements can be realized by taking care that the code seldom requires threads in a warp to diverge）。在实践中，这类似于传统代码中缓存行的作用：在设计正确性时可以安全地忽略缓存行大小，但在设计峰值性能时必须考虑代码结构。另一方面，向量架构要求软件将加载合并为向量，并手动管理发散。

详细解释下上述这段：

为了确保程序的正确性，程序员不必深入了解SIMT（单指令多线程）的行为。也就是说，即使程序员不完全理解SIMT的工作原理，他们编写的代码仍然可以正确运行
在GPU中，一个warp是一组线程，这些线程同时执行相同的指令。但是，如果某些线程需要执行不同的指令（例如，由于条件分支），那么这些线程就会“发散”。当warp中的线程发散时，它们不能同时执行，这会降低性能。因此，为了实现最佳性能，程序员应尽量避免这种发散
传统的CPU编程中，缓存行是内存的一个块，它在物理上连续，并且作为一个整体被加载到缓存中。程序员在编写代码时，为了确保程序的正确性，不需要考虑缓存行的大小。但是，为了获得最佳性能，他们需要考虑如何组织数据，以便有效地利用缓存行，减少缓存未命中
向量架构，如SIMD（单指令多数据）架构，要求程序员明确地管理数据的并行性。在这种架构中，数据被组织成向量，并且一次操作可以同时应用于整个向量。为了获得最佳性能，程序员需要确保数据加载是“合并”的，这意味着数据应该在内存中连续存放。此外，程序员还需要手动管理线程或数据元素之间的发散

在NVIDIA Volta之前，warps使用一个由warp中的所有32个线程共享的单一程序计数器，以及一个指定warp的活动线程的活动掩码。因此，来自同一warp的处于发散区域或不同执行状态的线程不能相互发送信号或交换数据，而且需要锁或互斥体保护的数据的细粒度共享的算法很容易导致死锁，这取决于争用线程来自哪个warp。

从NVIDIA Volta架构开始，独立线程调度允许线程之间完全并发，无论warp如何。通过独立线程调度，GPU为每个线程维护执行状态，包括程序计数器和调用堆栈，并可以以每个线程的粒度产生执行，要么更好地利用执行资源，要么允许一个线程等待另一个线程产生的数据。一个调度优化器确定如何将来自同一warp的活动线程组合成SIMT单元。这保留了与之前的NVIDIA GPU中的SIMT执行相同的高吞吐量，但具有更大的灵活性：线程现在可以在子warp粒度上发散和重新汇合。

独立线程调度可能导致参与执行代码的线程集合与开发人员对先前硬件架构的warp同步性[37]做出的假设大不相同。特别是，任何warp同步的代码（例如无同步、warp内规约）都应重新检查，以确保与NVIDIA Volta及更高版本的兼容性。有关更多详细信息，请参阅Compute Capability 7.x[38]。

参与当前指令的warp的线程被称为活动线程，而不在当前指令上的线程是非活动的（已禁用）。线程可能因各种原因处于非活动状态，包括比其warp的其他线程更早地退出，采取与warp当前执行的分支路径不同的分支路径，或者是线程数不是warp大小的倍数的块的最后线程。如果由warp执行的非原子指令为warp的多个线程的相同位置在全局或共享内存中写入，那么发生在该位置的序列化写入的数量取决于设备的计算能力（参见计算能力5.x、计算能力6.x和计算能力7.x），以及执行最终写入的线程是未定义的。如果由warp执行的原子指令为warp的多个线程在全局内存中的相同位置读取、修改和写入，那么对该位置的每个读/修改/写都会发生，它们都会被序列化，但它们发生的顺序是未定义的。

Hardware Multithreading

由多处理器处理的每个warp的执行上下文（程序计数器、寄存器等）在warp的整个生命周期中都在芯片上维护。因此，从一个执行上下文切换到另一个执行上下文没有成本，每次指令发出时，warp调度器选择一个有线程准备执行其下一指令的warp（warp的活动线程[39]）并向这些线程发出指令。

特别地，每个多处理器都有一组32位寄存器，这些寄存器在warp之间进行划分，以及一个并行数据缓存或共享内存，该内存在线程块之间进行划分。

对于给定的内核，可以在多处理器上驻留和一起处理的块和warp的数量取决于内核使用的寄存器和共享内存的数量以及多处理器上可用的寄存器和共享内存的数量。每个多处理器还有一个最大的驻留块数和一个最大的驻留warp数。这些限制以及多处理器上可用的寄存器和共享内存的数量是设备的计算能力的函数，并在计算能力[40]中给出。如果每个多处理器没有足够的寄存器或共享内存来处理至少一个块，内核将无法启动。

一个块中的warp总数如下：

T 是每个块的线程数，
Wsize 是warp大小，等于32，
ceil(x, y) 等于x四舍五入到最接近的y的倍数。

为块分配的寄存器总数和共享内存总量在CUDA工具包中提供的CUDA占用率计算器中有记录。

术语warp-synchronous指的是隐式地假设在同一个warp中的线程在每个指令上都是同步的代码。

参考

https://docs.nvidia.com/cuda/cuda-c-programming-guide/#hardware-implementation

参考资料

[1] C++语言扩展中: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#c-language-extensions

[2] 使用nvcc进行编译: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compilation-with-nvcc

[3] CUDA运行时: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-c-runtime

[4] Driver API: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#driver-api

[5] 执行配置中更详细地描述: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#execution-configuration

[6] 应用程序兼容性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#application-compatibility

[7] CUDA环境变量中所述: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#env-vars

[8] 计算能力: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability

[9] 二进制兼容性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#binary-compatibility

[10] PTX兼容性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility

[11] 即时编译: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#just-in-time-compilation

[12] 异构编程: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#heterogeneous-programming

[13] 设备内存: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-memory

[14] 共享内存: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory

[15] 线程层次结构: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#thread-hierarchy

[16] 页锁定的主机内存: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#page-locked-host-memory

[17] 异步并发执行: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-concurrent-execution

[18] 多设备系统: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#multi-device-system

[19] 错误检查: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#error-checking

[20] 调用堆栈: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#call-stack

[21] 纹理和表面内存: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#texture-and-surface-memory

[22] 图形互操作性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#graphics-interoperability

[23] 上下文: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#context

[24] 即时编译: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#just-in-time-compilation

[25] 运行时和驱动API之间的互操作性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#interoperability-between-runtime-and-driver-apis

[26] 设备选择: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-selection

[27] 异构编程: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#heterogeneous-programming

[28] Texture and Surface Memory: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#texture-and-surface-memory

[29] SIMT架构: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#simt-architecture

[30] 硬件多线程: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#hardware-multithreading

[31] SIMT架构: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#simt-architecture

[32] 硬件多线程: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#hardware-multithreading

[33] 计算能力5.x: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-5-x

[34] 计算能力6.x: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-6-x

[35] 计算能力7.x: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-7-x

[36] 线程层次结构: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#thread-hierarchy

[37] warp同步性: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#fn2

[38] Compute Capability 7.x: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-7-x

[39] 活动线程: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#simt-architecture-notes

[40] 计算能力: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities

八、部署Windows环境

肯定是个老方法而且我训练机确实是win的所以i就搬来了, 其实我还是觉得在windows上部署挺简单的, 可能之前弄过迁移的就具简单尤其是py的就是显卡好了 python复制粘贴一下~ 感觉只做训练是够了

硬件和软件的最低要求

安装 Python 和所需工具 (这里说的是用conda 不过本人还是pip 感觉windows下还是好简单各种驱动支持也好)

设置开发环境

GPU 术语

安装 GPU 驱动

安装 TensorFlow（CPU 和 GPU）

安装 PyTorch（CPU 和 GPU）

验证安装情况

硬件和软件的最低要求

如果你要按照本指南操作并且计划使用 GPU，你必须使用英伟达 GPU。

开发深度学习应用涉及到训练神经网络，这自然需要执行大量计算。也因此，我们需要越来越多的并行运算，而 GPU 正好能够满足我们的需求。这也是当前 GPU 需求旺盛的主要原因之一。大多数深度学习框架都自带 GPU 加速支持，这让开发者和研究者无需执行任何 GPU 编程就能在几分钟内使用 GPU 进行计算。

大部分这些框架都（只）支持 CUDA，而这只能在英伟达 GPU 上使用，这也是你需要使用英伟达 GPU 的原因。但是，使用 AMD 的 GPU 也不是不可能，相关信息可参阅：https://rocmdocs.amd.com/en/latest/。

不过，就算你没有 GPU，也依然可以继续本教程。但为了有效进行深度学习，至少你要有好用的 CPU、内存和存储空间。

我的硬件——笔记本电脑的配置如下：

CPU——AMD Ryzen 7 4800HS 8C -16T@ 4.2GHz on Turbo

RAM——16 GB DDR4 RAM@ 3200MHz

GPU——Nvidia GeForce RTX 2060 Max-Q @ 6GB GDDR6 显存

对于硬件配置，我推荐至少使用 4 核 2.6 GHz 的 CPU、16GB 内存和 6GB 显存的英伟达 GPU。

另外，对于本教程，你当然需要使用 Windows 10 系统。我也假设你对 Python 软件包和环境具备基本认知。不管怎样，后面都会给出解释。

推荐使用的 Windows 版本是最新的 64 位 Windows 10 稳定版。

本教程假设你的操作系统是刚装好的，没有执行过额外的修改。不过只要你知道自己在做什么，依然可以参考本教程。

安装 Python 和所需工具

第一步当然是安装 Python。我建议使用 Mini-Conda 来安装 Python。先给刚入门的新手解释一下原因。

Conda 是一个软件包管理工具，可以帮助你安装、管理和移除各种不同的软件包。不过 Conda 并不是唯一的选择，还有 pip——这是我很喜欢的 Python 默认软件包管理工具。这里我们选择 Conda 的原因是在 Windows 上使用它更简单直接。

Anaconda 和 Mini-Conda 都是 Conda 的软件发行版，其中预安装了一些非常有用的数据科学 / 机器学习软件包，能节省很多时间。Anaconda 包含 150 多个在数据科学和机器学习中有用的软件包，基本上包含了你可能需要的一切，而 Mini-Conda 仅包含一些必需的工具和软件包。

我推荐使用 Mini-Conda，因为我喜欢对所安装的软件包有（几乎）完整的控制权。清楚地了解你所安装的东西完全不是坏事。当然这还能帮你节省一些存储空间，你也不会装上几十个你可能永远也用不上的奇怪软件包。

要安装 Mini-Conda，请访问：https://docs.conda.io/en/latest/miniconda.html

下载 Windows 64 位版本的 Python3 安装工具，然后像安装其它 Windows 软件一样安装它。一定要勾选询问你是否要将 Conda 和 Python 加入到 PATH 的勾选框。

现在你可以通过以下命令检查 Conda 和 Python 是否安装成功。如果安装成功，则会显示版本号；否则你可能需要再次正确安装 Mini-Conda 并将其加入到 PATH。

> python —versionPython 
3.8.3
> conda —versionconda 
4.8.4

下一步是安装 jupyter-notebook，请在命令行界面使用以下命令：

> conda install -y jupyter

你可以通过运行 jupyter notebook 来验证安装，这会帮你在浏览器上打开 Jupyter Notebook。

设置开发环境

这一步很重要，但很多人会忽视它。使用 Anaconda 这种包含所有已知软件包的工具是可以理解的，但如果要开发自己的项目，真正构建一些东西，你可能还是需要一个专门针对该项目或你的工作性质的定制开发环境。使用专门虚拟环境的另一大优势是你可以将软件包与全局设置隔离开。这样，就算你在该环境中使用软件包时搞错了，你也可以轻松地丢弃它们，而不对全局软件包产生任何影响。

这也能让你灵活地使用任何之前版本的 Python 创建环境。这样，你就可以避免使用那些还不稳定的新特性，之后再根据支持情况选择是否升级。

创建 Conda 环境还算简单。为了方便解释，我创建了一个名为 tensorflow 的环境，你可以将其改为任何名称。我将使用 Python 3.7，因为我知道 TensorFlow 对其有很好的支持。顺便一提，这将是安装 TensorFlow 的位置，我还会创建一个名为 torch 的环境来安装 PyTorch。

> conda create --name tensorflow python=3.7

环境创建完成之后，你可以使用以下命令进入该环境，其中的 tensorflow 只是我们之前提供给该环境的名称。

> conda activate tensorflow

进入环境之后，你会在提示框的左边看到类似这样的信息：

如果你没在 Powershell 上看到这个信息，那么你可能需要先在 Powershell 初始化 conda 一次：

> conda init powershell

然后，你可能会在左边看到 (base)，如上图所示，此时你已不在任何环境中。之后，你再进入任何环境，你应该都会看见环境名。

此外，你还可以在环境中安装 nb 工具，并将其链接到我们之前安装的 Jupyter Notebook。

> conda install nb_conda

要将该环境注册到 Jupyter Notebook，可运行以下命令：

> python -m ipykernel install --user --name tensorflow --display-name “Python 3.7 (tensorflow)”

要退出 Conda 环境，则运行以下命令：

> conda deactivate

现在按照同样的步骤创建一个名为 torch 的环境：

> conda create --name torch python=3.7
> conda activate torch> conda install nb_conda
> python -m ipykernel install --user --name torch --display-name “Python 3.7 (torch)”

如果环境设置成功，你可以在环境列表中看到它们。

> conda env list

要验证每个环境是否都已安装了各自的软件包，你可以进入各个环境，执行 conda list，这会显示该环境中已安装的所有软件包。

不要因为这个列表很长而感到困扰。Conda 已经妥善地处理了主要部分和依赖包。

GPU 术语

在安装 GPU 相关软件之前，我们有必要了解这些软件是什么，以及你需要它们的原因。

GPU 驱动：顾名思义，GPU 驱动是让操作系统及程序能使用 GPU 硬件的软件。游戏玩家肯定很熟悉这个。如果你喜欢打游戏，你可能需要让这个软件保持最新以获得最好的游戏体验。

CUDA：简单来说，这是英伟达开发的一个编程接口层，能让你调用 GPU 的指令集及其并行计算单元。

自 2010 年代末的 GeForce 8 系列 GPU 以来，几乎所有 GPU 都兼容 CUDA。要想了解你的 GPU 是否启用 CUDA，可以访问英伟达的网站。

举个例子，如果你有一台消费级 GPU，不管是 GeForce 系列还是 Titan 系列，你都可以在下图中看到你的 GPU 是否支持 CUDA。

数据截至 2020 年 9 月，截图仅含部分型号。

如果你的电脑是笔记本，你应该看右边的列表；如果你的电脑是台式机，你显然就该看左边的列表。

之前已经提到，我的 GPU 是右侧列表中的 RTX 2060 Max-Q。另外，你不必在意显卡型号名称是否与该列表中的名称完全匹配，Max-Q 和 Super 的底层架构一样，只在 TDP、CUDA 核及张量核数量方面有一些差异。

比如，不管你的 GPU 是 RTX 2080 Super 还是 2080 Max-Q 又或是 2080 Super Max-Q，看列表中的 RTX 2080 就够了。但如果你的 GPU 是 RTX 2080Ti 或其它加了 Ti 的型号，则说明你的 GPU 是该系列中最高端的那一款，这些 GPU 通常在显存大小和 CUDA 核及张量核数量方面更具优势。

截至 2020 年 9 月，要使用 TensorFlow 2.0，显卡计算能力必须高于 3.5，但建议使用计算能力至少为 6 的显卡以获得更好的体验。TensorFlow 2.0 还需要 CUDA 10 版本，而这又进一步要求驱动版本至少为 418.x。

PyTorch 需要的 CUDA 版本至少为 9.2，但也支持 10.1 和 10.2。所需的计算能力至少要高于 3.0。

CuDNN：即 CUDA Deep Neural Network 软件库，这是一个用于深度神经网络的 GPU 加速原语库。cuDNN 为前向和反向卷积、池化、归一化和激活层等标准例程提供了经过高度微调的实现。

（可选）TensorRT：NVIDIA TensorRT 是一套用于高性能深度学习接口的 SDK。其包含深度学习接口优化器和运行时优化器，能为深度学习接口应用提供低延迟和高通量的特性。

安装 GPU 驱动

首先，你需要搞清楚所使用的 GPU 型号，而且你的 GPU 必须启用了 CUDA。

如果你还没有安装驱动，你可能需要运行一次 Windows 更新，它会自动处理有用软件的安装过程，比如英伟达控制面板。这能帮助你获悉 GPU 的相关信息，还有一些与本文无关的设置。

英伟达控制面板就绪之后，你可以在开始菜单打开它，也可以右键点击桌面，然后选择英伟达控制面板。

打开之后，你可以点击「帮助→系统信息」来查看 GPU 驱动版本。驱动版本号列在「细节」窗口的顶部。

如上图所示，我的驱动版本是 456.x，远超过 418.x 的最低要求，所以我不必安装新驱动。

但你的电脑可能不是这样的。要安装最新版的驱动，可访问 https://www.nvidia.com/Download/index.aspx，然后输入 GPU 信息，下载合适的驱动。

驱动下载完成后，运行安装包，选择快速安装会更轻松。驱动安装完成之后，可使用英伟达控制面板进行验证。

另一个安装驱动的方法是使用英伟达的 GeForce Experience 应用程序。只要你购买的是主打游戏的电脑，应该都预装了该软件。安装过程很简单。

这一步是可选的。如果你已经按照上面的步骤安装了驱动，或你的电脑没有预装该软件，那就不用在乎这个步骤。

你可在这里下载该程序：https://www.nvidia.com/en-in/geforce/geforce-experience/，然后跟着安装流程将其安装到电脑上。安装完成，打开它，进入驱动选项卡，检查更新并安装新驱动。你也可以在该应用中查看驱动的版本号。

GeForce Experience 演示

现在安装驱动过程中最重要的步骤已经完成，你可以选择手动安装 CUDA 工具包，也可以选择在安装 TensorFlow 或 PyTorch 时留给 Conda 来安装（强烈推荐后者）。

如果决定手动安装，你可以从这里下载安装包：https://developer.nvidia.com/cuda-downloads，然后跟着指示操作即可。

安装 CUDA 工具包

CUDA 工具包装好之后，你可以在 cmd 或 Powershell 中执行 nvidia-smi 命令进行验证。

nvidia-smi 的输出当然还能nvidia-smi -L

windows上装驱动真的太简单了不像在linux下这么麻烦~~~

安装 TensorFlow

现在终于来到本教程的关键了。如果你已经完成了前述步骤，那么这一步会非常简单。

我们通过 Conda 来安装 TensorFlow 2.x。

要注意，首先进入我们之前创建的 tensorflow 环境，然后再进行操作。

> conda activate tensorflow

如果你需要 GPU 支持，就运行以下命令：

> conda install -c anaconda tensorflow-gpu

通过 anaconda 通道安装 TensorFlow 的 GPU 支持软件。使用 conda 而非 pip 安装 TensorFlow 的一大优势是 conda 的软件包管理系统。使用 conda 安装 TensorFlow 时，conda 还会安装所有必需和兼容的依赖包。这个过程是自动的，用户无需通过系统软件包管理器或其它方式安装任何其它软件。

其中也包含 TensorFlow 或 PyTorch 所需的版本合适的 CUDA 工具包。因此，使用 conda 能让这个过程变得非常简单。

我们只能在安装了 TensorFlow GPU 的环境中看到所安装的 CUDA 工具包。这既不会影响到全局系统的 CUDA 版本，同时也能满足 TensorFlow 和 PyTorch 的不同版本 CUDA 需求。这就是使用虚拟环境的最大好处，它能让不同的虚拟环境完全隔离开。

如果一切顺利，你不会在安装过程中看到任何报错信息。

要验证 TensorFlow 和所需的软件包是否成功安装，你可以执行 conda list，这会显示已安装软件包的列表，你应该能在其中找到与 TensorFlow 相关的软件包以及 CUDA 工具包。

你也可以打开 Python prompt 来验证是否已安装 TensorFlow。

>>> import tensorflow as tf
>>> tf.version
'2.1.0'

如果你看到了版本号，那么恭喜你，TensorFlow 已安装成功！任务完成。

在 Python prompt 中验证 TensorFlow 的安装情况

你在 Python prompt 中使用 TensorFlow 时可能会看到这样的信息：「Opened Dynamic Library」，但这并不是坏消息。这只是一条日志消息，说明 TensorFlow 可以打开这些软件库。

GPU 上的安装情况验证将在下文中介绍。

如果要安装仅使用 CPU 的 TensorFlow，你需要对安装命令进行简单的修改。

> conda install -c anaconda tensorflow

这将会安装没有 CUDA 工具包和 GPU 支持的 TensorFlow。

安装 PyTorch

安装 PyTorch 的过程与安装 TensorFlow 其实没太大差异。conda 让这一切都变得非常简单。

首先，进入我们创建的 torch 环境。

> conda activate torch

如果你想安装支持 CUDA 的 PyTorch，使用以下命令：

> conda install pytorch torchvision cudatoolkit -c pytorch

该命令会通过 Conda 的 PyTorch 通道安装兼容 CUDA 的 PyTorch。

至于仅使用 CPU 的 PyTorch，只需从以上命令中移除 cudatookit 即可：

> conda install pytorch torchvision cpuonly -c pytorch

这会安装无 CUDA 支持的 PyTorch。

和之前一样，你可以使用 conda list 验证安装情况，也可使用以下代码在 Python 上执行验证。

>>> import torch
>>> torch.version
'1.6.0'

如果返回版本号，则说明已成功安装 PyTorch。其实用whl也挺简单的~ windows上还是这么简单不和arm 和linux上一样安装麻烦

验证安装情况

有时候，你觉得一切都很顺利，准备开始使用这些工具时却遇到了一些重大错误。如果你正好遇到了这种情况，有可能是机器的问题，也可能是流程出错了，不能一概而论，要具体问题具体分析。

为了帮助你更好地验证安装情况，并确保 TensorFlow 和 PyTorch 使用的是指定的硬件，这里分享一些笔记。

你可以在 https://github.com/abhinand5/blog-posts 的 dl-setup-win 文件夹中找到它们。你可以克隆这些笔记然后运行其中的代码。如果返回的信息正确，你就可以放手开发了。

下图是该笔记的代码示例：

注：如果你没有从正确的环境启动 Jupyter Notebook，就可能会遇到一些错误。例如，如果你想使用 tensorflow 环境，你可以从 base 环境启动 notebook，然后将核改到 tensorflow 环境，但我在这样操作时遇到过报错。因此，如果你要运行 TensorFlow，就在 tensorflow 环境里启动 Notebook；如果你要运行 PyTorch，就在 torch 环境中启动 Notebook。不要从 base 或其它地方启动。

其他

还一种方式：https://developer.nvidia.com/cuda/wsl，其中涉及在 WSL（Windows Subsystem for Linux）中启用 CUDA 和英伟达驱动以便使用 GPU 来进行深度学习训练。目前这个功能还在预览阶段，但一旦官方发布，必将为深度学习实践者带来重大影响。这能将让人惊喜的 WSL 与 CUDA/GPU 驱动结合到一起。

我是windows上装docker 然后当开发环境和我们公司一样也挺方便的感觉不容易把系统弄坏感觉这样更好一些~~~

只是wsl稍微有点坑~~

九、CUDA-Free Inference for LLMs

blog链接：https://pytorch.org/blog/cuda-free-inference-for-llms/

1无CUDA的LLM推理

作者：Adnan Hoque, Less Wright, Raghu Ganti 和 Mudhakar Srivatsa

在这篇博客中，我们讨论了如何使用OpenAI的Triton语言实现流行的LLM模型（如Meta的Llama3-8B和IBM的Granite-8B Code）的FP16推理，其中 100% 的计算都是使用Triton语言完成的。对于使用我们基于Triton kernel的模型进行单个token生成的时间，我们能够在Nvidia H100 GPU上达到相对于CUDA kernel主导工作流的0.76-0.78x性能，在Nvidia A100 GPU上达到0.62-0.82x性能。

为什么要探索使用100%的Triton？Triton为LLM在不同类型的GPU（如NVIDIA、AMD，以及未来的Intel和其他基于GPU的加速器）上运行提供了一条路径。它还为GPU编程提供了更高层次的Python抽象，使我们能够比使用特定供应商的API更快地编写高性能kernel。在本文的其余部分，我们将分享我们如何实现无CUDA的计算，对单个kernel进行微基准测试以进行比较，并讨论我们如何进一步改进未来的Triton kernel以缩小差距。

图1. 在NVIDIA H100和A100上，Llama3-8B和Granite-8B的Triton和CUDA变体的推理吞吐量基准测试 设置：批量大小 = 2，输入序列长度 = 512，输出序列长度 = 256

2 Transformer块的组成

我们从Transformer模型中发生的计算分解开始。下图显示了一个典型Transformer块的“kernels”。

图2. 按核心kernels划分的Transformer块

Llama3架构的核心操作总结如下：

这些操作中的每一个都是通过在GPU上执行一个（或多个）kernels来计算的。尽管这些kernels的具体细节在不同的transformer模型中可能有所不同，但核心操作保持不变。例如，IBM的Granite 8B Code模型在MLP层中使用了偏置，这与Llama3不同。这种变化确实需要对kernels进行修改。一个典型的模型是由这些transformer块堆叠在一起，并通过嵌入层连接起来的。

3模型推理

典型的模型架构代码与一个由PyTorch启动的python model.py文件共享。在默认的PyTorch eager执行模式下，这些kernel都是使用CUDA执行的。为了实现Llama3-8B和Granite-8B端到端推理的100% Triton，我们需要编写和集成手写的Triton kernel，并利用torch.compile（生成Triton操作）。首先，我们用编译器生成的Triton kernel替换较小的操作，其次，我们用手写的Triton kernel替换更昂贵和复杂的计算（例如矩阵乘法和flash attention）。

Torch.compile自动为RMSNorm、RoPE、SiLU和Element Wise Multiplication生成Triton kernel。使用Nsight Systems等工具，我们可以观察这些生成的kernel；它们在矩阵乘法和注意力之间显示为微小的深绿色kernel。

图3. Llama3-8B 使用 torch.compile 的跟踪，显示用于矩阵乘法和 flash attention 的 CUDA kernels

对于上述跟踪，我们注意到在 Llama3-8B 风格的模型中，构成 80% 端到端延迟的两个主要操作是矩阵乘法和注意力 kernels，并且这两个操作仍然是 CUDA kernels。因此，为了缩小剩余的差距，我们用手写的 Triton kernels 替换了矩阵乘法和注意力 kernels。

4Triton SplitK GEMM Kernel

对于线性层中的矩阵乘法，我们编写了一个自定义的FP16 Triton GEMM（通用矩阵-矩阵乘法）kernel，该kernel利用了SplitK工作分解（https://pytorch.org/blog/accelerating-moe-model//#30-work-decomposition---splitk）。我们之前在其他博客中讨论过这种并行化方法，作为加速LLM推理解码部分的一种方式。

这里对上面博客中的 Work Decomposition - SplitK 一节也翻译一下

工作分解 - SplitK

我们之前已经证明，对于LLM推理中发现的矩阵问题大小，特别是在W4A16量化推理的背景下，通过应用SplitK工作分解(https://arxiv.org/abs/2402.00025)，GEMM内核可以加速。因此，我们通过在vLLM MoE kernel(https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py)中实现SplitK，开始了我们的MoE加速研究，这相对于数据并行方法产生了大约18-20%的加速。

这一结果表明，SplitK优化可以作为在推理设置中改进/开发Triton kernel的更公式化方法的一部分。为了建立对这些不同工作分解的直觉，让我们考虑一个简单的例子，即两个4x4矩阵的乘法和SplitK=2。

在下图中显示的数据并行GEMM kernel中，输出矩阵的单个块的计算将由1个线程块TB0处理。

Figure 2. Data Parallel GEMM

相比之下，在SplitK kernel中，计算输出矩阵中单个块所需的工作被“分割”或共享给两个线程块TB0和TB1。这提供了更好的负载均衡和增加的并行性。

Figure 3. SplitK GEMM

关键思想是我们将并行性从MN增加到MN*SplitK。这种方法确实会带来一些成本，例如通过原子操作增加线程块间通信。然而，这些成本相对于节省的其他受限GPU资源（如共享内存和寄存器）来说是微不足道的。最重要的是，SplitK策略为瘦矩阵（如MoE推理中的情况）提供了优越的负载均衡特性，并且在解码和推理期间是常见的矩阵配置文件。

5GEMM Kernel 调优

为了实现最佳性能，我们使用了穷举搜索方法来调优我们的SplitK GEMM kernel。Granite-8B和Llama3-8B的线性层具有以下形状：

Figure 4. Granite-8B and Llama3-8B Linear Layer Weight Matrix Shapes

这些线性层具有不同的权重矩阵形状。因此，为了获得最佳性能，Triton kernel必须针对每种形状配置进行调优。在对每个线性层进行调优后，我们能够在Llama3-8B和Granite-8B上实现1.20倍的端到端加速，相比于未调优的Triton kernel。

6Flash Attention Kernel

我们评估了一系列具有不同配置的现有Triton flash attention kernels，分别是：

AMD Flash(https://github.com/ROCm/triton/blob/triton-mlir/python/perf-kernels/flash-attention.py)
OpenAI Flash(https://github.com/triton-lang/triton/blob/main/python/tutorials/06-fused-attention.py)
Dao AI Lab Flash(https://github.com/Dao-AILab/flash-attention/blob/3669b25206d5938e3cc74a5f7860e31c38af8204/flash_attn/flash_attn_triton.py#L812)
XFormers Flash(https://github.com/facebookresearch/xformers/blob/fae0ceb195a41f2ab762d89449c6012fbcf2ffda/xformers/ops/fmha/triton_splitk.py#L96)
PyTorch FlexAttention(https://github.com/pytorch/pytorch/blob/e7b870c88bc3b854a95399a96a274d2f1f908172/torch/nn/attention/flex_attention.py#L800)

我们评估了每个kernel的文本生成质量，首先在eager模式下进行评估，然后（如果我们能够使用标准方法对kernel进行torch.compile）在编译模式下进行评估。对于kernel 2-5，我们注意到以下几点：

图5. 我们尝试的不同Flash Attention Kernels的组合表

上表总结了我们开箱即用的观察结果。我们预计通过一些努力，kernel 2-5可以被修改以满足上述标准。然而，这也表明，拥有一个适用于基准测试的kernel通常只是使其可用作端到端生产kernel的开始。我们选择在后续测试中使用AMD flash attention kernel，因为它可以通过torch.compile进行编译，并且在eager模式和编译模式下都能产生可读的输出。

为了满足AMD flash attention内核与torch.compile的兼容性，我们必须将其定义为torch自定义操作符。这个过程在这里有详细解释。教程链接讨论了如何包装一个简单的图像裁剪操作。然而，我们注意到包装一个更复杂的flash attention内核遵循类似的过程。两个步骤如下：

将函数包装成PyTorch自定义操作符

为操作符添加一个FakeTensor Kernel，该Kernel根据flash（q、k和v）输入张量的形状提供一种计算flash kernel输出形状的方法

在将Triton flash kernel定义为自定义操作符后，我们能够成功地为我们的端到端运行进行编译。

图6。在替换Triton matmul和Triton flash attention kernel后，Llama3-8B的torch.compile跟踪

从图6中，我们注意到，在整合了SplitK矩阵乘法kernel、torch操作符包装的flash attention kernel，并运行torch.compile后，我们能够实现一个使用100% Triton计算kernel的前向pass。

7End-to-End Benchmarks

我们在NVIDIA H100s和A100s（单GPU）上对Granite-8B和Llama3-8B模型进行了端到端测量。我们使用两种不同的配置进行了基准测试。

Triton kernel配置使用：- Triton SplitK GEMM - AMD Triton Flash Attention

CUDA kernel配置使用：- cuBLAS GEMM - cuDNN Flash Attention - Scaled Dot-Product Attention（SDPA）

我们发现在典型的推理设置下，eager模式和torch编译模式下的吞吐量和token间延迟如下：

图7。Granite-8B和Llama3-8B在H100和A100上的单token生成延迟，（批量大小 = 2，输入序列长度 = 512，输出序列长度 = 256）

总结来说，Triton模型在H100上可以达到CUDA模型性能的78%，在A100上可以达到82%。

性能差距可以通过我们在下一节中讨论的矩阵乘法和flash attention的kernel延迟来解释。

8Microbenchmarks

图8. Triton 和 CUDA kernel 延迟比较（Llama3-8B 在 NVIDIA H100 上）输入是一个任意提示（bs=1， prompt = 44 seq length），解码延迟时间

从上述结果中，我们注意到以下几点：

Triton 矩阵乘法 kernel 比 CUDA 慢 1.2-1.4 倍
AMD 的 Triton Flash Attention kernel比 CUDA SDPA 慢 1.6 倍

这些结果突显了进一步提高GEMM和Flash Attention等核心原语kernel性能的必要性。我们将其留作未来的研究，因为最近的工作（例如FlashAttention-3(https://pytorch.org/blog/flashattention-3/)，FlexAttention(https://pytorch.org/blog/flexattention/)）提供了更好地利用底层硬件的方法，以及我们希望能够在其基础上构建以实现更大加速的Triton路径。为了说明这一点，我们将FlexAttention与SDPA和AMD的Triton Flash kernel进行了比较。

我们正在努力验证FlexAttention的端到端（E2E）性能。目前，使用Flex进行的初步微基准测试在处理较长上下文长度和解码问题形状（其中查询向量较小）时显示出了良好的前景：

图9。在NVIDIA H100 SXM5 80GB上的FlexAttention kernel基准测试（批量大小=1，头数=32，序列长度=seq_len，头维度=128）

9Future Work

未来的工作计划包括探索进一步优化我们的矩阵乘法，以更好地利用硬件，例如我们发表的关于在H100上利用TMA的博客(https://pytorch.org/blog/hopper-tma-unit/)，以及不同的工作分解（如持久内核技术如StreamK等）以获得更大的速度提升。对于flash attention，我们计划探索FlexAttention和FlashAttention-3，因为这些kernel中使用的这些技术可以帮助进一步缩小Triton和CUDA之间的差距。我们还注意到我们之前的研究表明，FP8 Triton GEMM kernel性能在与cuBLAS FP8 GEMM相比时前景光明，因此在未来的帖子中，我们将探索端到端的FP8 LLM推理。

十、使用PTX指令更高效地加载和存储矩阵

使用PTX指令更高效地加载和存储矩阵

原始地址为：https://veitner.bearblog.dev/load-and-store-matrices-efficently-with-ptx-instructions/

本文实验cuda代码见：https://github.com/simveit/load_and_store

使用PTX指令更高效地加载和存储矩阵

ldmatrix

从PTX文档(https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-ldmatrix)中我们可以看到,ldmatrix可以用于从共享内存中集体加载一个或多个矩阵,以供mma指令使用。

指令格式如下

ldmatrix.sync.aligned.shape.num{.trans}{.ss}.type r, [p];  ldmatrix.sync.aligned.m8n16.num{.ss}.dst_fmt.src_fmt        r, [p];  
ldmatrix.sync.aligned.m16n16.num.trans{.ss}.dst_fmt.src_fmt r, [p];  .shape   = {.m8n8, .m16n16};  
.num     = {.x1, .x2, .x4};  
.ss      = {.shared{::cta}};  
.type    = {.b16, .b8};  
.dst_fmt = { .b8x16 };  
.src_fmt = { .b6x16_p32, .b4x16_p64 };

该指令将从.shared空间集体加载一个或多个矩阵到寄存器中。

p: 在.shared空间中的地址操作数
r: 目标寄存器
shape: 加载矩阵的形状
num: 加载1、2或4个矩阵

可能的数据类型如下：

注意，目前只有sm_100及更高版本的GPU支持m16n16和m8n16的形状。由于我目前没有访问权限，我们将专注于m8n8指令。

下表显示了每个线程组负责哪些矩阵。每个地址对应矩阵中的一行。每个"线程组"（即0-7、8-15、16-23和24-31）加载一个单独的矩阵。连续的行应该在内存中连续存储。

下图展示了使用ldmatrix加载的8x8矩阵的fragment布局：

// 使用64位地址加载一个8x8矩阵  
.reg .b64 addr;  
.reg .b32 d;  
ldmatrix.sync.aligned.m8n8.x1.shared::cta.b16 {d}, [addr];  // 加载两个8x8矩阵，以列主格式  
.reg .b64 addr;  
.reg .b32 d<2>;  
ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {d0, d1}, [addr];  // 加载四个8x8矩阵  
.reg .b64 addr;  
.reg .b32 d<4>;  
ldmatrix.sync.aligned.m8n8.x4.b16 {d0, d1, d2, d3}, [addr];

实现

如上所述，指针应该位于.shared空间中。有多种方法将通用指针转换为.shared空间。最简单的方法如下(https://forums.developer.nvidia.com/t/problem-about-ptx-instruction-cp-async-ca-shared-global/224219/2):

size_t asl = __cvta_generic_to_shared(smem+threadIdx.x);

我们也可以使用内联汇编：

asmvolatile(".reg .u64 smem_ptr64; cvta.to.shared.u64 smem_ptr64, %0;n" :: "l"(smem+threadIdx.x));

或者像这样：

asmvolatile(".reg .u64 smem_ptr64; cvta.to.shared.u64 smem_ptr64, %0;n" :: "l"(smem+threadIdx.x));   
asmvolatile(".reg .u32 smem_ptr32; cvt.u32.u64 smem_ptr32, smem_ptr64;n" ::);

我们也可以参考CUTLASS库(https://github.com/NVIDIA/cutlass/blob/ad7b2f5e84fcfa124cb02b91d5bd26d238c0459e/include/cute/arch/copy_sm75.hpp#L39)来获取实现思路。从这里开始,实现就比较直接了:

#include<cstdint>  
#include<iostream>  // 定义一个设备端内联函数，用于从共享内存加载8x8矩阵  
// d0: 输出参数，用于存储加载的数据  
// address: 输入参数，共享内存中的地址  
__device__ __forceinline__ voidldmatrix_sync_aligned_m8n8_x1_b16(  
uint32_t &d0, constuint32_t &address){  
// 使用内联PTX汇编指令加载矩阵  
// ldmatrix.sync.aligned.m8n8.x1.shared.b16: 同步加载8x8矩阵，每个元素16位  
// {%0}: 输出寄存器，存储加载的数据  
// [%1]: 输入寄存器，指定共享内存地址  
asmvolatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];"  : "=r"(d0)    // 输出约束，表示d0是一个输出寄存器  : "r"(address)); // 输入约束，表示address是一个输入寄存器  
}  // CUDA核函数，用于演示矩阵加载  
__global__ voidldmatrix(uint16_t *value){  
// 定义共享内存大小  
constexprint N = 64;  
// 声明共享内存数组  __shared__ uint16_t smem[N];  
// 获取当前线程ID  
auto tid = threadIdx.x;  // 计算行偏移量：每个线程负责8个元素，所以乘以8  
constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
// 计算最终地址：共享内存基址 + 行偏移  
constuint32_t address = __cvta_generic_to_shared(smem) + offset_rows;  // 初始化共享内存  
for (uint32_t i = tid; i < N; i += blockDim.x) {  smem[i] = i;  }  
// 同步所有线程，确保共享内存初始化完成  __syncthreads();  // 声明用于存储加载数据的变量  
uint32_t frag;  
// 调用矩阵加载函数  ldmatrix_sync_aligned_m8n8_x1_b16(frag, address);  // 再次同步，确保所有线程都完成加载  __syncthreads();  // 从32位数据中提取两个16位值  
// 提取低16位  
uint16_t number1 = static_cast<uint16_t>(frag & 0xFFFF);  
// 提取高16位  
uint16_t number2 = static_cast<uint16_t>((frag >> 16) & 0xFFFF);  
// 打印结果  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  
}  // 主函数  
intmain(){  
// 声明设备端指针  
uint16_t *d_value;  
// 分配设备内存  cudaMalloc(&d_value, sizeof(uint16_t));  
// 启动核函数，使用1个块，32个线程  ldmatrix<<<1, 32>>>(d_value);  
// 等待设备完成  cudaDeviceSynchronize();  
// 释放设备内存  cudaFree(d_value);  
return0;  
}

注意，根据上表，线程0-7需要对应于前8行的地址：

constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset_rows;

我们将在加载时传递地址和fragment。注意，每个fragment有32bit，我们可以通过先使用全16位掩码来提取最后16位，然后右移并再次执行相同的操作来提取前16位来输出加载的fragment。

0 -> 0  0   1     
1 -> 2  2   3     
2 -> 4  4   5     
3 -> 6  6   7     
4 -> 8  8   9     
5 -> 10  10   11     
6 -> 12  12   13     
7 -> 14  14   15     
8 -> 16  16   17     
9 -> 18  18   19     
10 -> 20  20   21     
11 -> 22  22   23     
12 -> 24  24   25     
13 -> 26  26   27     
14 -> 28  28   29     
15 -> 30  30   31     
16 -> 32  32   33     
17 -> 34  34   35     
18 -> 36  36   37     
19 -> 38  38   39     
20 -> 40  40   41     
21 -> 42  42   43     
22 -> 44  44   45     
23 -> 46  46   47     
24 -> 48  48   49     
25 -> 50  50   51     
26 -> 52  52   53     
27 -> 54  54   55     
28 -> 56  56   57     
29 -> 58  58   59     
30 -> 60  60   61     
31 -> 62  62   63

我们可以看到每个寄存器包含两个值。

我们可以在一个warp中同时写入两个矩阵。我们需要考虑到地址是按线程组提供的:

使用x2的ldmatrix语法如下所示

__device__ __forceinline__ voidldmatrix_sync_aligned_m8n8_x2_b16(  
uint32_t &d0, uint32_t &d1, constuint32_t &address){  
asmvolatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];"  : "=r"(d0), "=r"(d1)  : "r"(address));  
}

注意，现在我们写入32bit fragment。我们可以将这个包装成一个kernel:

__global__ voidldmatrix(uint16_t *value){  
constexprint N = 64;  __shared__ uint16_t smem[2 * N];  
auto tid = threadIdx.x;  constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 2) * 64;  
constuint32_t offset = offset_rows + offset_matrix;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;  for (uint32_t i = tid; i < N; i += blockDim.x) {  smem[i] = i;  smem[i + 64] = i + 64;  }  __syncthreads();  uint32_t frag1;  
uint32_t frag2;  ldmatrix_sync_aligned_m8n8_x2_b16(frag1, frag2, address);  __syncthreads();  uint16_t number1 = static_cast<uint16_t>(frag1 & 0xFFFF);  
uint16_t number2 = static_cast<uint16_t>((frag1 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  
}

计算地址的逻辑如下:

constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 2) * 64;  
constuint32_t offset = offset_rows + offset_matrix;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;

我们需要计算行偏移和矩阵偏移。前8个线程提供第一个矩阵的地址。接下来的8个线程提供第二个矩阵的地址。

我们可以非常相似地加载4个8x8矩阵。语法如下:

__device__ __forceinline__ voidldmatrix_sync_aligned_m8n8_x2_b16(  
uint32_t &d0, uint32_t &d1, uint32_t &d2, uint32_t &d3,  
constuint32_t &address){  
asmvolatile(  
"ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];"  : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3)  : "r"(address));  
}

完整的kernel如下所示:

__global__ voidldmatrix(uint16_t *value){  
constexprint N = 64;  __shared__ uint16_t smem[4 * N];  
auto tid = threadIdx.x;  constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 4) * 64;  
constuint32_t offset = offset_rows + offset_matrix;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;  for (uint32_t i = tid; i < N; i += blockDim.x) {  smem[i] = i;  smem[i + 64] = i + 64;  smem[i + 128] = i + 128;  smem[i + 192] = i + 192;  }  __syncthreads();  uint32_t frag1;  
uint32_t frag2;  
uint32_t frag3;  
uint32_t frag4;  ldmatrix_sync_aligned_m8n8_x2_b16(frag1, frag2, frag3, frag4, address);  __syncthreads();  uint16_t number1 = static_cast<uint16_t>(frag1 & 0xFFFF);  
uint16_t number2 = static_cast<uint16_t>((frag1 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  
uint16_t number3 = static_cast<uint16_t>(frag2 & 0xFFFF);  
uint16_t number4 = static_cast<uint16_t>((frag2 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 64]), (int)number3,  (int)number4);  
uint16_t number5 = static_cast<uint16_t>(frag3 & 0xFFFF);  
uint16_t number6 = static_cast<uint16_t>((frag3 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 128]),  (int)number5, (int)number6);  
uint16_t number7 = static_cast<uint16_t>(frag4 & 0xFFFF);  
uint16_t number8 = static_cast<uint16_t>((frag4 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 192]),  (int)number7, (int)number8);  
}

地址计算类似。我们再次有8个线程组，每个线程组提供4个矩阵的8行地址，因此总共32个线程在warp中提供地址。

constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 4) * 64;  
constuint32_t offset = offset_rows + offset_matrix;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;

每个kernel可以像这样调用:

intmain(){  
uint16_t *d_value;  cudaMalloc(&d_value, sizeof(uint16_t));  ldmatrix<<<1, 32>>>(d_value);  cudaDeviceSynchronize();  cudaFree(d_value);  
return0;  
}

stmatrix

stmatrix是一个PTX指令，用于将一个或多个矩阵集体存储到共享内存中。

stmatrix.sync.aligned.shape.num{.trans}{.ss}.type [p], r;  .shape  = {.m8n8, .m16n8};  
.num    = {.x1, .x2, .x4};  
.ss     = {.shared{::cta}};  
.type   = {.b16, .b8};

如你所见，指令与ldmatrix类似。.m8n8在Hopper中可用，m16n8在Blackwell GPU中可用。

地址的提供方式与上面相同。这次我们提供地址来知道提供寄存器(s)的内容存储到哪个位置。

实现

一旦我们正确理解了上面的ldmatrix指令，实现并不困难。请确保您理解了上面的代码，然后再继续阅读。

下面的代码给出了一个简单的PTX指令包装器，并存储一个8x8矩阵。

__device__ __forceinline__ voidstmatrix_sync_aligned_m8n8_x1_b16(  
uint32_t &d0, constuint32_t &address){  
asmvolatile(  
"stmatrix.sync.aligned.x1.m8n8.shared.b16 [%0], {%1};n" ::"r"(address),  
"r"(d0));  
}

我们可以将这个包装成一个kernel:

__global__ voidstmatrix(uint16_t *value){  
constexprint N = 64;  __shared__ uint16_t smem[N];  
auto tid = threadIdx.x;  constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset_rows;  uint32_t frag = 0x00000000;  frag |= (tid * 2 + 0);  frag |= (tid * 2 + 1) << 16;  __syncthreads();  stmatrix_sync_aligned_m8n8_x1_b16(frag, address);  __syncthreads();  uint16_t number1 = static_cast<uint16_t>(frag & 0xFFFF);  
uint16_t number2 = static_cast<uint16_t>((frag >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  
}

大部分代码与上面类似。但这次我们定义一个fragment，并将其存储到共享内存中的地址。

下面的代码初始化一个32位无符号整数。我们将首先初始化前16位为2 * tid + 0，然后初始化最后16位为2 * tid + 1。这主要是为了与ldmatrix示例中的相同结果。

uint32_t frag = 0x00000000;  
frag |= (tid * 2 + 0);  
frag |= (tid * 2 + 1) << 16;

我们将片段存储到给定的地址。这将输出:

0 -> 0  0   1     
1 -> 2  2   3     
2 -> 4  4   5     
3 -> 6  6   7     
4 -> 8  8   9     
5 -> 10  10   11     
6 -> 12  12   13     
7 -> 14  14   15     
8 -> 16  16   17     
9 -> 18  18   19     
10 -> 20  20   21     
11 -> 22  22   23     
12 -> 24  24   25     
13 -> 26  26   27     
14 -> 28  28   29     
15 -> 30  30   31     
16 -> 32  32   33     
17 -> 34  34   35     
18 -> 36  36   37     
19 -> 38  38   39     
20 -> 40  40   41     
21 -> 42  42   43     
22 -> 44  44   45     
23 -> 46  46   47     
24 -> 48  48   49     
25 -> 50  50   51     
26 -> 52  52   53     
27 -> 54  54   55     
28 -> 56  56   57     
29 -> 58  58   59     
30 -> 60  60   61     
31 -> 62  62   63

这证实了我们的实现与上面的ldmatrix操作相反。存储到2或4个矩阵的实现非常相似:

__device__ __forceinline__ voidstmatrix_sync_aligned_m8n8_x2_b16(  
uint32_t &d0, uint32_t &d1, constuint32_t &address){  
asmvolatile(  
"stmatrix.sync.aligned.m8n8.x2.shared.b16 [%0], {%1, %2};" ::"r"(address),  
"r"(d0), "r"(d1));  
}  // CUDA核函数，用于演示矩阵存储  
__global__ voidstmatrix(uint16_t *value){  
// 定义共享内存大小  
constexprint N = 64;  
// 声明共享内存数组，大小为2*N以存储两个矩阵  __shared__ uint16_t smem[2 * N];  
// 获取当前线程ID  
auto tid = threadIdx.x;  // 计算行偏移量：每个线程负责8个元素，所以乘以8  
constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
// 计算矩阵偏移量：根据线程组(0-7或8-15)选择第一个或第二个矩阵  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 2) * 64;  
// 计算总偏移量  
constuint32_t offset = offset_rows + offset_matrix;  
// 计算最终地址：共享内存基址 + 总偏移  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;  // 初始化第一个数据片段  
uint32_t frag1 = 0x00000000;  
// 设置低16位为 2*tid  frag1 |= (tid * 2 + 0);  
// 设置高16位为 2*tid+1  frag1 |= (tid * 2 + 1) << 16;  // 初始化第二个数据片段  
uint32_t frag2 = 0x00000000;  
// 设置低16位为 2*tid+64  frag2 |= (tid * 2 + 0 + 64);  
// 设置高16位为 2*tid+65  frag2 |= (tid * 2 + 1 + 64) << 16;  // 同步所有线程，确保数据准备完成  __syncthreads();  // 调用矩阵存储函数，将两个片段写入共享内存  stmatrix_sync_aligned_m8n8_x2_b16(frag1, frag2, address);  // 再次同步，确保所有线程都完成存储  __syncthreads();  // 从第一个32位片段中提取两个16位值  
uint16_t number1 = static_cast<uint16_t>(frag1 & 0xFFFF);  // 提取低16位  
uint16_t number2 = static_cast<uint16_t>((frag1 >> 16) & 0xFFFF);  // 提取高16位  
// 打印第一个矩阵的结果  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  // 从第二个32位片段中提取两个16位值  
uint16_t number3 = static_cast<uint16_t>(frag2 & 0xFFFF);  // 提取低16位  
uint16_t number4 = static_cast<uint16_t>((frag2 >> 16) & 0xFFFF);  // 提取高16位  
// 打印第二个矩阵的结果  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 64]), (int)number3,  (int)number4);  
}

对于四个矩阵的情况

__device__ __forceinline__ voidstmatrix_sync_aligned_m8n8_x4_b16(  
uint32_t &d0, uint32_t &d1, uint32_t &d2, uint32_t &d3,  
constuint32_t &address){  
asmvolatile(  
"stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};" ::"r"(  address),  
"r"(d0), "r"(d1), "r"(d2), "r"(d3));  
}  __global__ voidstmatrix(uint16_t *value){  
constexprint N = 64;  __shared__ uint16_t smem[4 * N];  
auto tid = threadIdx.x;  constuint32_t offset_rows = sizeof(uint16_t) * (tid % 8) * 8;  
constuint32_t offset_matrix = sizeof(uint16_t) * ((tid / 8) % 4) * 64;  
constuint32_t offset = offset_rows + offset_matrix;  
constuint32_t address = __cvta_generic_to_shared(smem) + offset;  uint32_t frag1 = 0x00000000;  frag1 |= (tid * 2 + 0);  frag1 |= (tid * 2 + 1) << 16;  
uint32_t frag2 = 0x00000000;  frag2 |= (tid * 2 + 0 + 64);  frag2 |= (tid * 2 + 1 + 64) << 16;  
uint32_t frag3 = 0x00000000;  frag3 |= (tid * 2 + 0 + 128);  frag3 |= (tid * 2 + 1 + 128) << 16;  
uint32_t frag4 = 0x00000000;  frag4 |= (tid * 2 + 0 + 192);  frag4 |= (tid * 2 + 1 + 192) << 16;  __syncthreads();  stmatrix_sync_aligned_m8n8_x4_b16(frag1, frag2, frag3, frag4, address);  __syncthreads();  uint16_t number1 = static_cast<uint16_t>(frag1 & 0xFFFF);  
uint16_t number2 = static_cast<uint16_t>((frag1 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid]), (int)number1,  (int)number2);  
uint16_t number3 = static_cast<uint16_t>(frag2 & 0xFFFF);  
uint16_t number4 = static_cast<uint16_t>((frag2 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 64]), (int)number3,  (int)number4);  
uint16_t number5 = static_cast<uint16_t>(frag3 & 0xFFFF);  
uint16_t number6 = static_cast<uint16_t>((frag3 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 128]),  (int)number5, (int)number6);  
uint16_t number7 = static_cast<uint16_t>(frag4 & 0xFFFF);  
uint16_t number8 = static_cast<uint16_t>((frag4 >> 16) & 0xFFFF);  
printf("%d -> %d  %d   %d   n", tid, (int)(smem[2 * tid + 192]),  (int)number7, (int)number8);  
}

我们需要做的就是初始化更多的fragments。当存储到2个矩阵时我们需要提供2个fragments,当存储到4个矩阵时我们需要提供4个fragments。

结论

我希望这篇博客对以下方面有所帮助:

详细理解ldmatrix和stmatrix指令
通过观察这两条指令之间的对称性来加深理解。如果你想了解更多,可以参考PTX文档。我很乐意收到你对这篇博客的反馈,并讨论CUDA及相关话题。如果你想联系我,可以在LinkedIn上添加我(https://www.linkedin.com/in/simon-veitner-174a681b6/)。代码可以在我的代码仓库(https://github.com/simveit/load_and_store)中找到。

查看全文

http://www.xdnf.cn/news/690067.html