读 ncnn 源码（XXXV）：fuse_binaryop_eltwise——识别加权求和并替换为 Eltwise

在 ncnnoptimize 的图优化篇章中，我们已经见证了多种“算子融合”（如 Conv+BN）和“算子消除”（如 eliminate_dropout）。本篇，我们将探讨一种更高级的优化：基于模式识别的算子替换 (Pattern-based Operator Replacement)。fuse_binaryop_eltwise 就是其典型代表。

它的核心任务是识别出一种由多个 BinaryOp 层构成的、在语义上等价于**“加权求和”的计算模式，并将其替换**为一个功能更强、性能更高的 Eltwise 专用层。

TL;DR

  1. 目标: 识别 Y = (A * C0) + (B * C1) 这样的加权求和模式。A 和 B 是特征图（Tensor），C0 和 C1 是标量（Scalar）。
  2. 模式匹配: ncnnoptimize 查找一个 BinaryOp(Add, Tensor, Tensor) 作为中心，然后反向查找其两个输入 A 和 B 的生产者：
    • Case 1 (Full): (A * C0) + (B * C1)。A 的生产者是一个 BinaryOp(Mul, A, C0)，B 的生产者也是一个 BinaryOp(Mul, B, C1)。
    • Case 2 (Partial): (A * C0) + B 或 A + (B * C1)。只有一个输入来自 BinaryOp(Mul)。
  3. 替换为 Eltwise: ncnn::Eltwise 层（op_type = SUM）天生支持 Y = A*coeff[0] + B*coeff[1] 这种加权求和操作。它在一个 Kernel 内完成所有计算，效率远高于三个 BinaryOp（两次乘法，一次加法）的组合。
  4. 融合机制:
    • 创建一个新的 Eltwise 层，并将其 op_type 设为 Eltwise::Operation_SUM。
    • 将标量 C0 和 C1“烘焙”到 eltwise->coeffs（一个 Mat) 中。
    • 将 Eltwise 层的 bottoms（输入）直接连接到 A 和 B。
    • 替换原 BinaryOp(Add) 层，并将上游的 BinaryOp(Mul) 层标记为 "ncnnfused"。
  5. 性能优势: 此次替换将三个层、两次中间 Mat 读写（A*C0 和 B*C1 的结果）优化为了一个层、零中间 Mat 读写，极大地节省了内存带宽和调度开G销。

1. 融合动机：低效的加权求和

在许多网络结构中（如 FPN 特征融合、ResNet 的 shortcut 连接），加权求和 Y = A * C0 + B * C1 是一种常见操作。如果使用通用的 BinaryOp 来实现它，计算图会是这样的：

  1. BinaryOp_Mul_0 (with_scalar=1, b=C0) : X0 = A * C0
  2. BinaryOp_Mul_1 (with_scalar=1, b=C1) : X1 = B * C1
  3. BinaryOp_Add_0 (with_scalar=0) : Y = X0 + X1

这个过程存在巨大的效率问题：

  • 三次 Kernel Launch: 需要调度三个独立的层。
  • 两次中间 Mat 读写: X0 和 X1 作为中间结果，必须被完整地写入主存，然后再被 BinaryOp_Add_0 读回。这是巨大的内存带宽浪费。

2. Eltwise 层：更优的解决方案

ncnn 提供的 Eltwise 层是一个更强大的多输入元素操作层。当其 op_type 被设为 Eltwise::Operation_SUM 时，它可以接收多个（N 个）输入 blob，并执行：

Y=∑i=0N−1bottomblob[i]⋅coeffs[i]Y = \sum_{i=0}^{N-1} \text{bottomblob}[i] \cdot \text{coeffs}[i]

coeffs 是一个存储在 Mat 中的系数向量。

显然，Eltwise(A, B, coeffs=[C0, C1]) 在数学上与上述三个 BinaryOp 组成的子图完全等价。Eltwise 层的 forward 实现会在一个高度优化的循环（或 SIMD Kernel）中，同时读取 A 和 B，分别乘以各自的系数，然后相加，最后写入 Y。这消除了所有的中间 Mat 读写。

fuse_binaryop_eltwise 的任务就是识别出这个低效的 BinaryOp 组合，并将其替换为高效的 Eltwise 实现。

3. 代码实现：模式匹配与替换

