揭秘 Jetson 上的统一内存

本文整理自NVIDIA GTC2022讲座[SE2600]

我们知道Jetson是一个被称为集成 gpu 的产品，这意味着我们的 cpu 和一个 gpu 共享一个物理统一的内存结果，这与你可能熟悉的典型独立 gpu 完全不同，独立gpu 显卡有自己的内存与cpu、内存的系统分开，所以cpu、内存和gpu内存之间有很多迁移。它也恰好是典型独立GPU 计算的最大瓶颈之一。因此，当我们在编写项目时，我们真的应该考虑到一些阴暗面，因为这意味着我们需要合并很多不必要的开销，并牺牲很多潜在的性能提升。

本次会议专为数据科学家、研究人员、开发人员以及主要对为 Jetson 开发应用程序代码感兴趣的任何人设计，因此不需要硬件、CUDA 或嵌入式系统方面的特殊专业知识。在本课程结束时，您将了解是什么让 jetson 内存架构如此特别，以及与典型的独立 gpu 设置有何不同。我们将讨论统一内存的含义，它以几种不同的方式使用，最后从所有这些知识中得到的实际收获是如何调整 Python 代码以在 jetson 上运行，我们将从一个简单的向量加法示例，然后看一些更复杂或更实用的东西，如何用 TensorRT 优化神经网络做推理。因为我们将看到很多 Python 代码示例。

所以当我们提到一个独立gpu显卡时，这意味着我们通常有一个cpu，一个gpu，每个都有自己的内存，而cpu通常被称为主机，gpu是设备，我们的cpu和gpu由 pci express bus (PCIe)连接，你会看到连接这些组件的箭头有不同的粗细，这只是为了说明有不同的带宽，与我们在这些组件之间移动信息的速度有关。

当我们并行化程序时，我们知道 cpu 和 gpu 擅长不同的任务，所以通常它可能看起来像这样：我们有一个程序开始在 cpu 中运行，然后当你达到一些计算密集型功能时，您移至 GPU，一旦 gpu 完成计算，您必须将其移回 cpu 以继续使用该输出，因此我们继续执行其余的串行 cpu 代码。

我之所以强调这一点是因为这意味着每次我们从 cpu 切换到 gpu 时，都会有一个三步处理过程，因为 gpu 不会自动访问 cpu 可以访问的数据，所以首先我们必须将输入数据从 cpu 内存复制到 gpu 内存。然后我们加载我们的 gpu 程序并执行。我们会将结果从 gpu 内存复制回 cpu 内存，这三步过程会一遍又一遍地出现.

所以我们将使用 PyCUDA，我只是设置一个非常简单的示例：

这是一个典型的PyCUDA代码，

如上图所示，首先分配CPU端Input和output memory，然后分配GPU端Input和Output memory.第一步将数据从CPU传输到GPU，第二步，在GPU端做计算，这里调用一个函数，做计算：

第三步，将计算结果输出到CPU端。

以前需要释放GPU端meory:

但由于程序一开始，我们import pycuda.autoinit，所以这两行代码就不需要了。

下一个明显的问题是，如果我们要重复从 cpu 到 gpu之间来回进行这种内存复制，它看起来确实有点像样板代码，因为我们知道无论如何我们都必须这样做，所以一个明显的问题是，如果我们不必编写那些代码行不是很好吗，这就是最初引入的 cuda 统一内存背后的动机。

2012年CUDA6.0发布，首次引入cuda 统一内存。这是一种减少开发人员工作量的抽象，所以不像我们在拥有主机和设备之前看到的那样，让这两个并行指针指向相同的数据，他们将有一个单一的分配，一个单一的指针，以某种方式可用于主机和设备代码，消除了对我们看到的那些显式内存副本的需要。所以他们真的很重要，你看到你有两个独立的系统memory和 gpu memory 然后在这个统一的内存方法中，它只是一个抽象，所以它改变了developer view，好像主机和主机之间只有一个共享内存设备，但实际上它仍然是两个不同的物理存储。

那么，如果我们想以我们之前的例子为例，只使用统一内存，看看这是否会减少样板代码的行数。

这是我们之前代码示例，现在我们要做的是看看我们如何适应使用统一内存。

所以第一件事是我们不再需要为Host和device分配内存，

我们只需要有一个Managed Memory分配，我们如何做到这一点呢？我们使用 cuda.managed_empty，我们用我们的输入数组填充它，然后我们为输出创建一个缓冲区。你会看到我们有这个额外的参数 mem_flags 并且它总是会在我们今天介绍的所有代码示例中采用这个值所以我只是为了篇幅而不会包括它

第一步，从Host复制到Device，现在这步骤没有了。

