存算一体芯片算子融合策略实战 - 第2节

一、算子是什么？先别急着背定义

说实话，我第一次接触「算子」这个词的时候，也觉得挺玄乎的。后来做AI芯片久了，慢慢就明白了——算子就是神经网络里最基本的计算单元。

你想想看，一个神经网络模型，不管它多复杂，拆到最底层，无非就是那么几种操作：乘一下、加一下、取个最大值、做个卷积。这些基础操作，我们就叫它算子（Operator）。

我习惯把算子比作乐高积木。你搭一个城堡，用的就是那几种标准积木块。你搭一个ResNet-50，用的也就是Conv、ReLU、Add这些标准算子。

嗯，这里要注意：不同框架对算子的定义略有差异。PyTorch里可能叫torch.nn.functional.conv2d，TensorFlow里叫tf.nn.conv2d，但底层干的事是一样的。

二、算子的三种计算模式

我个人习惯把算子按计算模式分成三类。为什么这么分？因为不同的计算模式，对硬件的访存和计算要求完全不同。你在做算子融合的时候，必须搞清楚每个算子的「脾气」。

2.1 Element-wise（逐元素算子）

这种算子最简单，也最常见。说白了就是：输入和输出形状一样，每个输出元素只依赖对应的输入元素。

典型的例子：

• ReLU：y = max(0, x)
• 加法：y = a + b
• Sigmoid、Tanh 等激活函数

我在项目中遇到过一个问题：有人把ReLU单独实现成一个kernel，结果发现整个模型里ReLU占了30%的kernel启动开销。其实ReLU这种算子，计算量极小，瓶颈全在访存和kernel launch上。

2.2 Reduction（规约算子）

规约算子的特点是：输出比输入小，多个输入元素合并成一个输出元素。

常见的规约算子：

• Sum：对某个维度求和
• Max：取最大值
• Mean：求平均
• Softmax：虽然复杂点，但本质也是规约+逐元素

规约算子有个麻烦的地方——它需要跨行/跨列的数据依赖。比如你要对一个矩阵的每一行求和，那就得把整行数据都读到才能算。这对存算一体架构来说，意味着你需要更大的局部存储或者更聪明的数据流设计。

为什么会这样？因为规约操作天然就有「等待」的特性。你算到一半，不能输出结果，必须等所有数据到齐。这跟element-wise那种「来一个算一个」完全不同。

2.3 Convolution（卷积算子）

卷积是CNN的绝对核心。它的特点是：一个输出元素依赖输入的一个局部窗口。

举个例子，3x3卷积：

// 伪代码示意
for (oh = 0; oh < output_height; oh++) {
  for (ow = 0; ow < output_width; ow++) {
    for (kh = 0; kh < 3; kh++) {
      for (kw = 0; kw < 3; kw++) {
        output[oh][ow] += input[oh+kh][ow+kw] * weight[kh][kw];
      }
    }
  }
}

卷积的访存模式是有规律的重用。同一个输入像素，会被多个输出窗口用到。这就是为什么卷积天然适合做数据重用优化。

算子类型	计算密度	访存模式	融合友好度
Element-wise	极低	1:1 映射	★★★★★
Reduction	中等	多:1 映射	★★★☆☆
Convolution	高	窗口滑动	★★★★☆

三、为什么要做算子融合？

好，现在你知道了算子是什么，也知道了三种计算模式。那问题来了——为什么要把它们融合在一起？

我直接说结论：算子融合的核心动机，就是减少数据搬运。

你想想看，在传统架构里，每个算子算完，结果都要写回DDR。下一个算子再从DDR读出来。这一写一读，时间全花在数据搬运上了，真正算的时间反而没多少。

我做过一个实测：一个简单的Conv+ReLU+Add组合，如果不做融合，数据搬运占了总时间的75%。融合之后，搬运时间降到20%以下。

3.1 收益一：减少访存开销

这是最直接的收益。融合之后，中间结果不用写回DDR，直接在片上缓存里流转。对于存算一体架构来说，这个收益更明显——数据根本不用离开存储阵列。

3.2 收益二：减少kernel启动开销

每个算子对应一个GPU kernel或者一个硬件指令。启动一个kernel是有固定开销的——配置寄存器、加载指令、同步等等。你把N个算子融合成1个，启动开销就变成原来的1/N。

3.3 收益三：提升计算密度

计算密度 = 计算时间 / (计算时间 + 访存时间)。融合后访存时间减少，计算密度自然就上去了。对于存算一体架构，计算密度直接决定了TOPS/W（每瓦性能）的上限。

四、一张图看懂算子融合

下面这张图，是我自己总结的算子融合知识框架。你可以把它当作本章的思维导图：

五、小结

这一章我们聊了三件事：

1. 算子是什么——神经网络的基本计算单元，硬件的最小调度单位
2. 三种计算模式——Element-wise（逐元素）、Reduction（规约）、Convolution（卷积），各有各的访存和计算特点
3. 融合的动机与收益——减少访存、减少启动开销、提升计算密度

说白了，算子融合就是用空间换时间，用局部性换效率。在存算一体架构里，这个思想会被发挥到极致——因为数据根本不用离开存储，融合的收益比传统架构大得多。

下一章，我们会深入具体的融合策略，看看不同计算模式的算子到底怎么「粘」在一起。到时候我会拿几个真实项目里的踩坑经历跟你分享。