fuse_binaryop_eltwise 的实现逻辑非常清晰：

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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
int NetOptimize::fuse_binaryop_eltwise()
{
    const size_t layer_count = layers.size();
    for (size_t i = 0; i < layer_count; i++)
    {
        // 1. 模式匹配：找到核心的 "ADD" 操作
        if (layers[i]->type != "BinaryOp") continue;
        if (layers[i]->bottoms.size() != 2) continue; // 必须是 2 输入
        ncnn::BinaryOp* binaryop = (ncnn::BinaryOp*)layers[i];
        if (binaryop->op_type != ncnn::BinaryOp::Operation_ADD) continue; // 必须是 ADD
        if (binaryop->with_scalar) continue; // 必须是 Tensor + Tensor

        // 2. 模式匹配：反向查找两个输入的生产者
        int bottom_blob_index_0 = binaryop->bottoms[0];
        int bottom_blob_index_1 = binaryop->bottoms[1];

        // 查找输入 0 的生产者 (j0)
        size_t j0 = 0;
        for (; j0 < i; j0++)
        {
            if (layers[j0]->type != "BinaryOp") continue;
            // 必须是 scalar mode (单输入) 的 MUL 操作
            if (layers[j0]->bottoms.size() != 1) continue; 
            if (((ncnn::BinaryOp*)layers[j0])->op_type != ncnn::BinaryOp::Operation_MUL) continue;
            if (layers[j0]->tops[0] == bottom_blob_index_0) break; // 找到了
        }
        // 查找输入 1 的生产者 (j1)
        size_t j1 = 0;
        // ... (与 j0 类似的循环) ...
        if (layers[j1]->tops[0] == bottom_blob_index_1) break;

        // 如果两个输入都不是 MUL, 则跳过
        if (j0 == i && j1 == i) continue; 

        // 3. 执行替换
        ncnn::BinaryOp* binaryop0 = (ncnn::BinaryOp*)layers[j0];
        ncnn::BinaryOp* binaryop1 = (ncnn::BinaryOp*)layers[j1];

        fprintf(stderr, "fuse_binaryop_eltwise %s %s %s\n", ...);

        // a) 创建新的 Eltwise 层
        ncnn::Eltwise* eltwise = (ncnn::Eltwise*)ncnn::create_layer_cpu("Eltwise");
        eltwise->type = "Eltwise";
        eltwise->name = binaryop->name;     // 继承 ADD 层的名称
        eltwise->bottoms = binaryop->bottoms; // 继承 ADD 层的输入 (稍后修改)
        eltwise->tops = binaryop->tops;       // 继承 ADD 层的输出
        ncnn::ParamDict pd; eltwise->load_param(pd);
        eltwise->op_type = ncnn::Eltwise::Operation_SUM; // 设置为 SUM
        eltwise->coeffs = ncnn::Mat(2);                  // 分配系数空间

        // b) 根据匹配到的模式，烘焙系数并重定向输入
        if (j0 != i && j1 != i) // Case 1: (A * C0) + (B * C1)
        {
            eltwise->coeffs[0] = binaryop0->b; // C0
            eltwise->coeffs[1] = binaryop1->b; // C1
            eltwise->bottoms[0] = binaryop0->bottoms[0]; // 连接到 A
            eltwise->bottoms[1] = binaryop1->bottoms[0]; // 连接到 B
            binaryop0->type = "ncnnfused"; // 标记 MUL0
            binaryop1->type = "ncnnfused"; // 标记 MUL1
        }
        else if (j0 != i && j1 == i) // Case 2: (A * C0) + B
        {
            eltwise->coeffs[0] = binaryop0->b; // C0
            eltwise->coeffs[1] = 1.f;          // C1 = 1.0
            eltwise->bottoms[0] = binaryop0->bottoms[0]; // 连接到 A
            // eltwise->bottoms[1] 保持不变 (已连接到 B)
            binaryop0->type = "ncnnfused"; // 标记 MUL0
        }
        else if (j0 == i && j1 != i) // Case 3: A + (B * C1)
        {
            eltwise->coeffs[0] = 1.f;          // C0 = 1.0
            eltwise->coeffs[1] = binaryop1->b; // C1
            // eltwise->bottoms[0] 保持不变 (已连接到 A)
            eltwise->bottoms[1] = binaryop1->bottoms[0]; // 连接到 B
            binaryop1->type = "ncnnfused"; // 标记 MUL1
        }

        // c) 在层列表中替换 ADD 层
        layers[i] = eltwise;
        delete binaryop; // 删除原 ADD 层
    }
    return 0;
}

4. 结语

fuse_binaryop_eltwise 是一个精妙的图优化 Pass。它超越了简单的层间融合，展现了 ncnn “理解” 计算图语义模式（Semantic Pattern）的能力。通过识别 (Tensor * Scalar) + (Tensor * Scalar) 这一常见模式，并将其替换为单一、高效的 Eltwise(SUM) 层，ncnnoptimize 能够将一个由三层 BinaryOp 构成的低效子图，优化为一个单一的高性能算子。

这不仅减少了 kernel launch 的开销，更重要的是消除了两次代价高昂的中间特征图内存读写，是优化内存带宽密集型网络（如 FPN）的关键步骤。