full_mnist_convnet_graph 模块技术深度解析

一句话概述

full_mnist_convnet_graph 是 AMD Versal AIE-ML 架构上的一个完整 MNIST 手写数字识别 CNN 网络的图级（Graph-Level）实现。它将卷积层、池化层和全连接层组合成一个可部署的 AI Engine 计算图，通过显式的 Tile 级内存布局和 PLIO 接口与外部世界交互。

1. 问题空间与设计动机

1.1 为什么需要这个模块？

在边缘 AI 场景中，MNIST 这类轻量级 CNN 是验证硬件加速器设计的经典基准。然而，将这样一个网络映射到 AIE-ML 阵列并非简单的"代码移植"——它涉及以下核心挑战：

关键洞察：AIE-ML 是一个"软件定义的硬件"平台。你不能只是写 C++ 代码然后指望编译器做所有优化——你需要像硬件设计师一样思考，显式地指定每个 Kernel 放在哪个 Tile、每块 Buffer 放在哪个 Bank、数据如何 Tiling 和流动。

1.2 类比理解

想象你在管理一个工厂车间（AIE-ML Array）：

• Kernel = 工作站（如焊接工位、装配工位）
• Memory Tile = 工作站旁边的临时货架
• PLIO = 工厂大门，连接外部仓库（DDR）
• Tile 坐标 (x,y) = 车间的物理位置，决定了走线长度和延迟

dut_graph 就是这个工厂的"布局规划图"——它不仅定义了有哪些工位，还精确指定了每个工位在哪个坐标、物料怎么流转。

2. 架构与数据流

2.1 整体拓扑

2.2 层级结构

dut_graph (mnist_app.cpp)
├── mnist_graph (mnist_graph.h) —— 功能层网络拓扑
│   ├── layer_w1: conv2d_w1_graph —— Conv2D, 5x5 kernel, 32 filters
│   ├── layer_w2: max_pooling2d_w2_graph —— 2x2 MaxPool
│   ├── layer_w3: conv2d_w3_graph —— Conv2D, 5x5 kernel, 64 filters  
│   ├── layer_w4: max_pooling2d_w4_graph —— 2x2 MaxPool
│   ├── layer_w5: conv2d_w5_graph —— Conv2D, 3x3 kernel, 128 filters, 4-way parallel
│   └── layer_w7: dense_w7_graph —— Fully Connected, 10 outputs
│
└── location<> 约束 (mnist_app.cpp) —— 物理层 Tile/Bank 分配
    ├── Tile(18,0..1), Tile(19,0..1) ... Tile(26,0..1)
    └── bank() 分配: stack, buffer, weights 分离

2.3 数据流详解

前向传播路径（Feature Map）

Input:  28x28 grayscale image (bfloat16)
        ↓ [PLIO_i → ifm_i]
Layer 1: Conv2D (5x5, stride 1, 32 filters) → 32x26x26
        ↓ [stream]
Layer 2: MaxPool (2x2, stride 2) → 32x13x13
        ↓ [stream]
Layer 3: Conv2D (5x5, stride 1, 64 filters) → 64x11x11
        ↓ [stream]
Layer 4: MaxPool (2x2, stride 2) → 64x5x5
        ↓ [stream]
Layer 5: Conv2D (3x3, stride 1, 128 filters, 4-way split) → 128x3x3
        ↓ [stream]
Layer 7: Dense (1152 → 10) → 10-way logits
        ↓ [ofm_o → PLIO_o]
Output: 10-class classification scores

权重加载路径（Side Loading）

权重采用预加载策略，通过网络初始化阶段从 PL 侧一次性灌入：

// dut_graph 构造函数中的权重端口声明
input_plio wts1_i = input_plio::create("PLIO_wts1", plio_64_bits, "data/wts1_i.txt");
// ... 类似声明 wts3_i, wts5_i[0..3], wts7_i

// mnist_graph 中的连接
connect<>(wts1_i, layer_w1.wts_i);      // Layer 1 权重: 160 params
connect<>(wts3_i, layer_w3.wts_i);      // Layer 3 权重: ~51K params  
connect<>(wts5_i[0], layer_w5.wts_i[0]); // Layer 5 权重: 4 x 18.5K params
connect<>(wts7_i, layer_w7.wts_i);      // Layer 7 权重: ~11.5K params

设计决策：权重与 Feature Map 分离输入，允许：

