RK3588 - RKNN（Rockchip 神经处理单元）的逆向工程

本文翻译自https://jas-hacks.blogspot.com/2024/02/rk3588-reverse-engineering-rknn.html
RK3588 NPU 的内部操作和功能主要隐藏在名为RKNPU2的闭源 SDK 中。由于对大型语言模型 (LLM) 的兴趣以及对transform模型最佳矩阵乘法的追求，想了解 RKNPU SDK 新引入的矩阵乘法 API (rknn_matmul_run) 的实现。对技术参考手册（TRM）中 RKNN NPU部分的查看，揭示了该NPU没有做矩阵乘法的原生机制（内部是通过倒一手的方法），尤其是对于向量。

为了了解发生了什么，第一步是了解 NPU 的工作原理。 TRM 提供了详细的寄存器列表以及构成 NPU 的核心单元的简要概述。它特别缺乏有关对寄存器进行编程以执行操作的基本信息。例如没有关于根据数据格式（如int8 与 float16）或输入数据或权重的大小等因素导出或计算寄存器值的具体细节。此外，没有关于如何构建 NPU 执行管道的信息，不过我从之前对V831 NPU 的逆向工程尝试中获得了一些优势。但是即使掌握了这些知识，它仍然需要几个月的反复对数据流的广泛分析、遇到一些死胡同以及无数次逆向工程尝试，最后才整明白了如何激活 NPU 并让它执行简单的操作。

RK3588 NPU 似乎是NVDLA架构的远亲，因为一些术语相似，核心单元与 NVDLA 具有相似的功能和管线，尽管它们的命名不同。其中主要区别之一是我们可以给 NPU一个要执行的任务列表（RKNN 术语），然后等待完成。例如，如果我有一个由3层的神经网络，每层由conv + bias组成，那么可以将 3 个任务（每个任务执行conv + bias）以及必要的输入、权重和偏置值提供给 NPU，随后我们只需等待 NPU 通知完成即可。
在这里插入图片描述
上面显示的图像是从 TRM 中提取的，并稍微修改了，因为 TRM 中提供的描述与它们的图表并不完全一致，更重要的是寄存器命名约定也不完全一致。这是我的解释，每个 NPU 由三个不同的单元组成：

CNA - 卷积网络加速器（包括 CORE 矩形），在TRM中它指的是神经网络加速引擎。
DPU-数据处理单元
PPU - 平面处理单元

基于上述，NPU 主要是为运行传统的卷积神经网络（CNN）而设计的。这是因为 CNA 核心功能是通过输入图像或特征数据以及相应的权重来执行卷积。官方提供的大多数 RKNPU2 demo（例如 YOLOX、Mobilenet 和 ResNet）都证明了对 CNN 的重视。 CNA 输出可以定向到 DPU，在 DPU 中可以执行加法、乘法和 RELU 等逐元素运算。随后，DPU 的输出可以传送到 PPU，在 PPU 中执行最小、最大和平均池化等操作。此外还可以选择直接将数据传输到 DPU 或 PPU，而无需卷积步骤。

为了有效地执行卷积，CNA 采用乘法累加 (MAC) 运算。 CNA 的性能部分取决于所使用的 MAC 单元的数量。根据 TRM，对于单个 NPU 核心，MAC 操作的计数取决于输入数据类型：

每个周期 1024 个 int8 MAC 操作
每个周期 512 个 float16 MAC 操作

每个 MAC 单元缓存 1x1x16 weight bytes，对于 int 8，它有 16 个值可用，而对于 fp16，它减少到 8 个。我们需要 2 个 MAC 单元来执行 fp16，因此每个周期的操作量减少。内部特征和权重数据必须符合 rk的NC1HWC2 格式，其中 C2 是weight bytes。然后，所有 MAC 单元共享一个 1x1x16 cube的特征数据，方便将其发送到累加器以计算部分和。在更高级别用法中，CNA 似乎执行块操作（block operation），如我在测试中观察到的，MAC 缓存32 channels的fp16 weight data。因此需要在每个kernel groups中layout 32 个channels的权重（不太准确，请看原文）。

性能还受到输入和权重数据的访问时间的影响，CNA 包含称为卷积缓冲区 (cbuf) 的二级缓存。在上图中，384KB 板载内存部分用于此目的。重要的是，MAC 单元的数量加上 CBUF 会影响一个task中可以完成的卷积size。

