conv2d_w5 模块技术深度解析

概述：这个模块解决什么问题？

conv2d_w5 是 MNIST 卷积神经网络中第二个卷积层的 AIE-ML 实现，负责将 64 通道的 8×8 特征图转换为 128 通道的 3×3 特征图。这是一个典型的深度学习计算密集型任务——一个 3×3 卷积核在 64 输入通道上滑动，生成 128 输出通道，计算量巨大。

想象一个工厂的生产线：输入是 64 层"原材料"（特征图），需要经过 128 组不同的"模具"（卷积核）加工，最终产出 128 层"成品"。问题在于，单个 AIE 内核无法在规定时间内完成全部计算，因此需要将工作并行拆分到多个内核上。

这个模块的核心设计洞察是：将 128 个输出通道均匀分配给 4 个 AIE 内核（kkA/kkB/kkC/kkD），每个内核负责计算 32 个输出通道。这种数据并行策略使得吞吐量提升 4 倍，同时通过流式连接（stream connection）将各内核的输出按正确顺序重组。

架构与数据流

整体架构图

关键组件角色

数据流动详解

阶段 1：输入分发（Input Fan-out） 输入特征图从 PL 端通过 input_plio 进入，写入 MT0 共享缓冲区。MT0 配置为多播模式——同一个数据源被广播到 kkA、kkB、kkC、kkD 四个内核。这类似于电视台信号被多个电视机同时接收。

阶段 2：权重加载（Weight Preload） 每个计算内核有独立的权重加载子图（wts_init_graph）。权重数据从 PL 端流入，被转换成 AIE 内核可直接访问的异步缓冲区格式。这种设计允许权重在计算开始前预加载，避免运行时内存带宽瓶颈。

阶段 3：并行计算（Parallel Computation） 四个内核同时执行相同的卷积算法，但操作不同的权重子集：

• kkA 使用 weights[0]，生成输出通道 0-31
• kkB 使用 weights[1]，生成输出通道 32-63
• kkC 使用 weights[2]，生成输出通道 64-95
• kkD 使用 weights[3]，生成输出通道 96-127

阶段 4：结果汇聚（Result Gathering） 内核之间通过流连接（stream connection）形成链式结构：

kkA.out[0] → kkB.in[2] → kkC.in[2] → kkD.in[2]

kkD 作为汇聚点，从前三个内核接收它们的结果，再与自己的计算结果按正确顺序交错写入 MT1。这种设计确保输出特征图的通道顺序符合预期（ch0, ch1, ch2, ... ch127）。

核心组件深度解析

1. dut_graph —— 顶层测试封装

位于 conv2d_w5_app.cpp，是仿真和测试的入口点。

class dut_graph : public graph {
public:
  conv2d_w5_graph                                    dut;
  input_plio                                         ifm_i;
  std::array<input_plio,4>                           wts_i;
  output_plio                                        ofm_o;
  // ...
};

设计意图：将 AIE 图与 PL IO 接口解耦。conv2d_w5_graph 是纯 AIE 逻辑，而 dut_graph 添加文件 I/O 能力用于仿真验证。这种分层使得同一套 AIE 图可以在不同系统集成方案中复用。

物理映射约束：

location<kernel>(dut.kkA) = tile(18,1);
location<kernel>(dut.kkB) = tile(19,1);
// ...
location<kernel>(dut.weights[0].kk) = tile(18,0);

注意权重加载内核与其对应的计算内核位于相邻的垂直位置（同列，行号相差 1）。这是有意为之——减少权重数据传输的延迟和布线复杂度。

2. conv2d_w5_graph —— 核心计算图

2.1 内核创建与调度配置

kkA = kernel::create_object<conv2d_w5A>();
source(kkA) = "conv2d_w5A.cpp";
runtime<ratio>(kkA) = 0.9;
repetition_count(kkA) = 4;

