读 ncnn 源码（XXXIV）：fuse_memorydata_binaryop——将“常量”烘焙进“算子”

在 ncnnoptimize 的图优化篇章中，我们已经见证了多种“算子融合”（Operator Fusion），它们将两个计算层（如 Conv 和 BN）合并为一个。本篇，我们将探讨一种不同但同样重要的优化：将数据层（MemoryData）与计算层（BinaryOp）融合。

这在本质上是一种常量折叠 (Constant Folding) 或 常量传播 (Constant Propagation)。fuse_memorydata_binaryop 的核心任务是识别出 BinaryOp（如加、减、乘、除）的一个输入是来自 MemoryData 的标量（Scalar），并将这个标量“烘焙”到 BinaryOp 的参数中，将其从一个昂贵的“张量-张量”操作，降级为一个高效的“张量-标量”操作。

TL;DR

  1. 目标: 识别 BinaryOp 的一个输入是来自 MemoryData 层的单个标量值的模式。
  2. 优化: 将 BinaryOp 转换为其高效的**“标量模式”**（with_scalar = 1），并将该标量值（memorydata->data[0]）直接存入 BinaryOp 层的 b 参数中。
  3. 图修改: 断开 MemoryData 层与 BinaryOp 层的连接（binaryop->bottoms.erase(...)），并将 MemoryData 层标记为 "ncnnfused" 以便后续移除。
  4. 关键技巧 (RSUB/RDIV): 为处理非交换操作（减法、除法），如果标量是第一个操作数（如 scalar - tensor），BinaryOp 的操作类型会被巧妙地切换为反向操作（如 RSUB，即 tensor - scalar），从而保持数学等价性。
  5. 两遍执行 (Two-Pass):
    • Pass 1: 处理简单的 MemoryData -> BinaryOp 模式。
    • Pass 2: 处理更复杂的 MemoryData -> Split -> BinaryOp 模式，即一个标量通过 Split 层广播给多个 BinaryOp。此 Pass 会逐个融合 BinaryOp，并最终清理掉 Split 层和 MemoryData 层。

1. 融合动机：昂贵的“张量-标量”操作

在神经网络中，有时会需要对整个特征图执行一个标量运算，例如 output = feature_map * 0.5。在 ncnn 的计算图中，这通常表示为：

  1. 一个 MemoryData 层，它持有一个 Mat（w=1, h=0, c=0），其中只包含一个 float 值 0.5。
  2. 一个 BinaryOp (Mul) 层，它接收 feature_map 和 MemoryData 的输出作为两个输入。

BinaryOp 的通用实现需要处理两个张量输入，它必须加载两个 Mat，处理复杂的广播逻辑。但如果它事先知道第二个输入永远是一个标量，它就可以切换到一个高度优化的“标量模式” (with_scalar = 1)：

  • BinaryOp(Mat, Mat): 昂贵。需要加载两个（可能很大的）Mat，内存带宽压力大。
  • BinaryOp(Mat, float): 高效。只需加载一个 Mat，标量 b 可以被立即加载到 SIMD 寄存器的所有通道中，计算和访存开销都大大降低。

fuse_memorydata_binaryop 的使命就是自动完成这种从“昂贵模式”到“高效模式”的转换。

2. Pass 1: 简单的 MemoryData -> BinaryOp 融合

函数的第一个循环处理最简单的 MemoryData -> BinaryOp 直连模式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
int NetOptimize::fuse_memorydata_binaryop()
{
    const size_t layer_count = layers.size();
    for (size_t i = 0; i < layer_count; i++)
    {
        // 1. 模式匹配：找到 MemoryData
        if (layers[i]->type != "MemoryData") continue;
        int top_blob_index = layers[i]->tops[0];

        // 2. 模式匹配：找到一个 BinaryOp，其输入之一是 MemoryData 的输出
        size_t j = i + 1;
        for (; j < layer_count; j++) {
            if (layers[j]->type != "BinaryOp") continue;
            if (layers[j]->bottoms.size() != 2) continue;
            if (layers[j]->bottoms[0] == top_blob_index || layers[j]->bottoms[1] == top_blob_index)
                break; // 找到了
        }
        if (j == layer_count) continue;

        ncnn::MemoryData* memorydata = (ncnn::MemoryData*)layers[i];
        ncnn::BinaryOp* binaryop = (ncnn::BinaryOp*)layers[j];

        // 3. 核心条件：MemoryData 必须是一个标量
        if (memorydata->w != 1 || memorydata->h != 0 || memorydata->c != 0)
        {
            continue; // 不是标量，无法融合
        }

        // 4. 处理非交换操作 (RSUB / RDIV)
        int memorydata_index = 1; // 假设标量是第二个输入 (index=1)
        if (binaryop->bottoms[0] == top_blob_index) // 如果标量是第一个输入 (index=0)
        {
            int op_type = binaryop->op_type;
            if (op_type == ncnn::BinaryOp::Operation_ADD || op_type == ncnn::BinaryOp::Operation_MUL || ...)
            {
                memorydata_index = 0; // 交换律：A+B = B+A, A*B = B*A
            }
            else if (op_type == ncnn::BinaryOp::Operation_SUB) // B = A(scalar) - B(tensor)
            {
                binaryop->op_type = ncnn::BinaryOp::Operation_RSUB; // 切换为 B = B(tensor) - A(scalar)
                memorydata_index = 0;
            }
            else if (op_type == ncnn::BinaryOp::Operation_DIV) // B = A(scalar) / B(tensor)
            {
                binaryop->op_type = ncnn::BinaryOp::Operation_RDIV; // 切换为 B = B(tensor) / A(scalar)
                memorydata_index = 0;
            }
            else { continue; } // 其他非交换操作 (如 ATAN2, POW) 无法简单切换
        }

        // 5. "烘焙" 标量参数
        float scalar = memorydata->data[0];
        binaryop->with_scalar = 1; // 开启标量模式
        binaryop->b = scalar;      // 将标量值存入 BinaryOp 的 'b' 参数

        // 6. 图结构修改
        fprintf(stderr, "fuse_memorydata_binaryop %s %s\n", ...);
        // 移除 BinaryOp 对 MemoryData blob 的依赖
        binaryop->bottoms.erase(binaryop->bottoms.begin() + memorydata_index);
        memorydata->type = "ncnnfused"; // 标记 MemoryData 为无效
    }
    // ... 第二个 Pass ...
}