第二步：执行GPU端代码，两边是一样的

第三步是内存复制从设备到主机，我们可以删除内存副本，但我们确实添加了强制同步调用。

我们之前没有提到过这一点，但是当你在gpu上运行时，cpu不会自动等待gpu完成，所以它会继续运行程序，即使gpu仍然很忙，所以我们需要告诉cpu等待gpu完成，因为否则我们不能在cpu代码中使用gpu的输出，所以内存拷贝设备会隐式同步并强制cpu等待，但当我们没有那个人工拷贝时，我们需要显式同步。

从上面的代码例子，我们正式完成了对原始代码的修改，使用cuda统一内存，你会看到我们减少了代码行的数量。我总结一下什么是统一内存，所以第一步和第三步虽然代码不存在，但它们仍然在执行任何关于底层硬件或设置的操作，只改变了developers view，因此 CUDA 统一内存首先是关于易于编程和程序员的生产力, 它主要不是一种使编写良好的 CUDA 代码运行得更快的技术，因为如果您具有 CUDA 专业知识，您可以指定很多数据移动数据分配以真正为应用程序定制，因此在大多数情况下，专业编写的 cuda 代码将仍然比统一内存做得更好，但它确实为程序员带来了便利和更少的开销。

Jetson的特殊地方在于它的"Unified Memory"是在物理层次真正统一的。我们之前用过的词语"Unified"的，实际上有两种意思。一种是从CUDA 6.X引入的概念上的，为了简化程序员的编码负担而引入的虚拟"Unified Memory": 这种如图，只是在概念和程序员的逻辑角度上是统一的，在物理存储上CPU的内存和GPU的显存依然是分开的。另外一种则是今天说的Jetson这种，从物理上，GPU和CPU共享同一个存储器芯片提供的内存/显存资源。这才是真正Jetson的GPU被称为集成GPU的原因。

统一内存在 Jetson 上特别酷，因为这次我们真的不需要复制数据，我们也不需要那些额外的分配从而节省宝贵的memory空间。因为现在讨论到了在物理芯片的层次是统一的，所以我们继续使用Unified Memory这个词，会造成困惑。因为（刚才）说过它有两方面的解释。所以我尽量不用这个词了，将CUDA Unified Memory, 我说成Managed Memory。用Managed Memory这种说法的时候，我将单指Jetson拥有的物理上统一的内存/显存存储，这样我们一说cuda managed memory，就是指jetson的.

下一步，我们将在向量加法示例中进一步探索这一点，我们通过定义一个非常大的输入数组 a 和一个非常大的输入数组 b 来设置它，我们将把元素相加以创建一个输出数组 C，

现在我们接着看蓝色的左侧代码，蓝色代码我们在独立GPU上的，使用原始的手工显存分配、数据复制过程。再看右侧绿色代码，我们还是使用独立GPU，改成使用Unified Memory。这个例子就像之前的那个例子一样，从一种做法，改成另外一种做法。

具体的说，左侧我们有2个输入，是Host上的A和B，外加一个C。然后呢，我们不仅仅要在Host上填充A，B, 还要分配设备上的A，B, C。你再看右侧，变到右侧，将A,B,C都改成用统一内存分配后，简单的填充了输入后，（就能跑kernel了），并没有Device上的A,B,C的分配过程了。

再看看具体的kernel使用，原本的三步走的，第一步传输到显存，第二步启动kernel，第三步再传输结果回来。其中第1步的内存复制，现在被消除了。第2步没变。第三步的设备数据回传Host，也被消除了。然后我们只需要再添加一步手工的显式同步就行了（原因我们之前提过）。

然后下图左侧，我去掉了一些空白（行），紧致了一下（代码）。所以左侧这么点东西，就是在独立显卡上，用Unified Memory的时候代码的样子了。然后接下来的问题，就是再到Jetson上，我们怎么用 Managed Memory了。需要注意，左侧的独立显卡情况下，因为GPU和CPU和各自的内存、显存，都是独立存在的，我们只是用Unified Memory将数据的复制从概念上给消除掉了，不是真消除了。我们再看右侧。我们想在Jetson上，不仅仅只是隐藏了这个，而是真的想要干掉这个数据移动的过程。而在Jetson上最酷的一点是，通过它的真物理统一架构，也就是Managed Memory这词，还是完全一样的用CUDA Unified Memory的代码。CUDA运行时就自动知道这是在Jetson上，自动去做了消除数据移动这种，非常正确的事情。和左侧的独立显卡还存在幕后的数据移动行为（形成了对比）。

Jetson是真就地使用了。无任何复制。