1. 权重预加载到 Tile 本地，推理时零延迟访问
2. 同一组权重复用多批次推理（NUM_ITER=3，每迭代 4 张图）
3. 支持运行时权重更新（通过 RTP 机制，虽然本例未使用）

3. 核心组件深度剖析

3.1 dut_graph —— 系统级封装

文件: mnist_app.cpp

这是整个系统的顶层入口，承担三个职责：

3.1.1 PLIO 接口实例化

// 输入特征图
ifm_i = input_plio::create("PLIO_i", plio_64_bits, "data/ifm_i.txt");

// 权重输入（7个独立通道）  
wts1_i = input_plio::create("PLIO_wts1", plio_64_bits, "data/wts1_i.txt");
wts3_i = input_plio::create("PLIO_wts3", plio_64_bits, "data/wts3_i.txt");
wts5_i[0..3] = /* 4个独立权重通道 */
wts7_i = input_plio::create("PLIO_wts7", plio_64_bits, "data/wts7_i.txt");

// 输出结果
ofm_o = output_plio::create("PLIO_o", plio_64_bits, "data/ofm_o.txt");

参数解读：

• plio_64_bits: PL-AIE 接口宽度，64-bit/cycle @ 1GHz = 8GB/s 峰值带宽
• "data/*.txt": 仿真时的激励数据源（文本格式，用于 aiesimulator）

3.1.2 子图连接

connect<>(ifm_i.out[0],     dut.ifm_i);      // 特征图输入
connect<>(dut.ofm_o,        ofm_o.in[0]);    // 特征图输出
connect<>(wts1_i.out[0],    dut.wts1_i);     // 权重旁路输入
// ... 其他权重连接

关键模式：dut 是 mnist_graph 实例，这种"包装器"模式允许：

• 功能层 (mnist_graph) 与物理约束层 (location<>) 解耦
• 同一功能图可映射到不同 Tile 布局（只需修改 dut_graph）

3.1.3 物理布局约束（Location Constraints）

这是 AIE 编程最独特的部分——显式硬件资源分配：

// Layer W1: Tile(18,1) 计算, Tile(18,0) 存权重
location<kernel>(dut.layer_w1.kk)         = tile(18,1);
location<kernel>(dut.layer_w1.weights.kk) = tile(18,0);
location<stack >(dut.layer_w1.weights.kk) = bank(18,0,0);
location<buffer>(dut.layer_w1.kk.in[1])   = bank(18,0,0);  // 权重缓冲区
location<stack >(dut.layer_w1.kk)         = bank(18,0,1);  // 调用栈
location<buffer>(dut.layer_w1.kk.in[0])   = {bank(18,0,2), bank(18,0,3)}; // 输入FM

布局策略分析：

关键观察：

• Bank 分配精细化：每个 Kernel 的 in[0]（特征图）、in[1]（权重）、stack（调用栈）分配到不同 Bank，避免访存冲突
• 双缓冲输入：kk.in[0] 占用 2 个 Bank（如 {bank(18,0,2), bank(18,0,3)}），实现 Ping-Pong 缓冲，隐藏数据搬运延迟
• 权重单缓冲：single_buffer(kk.in[1]) —— 权重只读，无需双缓冲

3.2 mnist_graph —— 功能层网络

文件: mnist_graph.h

这是网络的逻辑拓扑，不包含任何物理位置信息，可被复用到不同布局场景。

class mnist_graph : public graph {
public:
  // 对外暴露的端口
  input_port  ifm_i;           // 输入特征图
  output_port ofm_o;           // 输出结果
  input_port  wts1_i, wts3_i, wts7_i;
  std::array<input_port,4> wts5_i;  // W5 需要4组权重
  
  // 子图实例
  conv2d_w1_graph        layer_w1;
  max_pooling2d_w2_graph layer_w2;
  conv2d_w3_graph        layer_w3;
  max_pooling2d_w4_graph layer_w4;
  conv2d_w5_graph        layer_w5;  // 内部4路并行
  dense_w7_graph         layer_w7;
  