• runtime<ratio>(kkA) = 0.9：该内核占用其所在 AIE Tile 90% 的计算周期，预留 10% 给数据搬运和其他开销
• repetition_count(kkA) = 4：每次图运行迭代中，该内核执行 4 次（处理 4 张图片）

2.2 共享缓冲区与分块策略

MT0 = shared_buffer<bfloat16>::create({64*8*5},1,1);
write_access(MT0.in[0]) = tiling(bdw0);  // 写分块：线性展开
read_access(MT0.out[0]) = tiling(bdr0);  // 读分块：三维遍历

这里体现了 AIE 编程的关键技巧：读写分块不对称。PL 端以扁平化格式写入（64 通道 × 8 行 × 5 列），但 AIE 内核需要以 (layer, col, row) 的三维顺序读取。tiling_parameters 结构体精确控制这种转换：

tiling_parameters bdr0 = { 
    .buffer_dimension = {64,8,5},      // 逻辑维度：(通道,行,列)
    .tiling_dimension = {8,8,1},       // 每次读取 8 通道 × 8 行 × 1 列
    .tile_traversal = {
        {.dimension=0,.stride= 8,.wrap=8},  // 通道方向：步进 8，包裹 8 次
        {.dimension=1,.stride= 8,.wrap=1},  // 行方向：步进 8，包裹 1 次
        {.dimension=2,.stride= 1,.wrap=5}   // 列方向：步进 1，包裹 5 次
    }
};

2.3 连接拓扑

connect<>( ifm_i,          MT0.in[0] );
connect<>( MT0.out[0],     kkA.in[0] );  dimensions(kkA.in[0]) = { 5*8*64 };
connect<>( weights[0].wts_o, kkA.in[1] ); dimensions(kkA.in[1]) = { 18464 };

dimensions() 声明是编译时契约，告诉编译器每次内核调用期望的数据量。这使得编译器可以优化 DMA 传输大小和缓冲区分配。

3. 计算内核（kkA/kkB/kkC/kkD）

四个内核类（conv2d_w5A/B/C/D）遵循相同的基本结构，但在数据传递行为上有微妙差异。

3.1 公共结构

每个内核都有三个主要方法：

void grab_wts(input_async_buffer<bfloat16>& wts_i);      // 权重预取
void filter_3x3(input_buffer<bfloat16>& ifm_i, 
                input_async_buffer<bfloat16>& wts_i);    // 核心卷积
void pass_outputs(...);                                   // 结果传递
void run(...);                                            // 入口调度

3.2 状态机模式

void grab_wts(input_async_buffer<bfloat16>& wts_i) {
  if (state == 0) {
    wts_i.acquire();  // 仅在第一次调用时获取权重
    state = 1;
  }
}

state 成员实现了简单的单触发状态机。由于 repetition_count=4，run() 会被调用 4 次，但权重只需加载一次。这是计算摊销的经典模式——昂贵的权重加载成本被多次输入复用。

3.3 核心卷积算法

filter_3x3 方法是性能关键路径，展示了 AIE-ML 向量指令的密集使用：

// 向量化缓冲区声明
std::array<aie::vector<bfloat16,32>,4> buffc;  // 输入特征缓存
std::array<aie::vector<bfloat16,32>,3> buffw;  // 权重缓存
std::array<aie::accum<accfloat,16>,6>  acc;    // 累加器（6 组并行累加）

算法流程：

1. 外层循环：遍历 3 个输出行（rr）
2. 中层循环：遍历 32 个输出通道组，每组 4 通道（oc += 4）
3. 内层循环：遍历 64 个输入通道，每组 8 通道（ic += 8）
4. 最内层：3 次迭代处理 3×3 卷积核的每一行（tt）

MAC 运算模式：

acc[0] = mac_4x8_8x4(buffc[0], buffw[0], acc[0]);
acc[1] = mac_4x8_8x4(aie::shuffle_down_fill(buffc[0],buffc[1], 8), buffw[1], acc[1]);