然后我们在NX上评测（基准）一下这个例子。注意我们的两个输入数组，每个都有16M个元素哈(1 << 24).然后将元素两两相加的话，如果在CPU上用常规的numpy ，得大约需要30ms+。然后如果我们用原始的独立显卡上的那种写法，有重复的内存/显存复制和数据传输的话，大约还是这么个时间。所以你知道数据的移动，抵消掉了性能上的提升，然后我们在用Unified Memory个例子，我们看到了大约6.5X的提速。更重要的需要说的事是，就算没有任何的性能提升，我们实际上将内存使用量(memory footprint)减半了，因为原始的(独立显卡）上的写法，我们实际上是创建了重复的内存/显存分配，并进行了多余的存储器（DRAM）芯片内部的倒腾，这点在Jetson上，这种内存容量有限的多的平台上，是一个大事

我们现在都说了这么多，还没算说到Jetson，这只不过说到了从非Unified, 到Unified Memory的通用做法。不管你在Jetson上运行，还是在独立显卡上运行。这种通用做法，上去就是改代码，改成单一次分配，改成使用单一的managed分配。这样我们就不需要两份分配在在CPU内存和GPU显存中的副本了。然后就是去消除掉host->device和device->host的复制传输。再然后第三点就是我们需要用某种形式的同步，确定GPU它搞完了(GPU was done), 再在CPU上干后续的处理。这应该用某种显式的同步或者其他类似的东西。

对于Jetson的在物理上的内存和显存统一的情况，除了使用Unified Memory, 我们还可以使用Pinned Memory. 我们将谈谈这个，注意这个和Unified Memory的主要区别是，Pinned Memory看起来需要很多背景专业知识，我们今天这里没法交代给你。但是我认为，对于你来说，知道Pinned Memory是什么，和Managed Memory有何不同，以及，怎么在代码里用它，还是非常重要的。

为了理解Pinned Memory, 我们先回到最开头的时候样子开始，先不讨论机器有独立显卡，有CPU，和它们各自的显存和内存。我们光从CPU开始，如果你知道现代操作系统的内存管理机制的话，你就肯定知道Pinned Memory. 具体的说，如果你像我一样，(不抵触了解这些的话），首先你应当知道OS的内存管理机制中的的虚拟内存。虚拟内存是什么呢？它有物理地址和虚拟地址之间的，映射功能的；还有有物理页面和逻辑页面管理功能。操作系统没必要将所有的页面，都保存在物理的内存中。操作系统总是在负责，某时某个逻辑页面被分配给了某个物理页面；总是在计划进程的对物理页面资源的使用优先与否。

在进程的地址空间上，也会有空白区域（white space), 也会允许你过度的超量分配内存。因为操作系统不仅仅代表了你（的进程），不停的在优化内存的分配使用；还会决策什么是重要不重要的。例如它会将最不重要的页面，例如你访问(touch)过的最遥远的时间之前的那些页面，交换到后备存储上去，例如，将内存页面交换到你的硬盘上去，并通过某种方式的地址本一样的东西，记录这种所被交换到磁盘上的页面。这种将内存页面进行换页的能力，也就是将页面在主系统内存和后背存储之间来回移动的能力，就叫做paging。

那我们说了这么多，又和GPU又什么关系呢？GPU（在传输数据）的时候，要使用一种叫DMA引擎的东西，也叫复制引擎(copy engine)，它负责了例如从系统的内存，移动数据到GPU的显存的任务。它工作的时候需要使用物理内存地址才能工作，而不能依赖于虚拟内存的地址。我们得强制要求一段内存地址被定住(pin)，才能保证这段内存具有一个不会改变的固定物理地址了。以及，CUDA使用这种被定住的，不能被换页出去的这种内存，会得到更快的传输速度和一些其他的非常酷的有用的后续效果，像是让一些操作能够并行执行，然而我们今天不准备具体说这些了。

我们不会非常深入的去说Pinned Memory和相关代码，如同刚才的Unified Memory那样，我想指出的是通常的大框是，你只需要对A,B,C分配1次即可，然后填充A,B，就能直接启动kernel了。但是这里有意思的一点是，你不能直接给PyCUDA用这里的A,B,C三个地址，你得调用一次get_device_pointer, 然后这样你有两个地址，分别是给CPU和GPU用的。注意我们只是这里得到了两个不同的地址，并非进行了两次分配，只是GPU上执行程序时候所用的一另外一个地址而已。需要指出并没有像Managed Memory那样的简单哈，需要大量的改动代码哈，我们研究3个GPU配置上的例子.

*译者注：在64-bit和UVA的环境下，无需单独获得一次设备上的地址，可以直接用的。实际上64-bit和UVA已经是现在的标配了（Since CUDA 3.2)，实际上用zero-copy(pinned memory)和unified memory在代码的书写形式上是完全一样的。。。