有些人可能已经推断出矩阵乘法 API 本质上是通过 2D 卷积执行的。例如，我们将矩阵 A 视为 [M x K]，将矩阵 B 视为 [K x N]。矩阵A表示以Mx1xK（hwc）格式排列的特征数据，而矩阵B表示以1x1xNxK（hwck）格式排列的权重数据。因此，所得矩阵 C [M x N] 排列为 Mx1xN。我正在运行一个简单的demo，要求 NPU 执行矩阵乘法。我使用从 GGML 测试用例 ( test-mul-mat.cpp ) 派生的矩阵数据来验证输出是否正确。要运行测试，请查看我的github，遗憾的是我仍在 Rock-5b 上针对内核 5.10 进行测试。如果测试运行输出应如下所示和上面的屏幕截图。

rock@rock-5b:~/rk3588-npu/build$ ./matmul_4_36_16
drm name is rknpu - 20220829 - RKNPU driver
input dma is ffff8000, output dma is ffffa000, weights dma is ffff9000
Size of npu_regs 112
RKNPU_SUBMIT returned 0
=========================================================================================================1224.0 1023.0 1158.0 1259.0 1359.0 1194.0 1535.0 1247.0 1185.0 1029.0  889.0 1182.0  955.0 1179.0 1147.01216.0 1087.0 1239.0 1361.0 1392.0 1260.0 1247.0 1563.0 1167.0 1052.0  942.0 1214.0 1045.0 1134.0 1264.01125.0  966.0 1079.0 1333.0 1287.0 1101.0 1185.0 1167.0 1368.0  990.0  967.0 1121.0  971.0 1086.0 1130.0999.0  902.0 1020.0 1056.0 1076.0  929.0 1029.0 1052.0  990.0 1108.0  823.0  989.0  759.0 1041.0 1003.0
=========================================================================================================

关于逆向工程已经达到了一个阶段，了解在处理特征数据作为输入时影响卷积的大多数寄存器设置。主要的不确定性在于确定特征/权重数据的库大小，但是我希望可以推断出这一点。在投入大量时间分析 NPU 后，我们应该知道以下关键点：

NPU 内的所有数据指针（例如输入、权重、输出、任务列表）均为 32 位，并且必须引用物理内存。因此，这将内存范围限制为 4GB，使得利用具有 16/32GB 内存的主板供 NPU 使用是不可能的。此外，它还可能对 NPU 上执行的模型类型施加限制。
6 TOPS 的说法有待验证。虽然每个 NPU 核心的额定速度为 2 TOPS，但有些寄存器可能可以在 3 个NPU核心之间启用卷积。但是在分析SDK生成的数据流后，似乎没有使用过该功能。此外，DPU/PPU 单元似乎没有类似的功能，这会限制其可用性。在我看来，最有效的方法是将它们视为单独的核心并在每个核心上执行模型，但是要考虑到前面提到的内存限制。
SDK提供矩阵乘法API在某些方面代表了NPU的低效利用，因为他存在内存分配、内核调用以及指示 NPU 执行单个卷积的开销。理想情况下，NPU 的任务应该是执行多个操作并为这些操作提供所有提供的数据。一般来说这就是运行 CNN 模型（即 YOLOvX）时使用 NPU 的方式。这里需要注意的是，转换后的模型仅限于包含 NPU 支持操作的层。
两个 fp16 [512 x 512] 矩阵相乘的初始基准测试表明我可以在 1 毫秒左右的时间内完成。但是由于CBUF限制，这将变成2个task发送给NPU。不幸的是，当使用矢量数据作为输入时，这只是流程的一部分，耗时操作是将矩阵转换为特征和权重数据格式，如果在运行时完成，则输出反之亦然。我努力创建一个高度优化的转换例程，用于矢量到特征数据的转换。根据我的基准测试，对于 fp16 [512 x 512] 矩阵，此过程大约需要 2 毫秒。我估计需要 12-15 毫秒来执行上述矩阵的所有转换。理想情况下，应提前转换权重数据的矩阵以减少转换开销，并且如果可能的话，应保留以供重用。
我希望能够使用可编程内核来执行自定义操作。不幸的是，情况并非如此，您只能使用 OpenCL 作为替代方案。如果您需要在 OpenCL 和 NPU 之间调整数据，这就会带来挑战。

关于其他单元（DPU/NPU）还有更多需要发现，我将花时间来做这件事。最后，TRM v1.0 包含 RKNN 的许多差距和不一致，如果有人有更新的版本，我们将不胜感激。
在这里插入图片描述