mac_4x8_8x4 执行 4×8 矩阵与 8×4 矩阵的乘加运算。shuffle_down_fill 用于提取滑动窗口中的重叠数据——这是卷积计算的核心操作，避免了重复加载内存。

3.4 内核间的差异

kkD 的特殊性在于它是链式结构的终点，需要从流中接收前三个内核的结果并重新排序：

// kkD::pass_outputs 伪代码逻辑
for (i = 0; i < 768/16; i++) data0[i] = read_from_stream(pass_i);  // 来自 kkA
for (i = 0; i < 768/16; i++) data1[i] = read_from_stream(pass_i);  // 来自 kkB
for (i = 0; i < 768/16; i++) data2[i] = read_from_stream(pass_i);  // 来自 kkC
// data3 是 kkD 自己的计算结果

// 按行交错写入输出缓冲区
for (row = 0; row < 3; row++) {
    write(data0[row]);  // ch 0-31
    write(data1[row]);  // ch 32-63
    write(data2[row]);  // ch 64-95
    write(data3[row]);  // ch 96-127
}

这种设计确保输出特征图的通道顺序符合深度学习框架的预期。

4. wts_init_graph —— 权重预加载子图

template<class T, unsigned SIZE>
class wts_init_graph : public graph {
    kernel kk;
    input_port  wts_i;
    output_port wts_o;
    // ...
    connect<>( wts_i,      kk.in[0] );
    connect<>( kk.out[0],  wts_o );
    dimensions(kk.out[0]) = {SIZE};  // 18464 for conv2d_w5
};

这是一个可复用的模板组件，将流式输入转换为异步缓冲区输出。其核心目的是：

1. 解耦权重加载时序：允许权重在计算开始前预加载
2. 统一接口：为上层图提供一致的 output_async_buffer 接口

设计决策与权衡

1. 为什么选择 4 个内核而不是更多或更少？

计算负载分析：

• 总 MAC 操作数：\(3 \times 3 \times 64 \times 128 \times 3 \times 3 = 265,420,800\)（约 265M）
• 考虑输出尺寸 3×3，实际为 \(3 \times 3 \times 64 \times 128 = 73,728\) 每像素，共 9 像素 = 663,552 MACs

实际上从代码看，循环结构是：

• 3 行 × 32 通道组 × (64/8) 输入组 × 3 核行 = 3 × 32 × 8 × 3 = 2304 次内层迭代
• 每次迭代执行 6 次 mac_4x8_8x4 = 6 × 128 = 768 MACs
• 总计：2304 × 768 ≈ 176 万 MACs 每内核

4 内核并行后，每个内核约 44 万 MACs，配合 90% runtime ratio，在典型 AIE 频率下可以满足 MNIST 实时推理需求。

为什么不更多？

• 128 输出通道不能被 8 整除（会导致负载不均）
• 更多的内核意味着更复杂的汇聚逻辑和更大的片上缓冲区压力
• 4 内核正好填满一个 2×2 的 Tile 区域，布局规整

2. 流连接 vs 缓冲连接的选择

内核间使用流连接（connect<>(kkA.out[0], kkB.in[2])）而非缓冲连接，原因：

选择流连接是因为：

1. 四个内核是同步启动、同步执行的（相同 repetition_count）
2. 数据量是确定且固定的（每张图 768 个 bfloat16）
3. 避免了 MT1 之前引入中间缓冲区，节省片上内存

3. 权重异步缓冲区的必要性

void run(input_buffer<bfloat16>& ifm_i,
         input_async_buffer<bfloat16>& wts_i,  // 异步缓冲区
         ...)

使用 input_async_buffer 而非普通 input_buffer 是关键设计：

• 普通 buffer：每次内核调用都会触发新的 DMA 传输
• 异步 buffer：由 grab_wts() 显式控制获取时机，权重数据驻留本地