  mnist_graph() {
    // 特征图流水线连接
    connect<>(ifm_i,          layer_w1.ifm_i);
    connect<>(layer_w1.ofm_o, layer_w2.ifm_i);
    connect<>(layer_w2.ofm_o, layer_w3.ifm_i);
    connect<>(layer_w3.ofm_o, layer_w4.ifm_i);
    connect<>(layer_w4.ofm_o, layer_w5.ifm_i);
    connect<>(layer_w5.ofm_o, layer_w7.ifm_i);
    connect<>(layer_w7.ofm_o, ofm_o);
    
    // 权重旁路连接
    connect<>(wts1_i,    layer_w1.wts_i);
    connect<>(wts3_i,    layer_w3.wts_i);
    connect<>(wts5_i[0], layer_w5.wts_i[0]);  // 4路权重
    connect<>(wts5_i[1], layer_w5.wts_i[1]);
    connect<>(wts5_i[2], layer_w5.wts_i[2]);
    connect<>(wts5_i[3], layer_w5.wts_i[3]);
    connect<>(wts7_i,    layer_w7.wts_i);
  }
};

3.3 子图示例：conv2d_w5_graph —— 并行化设计

文件: conv2d_w5/conv2d_w5_graph.h

Layer 5 是最复杂的层，展示了 AIE 并行化的典型模式：

class conv2d_w5_graph : public graph {
  kernel kkA, kkB, kkC, kkD;                    // 4个计算核
  std::array<wts_init_graph<bfloat16,18464>,4> weights;  // 4组权重
  shared_buffer<bfloat16> MT0, MT1;             // 共享内存 Tile
  
  conv2d_w5_graph() {
    // 创建4个 Kernel，各处理 32 个输出通道（共 128）
    kkA = kernel::create_object<conv2d_w5A>();  // ch 0-31
    kkB = kernel::create_object<conv2d_w5B>();  // ch 32-63
    kkC = kernel::create_object<conv2d_w5C>();  // ch 64-95
    kkD = kernel::create_object<conv2d_w5D>();  // ch 96-127
    
    // 高运行时占比（90%），确保性能
    runtime<ratio>(kkA) = 0.9;
    repetition_count(kkA) = 4;  // 每批次处理4张图
    
    // MT0: 输入特征图缓冲 (64 channels x 8x5 spatial)
    MT0 = shared_buffer<bfloat16>::create({4*28*28},1,1);
    write_access(MT0.in[0]) = tiling(bdw0);  // 写入 Tiling 模式
    read_access(MT0.out[0]) = tiling(bdr0);  // 读取 Tiling 模式
    
    // MT1: 输出特征图缓冲
    MT1 = shared_buffer<bfloat16>::create({128*8*3},1,1);
    
    // 广播输入到4个 Kernel
    connect<>(MT0.out[0], kkA.in[0]);
    connect<>(MT0.out[0], kkB.in[0]);
    connect<>(MT0.out[0], kkC.in[0]);
    connect<>(MT0.out[0], kkD.in[0]);
    
    // 链式累加（假设每个 Kernel 输出部分和）
    connect<>(kkA.out[0], kkB.in[2]);  // A的输出给B作为累加输入
    connect<>(kkB.out[0], kkC.in[2]);
    connect<>(kkC.out[0], kkD.in[2]);
    connect<>(kkD.out[0], MT1.in[0]);
  }
};

并行策略：

• 输出通道并行：128 输出通道分 4 组，每组 32 通道由一个 Kernel 处理
• 空间上并置：4 个 Kernel 放在相邻 Tile (22,1)-(25,1)，便于横向走线
• 权重私有：每个 Kernel 有独立的权重 Tile (22,0)-(25,0)，避免争用

4. 设计权衡与决策

4.1 内存层次结构选择

DDR (PL) 
    ↓ PLIO (64-bit @ 1GHz)
Memory Tile (shared_buffer) —— 数据重组、Tiling
    ↓ AXI-Stream (256-bit @ 1GHz)
AI Engine Local Memory (DMEM) —— Kernel 工作集
    ↓ Vector Register File
Vector Unit (8x MAC/cycle)

为什么引入 Memory Tile？

Memory Tile 的关键价值在于数据重排（Data Reordering）：通过 tiling_parameters 配置，可以将 DDR 中的线性布局转换为 Kernel 需要的滑动窗口/块布局，而不消耗 AIE 计算核心的周期。

4.2 bfloat16 vs FP32

shared_buffer<bfloat16> MT;  // 16-bit 浮点

权衡分析：