我们对比pinned memory和其他方式，加速比差不多。

这次我想做一张幻灯片，说说Managed Memory和Pinned Memory之间的比较，在与各种CUDA专家交谈并尝试回答这个问题之后，我发现基本上没有黄金法则,你必须分析你的代码，因为这对你的特定应用程序很重要，而且 CUDA for tegra 文档一直说我基本上是Pinned Memory或统一内存可用于减少数据传输开销，在最后一句话中，它说"评估影响以确定正确的内存选择"，所以没有黄金法则。如果你需要一些非常严重的优化，你总是可以尝试两者并分析代码。

下一个示例的重点是在独立gpu卡上编写代码，您如何调整该代码在 Jetson 上运行得最好，我们将看到的用 TensorRT 优化神经网络进行推理。TensorRT 是一个用于加速深度学习推理的 NVIDIA 库，所以它有两个主要组成部分，第一个是它需要一些经过训练的神经网络并将其转换为优化 TensorRT 引擎，基本上只是模型权重加上一些说明，就是如何在该模型上最佳运行，因此当您将其转换为TensorRT引擎时，嗯，它针对您指定的目标硬件进行了优化，您指定了必要的精度，因此基本上它包含了一个模型加上有关如何在最终目标硬件上运行它的说明。我的做法是我在 tensorflow 中训练了一个神经网络，利用该模型调整为 ONNX 格式，然后采用 ONNX 模型，并将其转换为 TensorRT 引擎，

为了今天的目的，你需要知道的是我们有 TensorRT 引擎，所以基本上一些小权重可以运行一些步骤，TensorRT 的第二个组件称为运行时，所以现在我们有了这个优化的 TensorRT 引擎，我们如何使用它进行推理，我们有一些选择，第一个是使用 NVIDIA 推理服务器：Triton ，还有TensorFlow API，很快就会使用 C++ 或只是使用普通 python，这就是我们今天要做的。

TensorRT 的示例与向量加法示例并没有什么不同，尽管它们会包含更多的代码，而且我真正想提请注意的是 TensorRT 引擎只是 CUDA 内核的一个特殊子集，所以想想它输入为批量图像的向量 a 和 b ，cuda 内核就是 TensorRT 引擎，现在输出将是我们批次中所有图像的预测类标签。虽然我们今天介绍的代码示例可以在这个 github 链接上找到，会提请你注意代码的哪些部分，但如果你真的想深入研究它，可以访问这里：https://github.com/annikabrundyn/JetsonUnifiedMemory

然后我们看看这个代码例子，几乎就是和GitHub上的TRT代码库里的例子一样。例子是为独立GPU写的，我们看看，基本上就是你从一个trt_engine.trt文件中，创建引擎。然后就接着创建缓冲区，基本上就在分配内存/显存和一点点其他的东西（创建流，还有bindings)等等。然后接着我要进行一次单独的推理调用。所以我们得将一个batch的输入图片，从CPU上传输到GPU。就是你看到的（图片中的for inp in inputs）的那些memcpy_htod_async(). 接着再下一步，我们就调用execute_async_v2()进行推理。然后第三步，我们将预测结果从GPU取回CPU，所以我们又有了一个memcpy_dtoh()。你看，我们还得有一步同步，一步显式的synchronize()调用，因为之前都是异步的(asynchronize)的memory copy（和推理执行）。

我们要进行的任务是从常规内存和显存分配，改成使用Managed Memory方案。所以我要提醒你的注意力，看这些代码，这些是我们要改动的，其他行先忽略不用看。刚才说过的，那个进行存储器分配和一点点其他事情的，那个allocate_buffer_dgpu()函数。（译者注：因为通篇在将如何改成Unified Memory/Managed Memory.所以其他事情就不用看。）

我们看这些黄色的行，就是原来缓冲区都是怎么分配的那里。这个循环，将循环每个输入和输出，（每次循环内部）都基本上是：用cuda.pagelocked_empty()创建内存缓冲区，然后它再分配一段显存上的缓冲区。所以我们要做的，就是将这两个，替换成单一的Managed分配。

剩下的就是如何改编这函数，让它继续还能工作。你看这里，用binding.append()添加设备缓冲区指针这行，我们得用Managed Memory的指针替换它。

完成从常规到Managed的替换，不仅仅要替换bindinds里的，还得替换inputs和outputs。然后还有一个部分就是说，我们在inputs和outputs里添加对应的内存和显存分配的那两行，得替换成单一的Managed Memory缓冲区。这样，我们就成功的从两份(duplicated)内存和显存分配，替换成了单一的Managed Memory分配了。就是图片里面的，黄色的部分变成了绿色了，就说明替换好了。