这使得 4 次 run() 调用复用同一套权重，大幅减少内存带宽需求。

4. single_buffer 的作用

single_buffer(kkA.in[1]);  // 权重输入端口使用单缓冲

默认情况下，AIE 编译器会为缓冲区插入双缓冲（ping-pong）以隐藏传输延迟。但对于权重数据：

• 数据量较大（18464 × 2 bytes ≈ 36KB）
• 只读一次，复用多次
• 已使用异步缓冲区语义

因此显式声明 single_buffer，节省 36KB 的宝贵片上内存。

新贡献者须知

关键常量与魔数解释

重要：修改这些数值会破坏数据流的一致性。如果需要调整（如改变输出通道数），必须同步修改：

1. conv2d_w5_graph.h 中的缓冲区大小和分块参数
2. 各内核中的循环边界（oc < 32, ic < 64 等）
3. wts_init_graph 的模板参数 SIZE

常见陷阱

陷阱 1：混淆 bfloat16 和 float 的字节数

alignas(32) bfloat16 data[3*8*128/4];  // conv2d_w5A/B/C

注意 /4 是因为 128 bit 向量可以容纳 8 个 bfloat16（16bit each），数组索引是按向量元素计算的。如果误以为是按字节，会导致缓冲区溢出。

陷阱 2：流连接的阻塞风险

如果 kkA 的计算速度显著快于 kkB，流缓冲区可能溢出。虽然当前设计中四个内核执行完全相同的计算，理论上是平衡的，但如果修改了某个内核的优化级别或插入了调试代码，可能破坏这种平衡。

诊断方法：在 x86sim 或 aiesim 中观察流深度（stream depth）指标。

陷阱 3：权重偏移量的硬编码

bias = aie::concat(
    aie::broadcast<bfloat16,8>(*(wts_i.data()+18432+oc)),
    // ...
);

18432 是 bias 数据在权重缓冲区中的起始偏移（假设 18464 总大小 - 32 个 bias = 18432）。如果权重格式改变（如增加/减少通道数），这个偏移必须同步更新。

陷阱 4：shuffle_T16_4x8 的模式依赖

bias = shuffle(bias, shuffle_T16_4x8);

shuffle_T16_4x8 是一个特定的置换模式，将 4 个 bias 值广播到正确的向量位置。这是一种针对 AIE 向量 ALU 数据通路的底层优化。修改 bias 加载逻辑时必须理解这个置换的语义。

调试建议

1. 使用 gen_vectors.ipynb：这个 Jupyter Notebook 用于生成测试向量。如果修改了网络结构，需要重新生成输入/权重/参考输出数据。

2. 分层验证：

   • 先单独验证 wts_init_graph 是否正确加载权重
   • 再验证单个内核（如 kkA）的输出
   • 最后验证完整的 4 内核链

3. 检查 location 约束：如果编译报错涉及路由或资源冲突，尝试调整 tile 位置。保持计算内核和权重内核的垂直邻近关系有助于减少布线压力。

相关模块

• conv2d_w1 —— 第一个卷积层（1×1 卷积）
• conv2d_w3 —— 第三个卷积层变体
• wts_init —— 权重预加载子图模板
• max_pooling2d_w2 —— 池化层实现
• dense_w7 —— 全连接分类层
• mnist —— 完整网络集成

总结

conv2d_w5 展示了 AIE-ML 编程的核心范式：数据并行分解 + 流式结果汇聚。通过将 128 输出通道拆分到 4 个内核，实现了计算吞吐量的线性扩展。内核间的流连接消除了中间缓冲区的内存开销，而异步权重缓冲区则优化了内存带宽利用率。

理解这个模块的关键心智模型是：想象四个工人并行组装产品的不同部件，然后通过传送带将部件按顺序传递给最后一个工人进行最终包装。每个工人有自己的工具（权重），但共享原材料（输入特征图），最终产出一个完整的产品（输出特征图）。