关键点：

  • memorydata->w != 1 ...: 精确识别标量。
  • RSUB / RDIV 切换: 这是处理非交换律操作的精妙之处，使得融合得以在更多情况下（即使标量是第一个操作数）发生。
  • with_scalar = 1: 这是 BinaryOp 层内部实现的“开关”，forward 函数会根据此标志，选择 (Mat, Mat) 还是 (Mat, float) 的计算路径。

3. Pass 2: 处理 MemoryData -> Split -> BinaryOp 模式

在某些网络中，一个常量（如 MemoryData）可能被 Split 层复制多份，供多个后续层使用。第一个 Pass 无法处理这种情况，因此需要第二个 Pass。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
// 第二个 for 循环
for (size_t i = 0; i < layer_count; i++)
{
    if (layers[i]->type != "MemoryData") continue;
    int top_blob_index = layers[i]->tops[0];

    // 1. 模式匹配：找到 MemoryData -> Split
    size_t j0 = i + 1;
    // ... (find Split j0) ...
    if (j0 == layer_count) continue;

    int split_top_blob_index = -1;
    size_t j1 = j0 + 1;
    // 2. 模式匹配：找到 BinaryOp j1，其输入来自 Split
    for (; j1 < layer_count; j1++)
    {
        // ... (find BinaryOp j1 consuming one of Split's outputs) ...
        for (int k = 0; k < (int)layers[j0]->tops.size(); k++) {
            if (layers[j1]->bottoms[0] == layers[j0]->tops[k] || layers[j1]->bottoms[1] == layers[j0]->tops[k]) {
                split_top_blob_index = k; // 记住是 Split 的第 k 个输出
                break;
            }
        }
        if (split_top_blob_index != -1) break; // 找到了
    }
    if (j1 == layer_count) continue;

    // ... (获取 memorydata, split, binaryop 指针) ...

    // 3. 核心条件：依然是检查 MemoryData 是否为标量
    if (memorydata->w != 1 || memorydata->h != 0 || memorydata->c != 0) continue;

    // 4. 处理非交换操作 (与 Pass 1 完全相同)
    // ... (RSUB / RDIV logic) ...

    // 5. "烘焙" 标量参数 (与 Pass 1 完全相同)
    float scalar = memorydata->data[0];
    binaryop->with_scalar = 1;
    binaryop->b = scalar;

    // 6. 图结构修改 (更复杂)
    fprintf(stderr, "fuse_memorydata_binaryop %s %s\n", ...);
    // a) 移除 BinaryOp 对 Split blob 的依赖
    binaryop->bottoms.erase(binaryop->bottoms.begin() + memorydata_index);
    // b) 从 Split 层的输出列表中移除这个 blob
    split->tops.erase(split->tops.begin() + split_top_blob_index);

    // c) 关键：如果 Split 层的输出列表空了，说明它不再有任何消费者
    if (split->tops.empty())
    {
        split->type = "ncnnfused";      // 标记 Split 为无效
        memorydata->type = "ncnnfused"; // 标记 MemoryData 也为无效
    }

    i--; // *重要*：因为修改了 Split 层，可能导致层列表变化或需要重新检查
}

关键点:

  • Split 层处理: 融合的核心逻辑（BinaryOp 变为标量模式）不变，但增加了对 Split 层的清理逻辑。
  • split->tops.erase: 每融合一个 BinaryOp，Split 的一个输出就被剪除。
  • split->tops.empty(): 当 Split 的所有输出都被融合后，Split 层本身和其上游的 MemoryData 层就成了“孤岛”，可以被一并标记为 "ncnnfused"。
  • i--: 由于 layers 列表可能在迭代中被隐式修改（split 被融合），i-- 确保 for 循环在下一次迭代时能重新检查当前索引 i，防止跳过紧邻的下一个 MemoryData。

4. 结语

fuse_memorydata_binaryop 是一个精巧的、基于常量传播思想的优化 Pass。它通过将 MemoryData 标量“烘焙”进 BinaryOp 的参数中，将昂贵的“张量-张量”操作降级为高效的“张量-标量”操作。该 Pass 通过两遍扫描，分别处理了简单直连和通过 Split 广播的复杂情况，充分展现了 ncnnoptimize 在简化计算图、降低内存带宽压力方面的细致考量。