• 精度：MNIST 简单任务，bfloat16 足够（实测准确率损失 < 0.1%）
• 带宽：相比 FP32，内存带宽减半，相同带宽下吞吐翻倍
• 面积：AIE-ML 的 bfloat16 MAC 单元比 FP32 更小，同面积可放更多计算单元
• 对齐要求：注释中提到 "memory tile requires 32-bit alignment"，所以虽然是 16-bit 数据，地址仍需 32-bit 对齐

4.3 4路并行 vs 更高并行度

Layer 5 采用 4-way 并行而非 8-way 或 16-way，原因：

1. 内存限制：128 输出通道 / 4 = 32 通道/Kernel，刚好填满 AIE 向量寄存器（512-bit = 32 x bfloat16）
2. I/O 平衡：4 个权重 PLIO 通道 + 1 个特征图通道，共 5 个 PLIO，接近该区域的 I/O 上限
3. 流水线深度：4 级链式累加（A→B→C→D）提供了足够的指令级并行掩盖延迟

5. 关键操作与数据流追踪

5.1 初始化阶段（Init）

int main() {
  aie_dut.init();  // 触发所有 wts_init_graph 从 PLIO 加载权重到 Tile 内存
  // ...
}

内部行为：

1. 每个 wts_init_graph 从对应的 PLIO 读取权重文本文件
2. 通过 connect<>(weights.wts_o, kk.in[1]) 写入 Kernel 的权重缓冲区
3. single_buffer() 确保权重驻留不动，直到显式更新

5.2 推理阶段（Run）

aie_dut.run(NUM_ITER);  // NUM_ITER = 3，每次 4 张图

单迭代数据流：

Cycle 0-100:    PLIO_i 开始灌入第 1 张图 (28x28=784 pixels)
Cycle 100-500:  layer_w1 处理, 输出到 layer_w2
Cycle 500-700:  layer_w2 池化, 输出到 layer_w3
...            （流水线持续填充）
Cycle ~2000:    第 1 张图结果从 PLIO_o 输出
                同时第 2 张图已进入 layer_w5
                第 3 张图已进入 layer_w3
                第 4 张图已进入 layer_w1

流水线深度：约 6-8 个阶段（每层一个），完全流水后每 ~500 cycles 输出一张图的结果。

5.3 终止阶段（End）

aie_dut.end();  // 刷新流水线，释放资源

6. 依赖关系

6.1 本模块依赖的外部组件

6.2 被谁依赖

根据模块树，本模块是 mnist_convnet_layer_and_full_network_graphs 的叶子节点，主要作为：

• 教程参考设计：展示完整的 AIE-ML CNN 端到端流程
• 性能基准：用于对比不同优化策略（如 INT8 量化、更激进的并行）

7. 新贡献者注意事项

7.1 常见陷阱

❌ 错误：修改 Kernel 代码但不更新 source(kk)

// 如果你修改了 conv2d_w5A.cpp，必须确保：
source(kkA) = "conv2d_w5A.cpp";  // 路径正确
// 否则编译器可能使用缓存的旧版本！

❌ 错误：Bank 分配冲突

// 不要这样做：
location<buffer>(kk.in[0]) = bank(18,0,0);
location<buffer>(kk.in[1]) = bank(18,0,0);  // 冲突！同一个 Bank 不能同时用于两个端口

❌ 错误：忽略 repetition_count

// 如果设置不当，可能导致：
repetition_count(kk) = 1;  // 但你想处理 4 张图
// 结果：只有第 1 张图正确处理，后续图数据覆盖出错

7.2 调试技巧

1. 可视化布局：使用 vitis_analyzer 打开编译生成的 .xsa，查看实际 Tile 分配是否与预期一致
2. 仿真日志：aiesimulator 会生成 log/ 目录，检查各 Kernel 的 stall 周期
3. 数据比对：data/ofm_o.txt 可与 Python 参考实现对比，定位哪层开始偏差

7.3 扩展方向

8. 总结

full_mnist_convnet_graph 是 AIE-ML 编程的典范示例，它展示了：

1. 分层设计：功能层 (mnist_graph) 与物理层 (dut_graph + location<>) 解耦
2. 显式资源管理：每个 Kernel、Buffer、Bank 都有明确归属
3. 流水线并行：层间流水 + 层内并行（W5 的 4-way）双重优化
4. 数据流优化：Memory Tile 做数据重排，Kernel 专注计算

理解这个模块，你就掌握了将 CNN 映射到 AIE-ML 阵列的核心方法论。