RISC-V端侧推理运行时设计指南：算子图融合、内存池与分块执行的协同设计

1. 引言：内存而非算力构成主要瓶颈

上篇讨论了量化的数值表示与算子融合，其中提到量化收益的主要来源是访存量下降。这一结论指向一个更基础的问题：在给定的内存预算下，一个计算图应当如何被组织成可执行的算子序列。

端侧设备的存储层次与数据中心差异显著：典型配置为数百 KB 至数 MB 的片上 SRAM 加数十 MB 外部 DRAM，两者带宽相差一个数量级。模型权重通常可整体放入 DRAM，但中间张量——算子之间的激活值——的数量与体积取决于图的拓扑结构而非参数量。一个参数量仅数 MB 的模型，按拓扑顺序逐算子执行时，中间张量的峰值占用可能达到数十 MB。

峰值内存的定义是：沿执行序列推进时，任一时刻同时存活的张量体积之和的最大值。设张量 t 从被生产者写出（first_use）到被最后一个消费者读取完毕（last_use）之间处于存活状态，则

peak = max over step k of Σ { size(t) : first_use(t) ≤ k ≤ last_use(t) }

降低该值的手段有三类：减少存活张量的数量（融合）、压缩单个张量的体积（量化与布局）、降低同时存活的规模（分块执行）。三者并非独立——融合改变生命周期，分块引入拼接开销，量化限制可用的融合模式。运行时的设计正处在三者交汇处。

---

2. 从计算图到执行计划

推理运行时的第一项工作是把模型文件中的计算图转换为执行计划：一个确定顺序的算子序列，每个算子绑定具体的实现内核、输入输出缓冲地址与分块参数。

这一转换可在设备端首次加载时完成，也可在主机端离线完成后序列化下发。离线编译能承担更重的搜索开销（融合方案枚举、分块调优），代价是失去对运行时内存预算的适应能力；在线编译可读取实际可用内存动态决定分块粒度，代价是引入一次加载延迟。

执行计划的产物包括算子序列（含内核选择）、缓冲布局表（每个张量在统一内存 arena 中的偏移与长度）与内存复用表（哪些张量共享同一段缓冲区）三项。第三项是运行时设计的核心，下文展开。

---

3. 图优化的三类变换

3.1 代数化简

常量折叠将输入全为常量的子图在编译期直接求值；算子消除则移除冗余节点，如连续两个转置、乘以 1 的缩放。这类变换的收益不在运行时开销，而在于缩短生命周期链——被消除的算子不再引入中间张量。

3.2 算子融合

垂直融合针对数据依赖链上的相邻算子，典型形态是归约型算子（矩阵乘、卷积）后接一串逐元素算子（加偏置、激活）。若不融合，链上每个算子都要完整读写一遍主存；融合后中间结果保留在寄存器或 L1 中，主存访问由 N 次降至 2 次。

水平融合针对共享同一输入或同一形状的多个独立算子，将它们合并为一次内核调用，收益来自减少内核启动与调度开销。

3.3 布局与数据类型插入

混合精度图中，类型转换算子常由编译期自动插入，数值上不改变结果却引入真实的访存开销。将其与相邻算子融合——例如把反量化吸收到矩阵乘的累加循环内——是量化推理的常规做法，其代数依据前文已述。

---

4. 融合边界的判定

融合并非越多越好，边界判定依据三条准则。

收益的量化。 一次融合的收益等于被消除的中间张量读写字节数：

saved_bytes = 2 × size(intermediate) × steps_eliminated

当中间张量规模远小于权重张量时（如大矩阵乘后的逐元素激活，其体积取决于批次与序列长度），其收益相对模型总访存量可能不足 1%，此时融合的意义在于减少内核启动次数而非带宽。

语义边界。 归约型算子跨越全局数据依赖，不能与消费其部分结果的算子随意融合。以 Softmax 为例，其分母依赖整行的最大值与指数和，属完整归约，后续算子须待归约完成后才能开始。将 Softmax 与逐元素缩放融合可行（后者不改变归约结构），再与归约型算子融合则会破坏语义。判定方法是检查候选组内是否存在"部分结果被消费"的依赖，即某算子输出的真子集被另一算子使用。

寄存器压力。 融合组越大，需同时驻留寄存器的中间值越多。组内并行度需求超过向量寄存器堆容量时，编译器会引入溢出（spill），把本应驻留寄存器的值写回栈，融合反而增加访存。该约束在向量长度可配的架构上尤为明显：VLEN 较小而融合组较宽的配置更易触发溢出。工程做法是对候选融合组设置上限，并以溢出计数作为编译期指标参与搜索。

---

5. 内存复用：生命周期分析

生命周期分析是复用分配的前提：对执行序列中的每个张量记录写入它的算子序号与最后读取它的算子序号，二者之间即为存活区间。

区间的重叠关系决定张量能否共享缓冲：两个张量的存活区间若不相交，其缓冲即可复用。这是一个区间图着色问题——张量视为顶点、区间重叠视为边，所需的最小缓冲段数等于任一位置上同时存活的张量数最大值。实践中最简单的策略是按写入位置排序后贪心分配：为每个张量选取不与任何已分配张量冲突的最小偏移，复杂度为 O(n²)，在端侧模型数百至数千算子的数量级下完全可接受。

复用率取决于图的拓扑。链式图的生命周期几乎不重叠，复用率接近理论最优；具有多分支输入的网络——残差连接、多尺度特征融合——会使多个张量长时间保持存活，从而抬高峰值，这是此类结构在内存受限设备上需特别处理的原因。

---

6. 内存池：为何运行时不做动态分配

确定全部缓冲偏移后，运行时只需在初始化阶段申请一块连续内存（arena），后续算子直接使用预计算的偏移。相比通用动态分配，其价值有三点：

• 无碎片。设备运行时长可达数周，反复分配释放会累积碎片，最终导致大块申请失败；
• 确定性时延。分配开销在初始化阶段一次付出，推理路径上不再有分配器调用；
• 可预算。arena 大小在编译期已知，可直接比对设备可用内存。

内存池通常按访问特性分层：权重张量在一次会话内保持不变，置于高带宽且容量受限的区域；激活缓冲按复用关系共享同一 arena；少量需跨层保留的张量单独分配。分层依据是访问频率与复用距离，而非数据大小。

---

7. 分块执行：内存不足时的降级策略

当可用内存无法容纳算子所需的工作集时，需将张量切分为若干块逐块计算。分块改变的是中间张量的驻留规模，而非计算总量。

以卷积为例，沿空间维度分块时，输出块的计算需要读取输入中略大于该块的区域，多出的边界称为 halo。对 k×k 卷积，halo 宽度为 (k-1)/2。分块后每块需额外读取边界输入，额外访存量与 halo 占比成正比：

overhead_ratio ≈ 1 - (BH × BW) / ((BH + k - 1) × (BW + k - 1))

块尺寸越大，halo 占比越低、访存效率越高，但占用内存更多。选取分块尺寸时可参考缓存容量与向量寄存器粒度的共同约束：块内需同时驻留的输入、权重与输出体积之和应能放入目标缓存层级，分块维度长度宜按向量寄存器位宽对齐以减少尾部掩码开销。运行时按当前可用内存动态调整分块粒度，是端侧运行时区别于静态编译器的常见设计。

---

8. 实战：生命周期分析与贪心分配器

以下实现接收一张算子依赖表，输出每个张量的缓冲偏移与 arena 总大小。代码可在 RV64 平台直接编译运行，不依赖特定硬件特性。

#include <stdio.h>

#define MAX_TENSOR 64

typedef struct {
    const char *name;
    int    first_use;   /* 写出它的算子序号 */
    int    last_use;    /* 最后读取它的算子序号 */
    size_t size;        /* 字节数 */
    int    resident;    /* 1 = 常驻（权重/输入），不参与复用 */
    size_t offset;      /* 分配结果 */
} tensor_t;

/* 有效存活区间：常驻张量视为全程占用，因而与所有激活张量避让 */
static void span_of(const tensor_t *t, int steps, int *f, int *l)
{
    if (t->resident) { *f = -1; *l = steps; }
    else             { *f = t->first_use; *l = t->last_use; }
}

static size_t align_up(size_t v, size_t a)
{
    return (v + a - 1) / a * a;
}

/* 贪心分配：常驻张量优先，其余按写入位置排序，逐次取最小可用偏移 */
static size_t assign(tensor_t *t, int n, int steps, size_t align,
                     size_t *resident_sz)
{
    int idx[MAX_TENSOR];
    size_t peak = 0;

    *resident_sz = 0;
    for (int i = 0; i < n; i++) idx[i] = i;

    /* 稳定插入排序：常驻张量排在最前，其余按 first_use 升序 */
    for (int i = 1; i < n; i++) {
        int k = idx[i], j = i - 1;
        int kb = t[k].resident ? -1 : t[k].first_use;
        while (j >= 0) {
            int ka = t[idx[j]].resident ? -1 : t[idx[j]].first_use;
            if (ka <= kb) break;
            idx[j + 1] = idx[j];
            j--;
        }
        idx[j + 1] = k;
    }

    for (int i = 0; i < n; i++) {
        tensor_t *cur = &t[idx[i]];
        int cf, cl, moved = 1;
        size_t off = 0;

        span_of(cur, steps, &cf, &cl);

        while (moved) {
            moved = 0;
            for (int j = 0; j < i; j++) {
                tensor_t *prev = &t[idx[j]];
                int pf, pl;
                span_of(prev, steps, &pf, &pl);
                if (!(cf <= pl && pf <= cl)) continue;   /* 生命周期不重叠 */
                if (off < prev->offset + prev->size &&
                    off + cur->size > prev->offset) {
                    off = align_up(prev->offset + prev->size, align);
                    moved = 1;
                }
            }
        }

        cur->offset = align_up(off, align);
        if (cur->offset + cur->size > peak)
            peak = cur->offset + cur->size;
        if (cur->resident)
            *resident_sz += cur->size;
    }
    return peak;
}

int main(void)
{
    /* 含残差分支的小型图：conv -> add -> relu -> gemm -> add -> softmax */
    tensor_t t[] = {
        { "input",  0, 0,   4096, 1, 0 },
        { "conv_w", 0, 1,  65536, 1, 0 },
        { "conv_y", 1, 2,   8192, 0, 0 },
        { "skip",   0, 2,   4096, 0, 0 },
        { "res",    2, 3,   8192, 0, 0 },
        { "gemm_w", 3, 4, 262144, 1, 0 },
        { "logits", 4, 5,   1024, 0, 0 },
        { "probs",  5, 5,   1024, 0, 0 },
    };
    int n = sizeof(t) / sizeof(t[0]), steps = 5;
    size_t resident_sz;
    size_t peak = assign(t, n, steps, 16, &resident_sz);

    size_t act_naive = 0;
    printf("%-8s %6s %10s %10s %8s\n",
           "tensor", "live", "offset", "size", "kind");
    for (int i = 0; i < n; i++) {
        printf("%-8s %2d-%-3d %10zu %10zu %8s\n", t[i].name,
               t[i].first_use, t[i].last_use, t[i].offset, t[i].size,
               t[i].resident ? "const" : "act");
        if (!t[i].resident) act_naive += t[i].size;
    }

    printf("\nresident (weight/input)   : %zu bytes\n", resident_sz);
    printf("activation peak (reused)  : %zu bytes\n", peak - resident_sz);
    printf("activation total (no-reuse): %zu bytes\n", act_naive);
    printf("arena total               : %zu bytes\n", peak);
    printf("saved by reuse            : %.1f%%\n",
           100.0 * (act_naive - (peak - resident_sz)) / act_naive);
    return 0;
}

编译与运行：

riscv64-unknown-linux-gnu-gcc -O2 -static -o alloc demo_alloc.c
qemu-riscv64 ./alloc

该图（8 个张量、5 个算子）的输出为：

tensor     live     offset       size     kind
input      0-0           0       4096    const
conv_w     0-1        4096      65536    const
conv_y     1-2      335872       8192      act
skip       0-2      331776       4096      act
res        2-3      344064       8192      act
gemm_w     3-4       69632     262144    const
logits     4-5      331776       1024      act
probs      5-5      332800       1024      act

resident (weight/input)    : 331776 bytes
activation peak (reused)   : 20480 bytes
activation total (no-reuse): 22528 bytes
arena total                : 352256 bytes
saved by reuse             : 9.1%

三点值得注意。

第一，常驻区与激活区严格分离。input、conv_w、gemm_w 三个张量占据 arena 前段并全程占用，其偏移不会与任何激活缓冲重合——这正是 span_of 把常驻张量的有效区间取为全程的原因。若省略这一步，分配器会把权重的地址误判为可复用。

第二，skip 无法被复用。它从算子 0 存活到算子 2，与 conv_y（1–2）、res（2–3）的区间均重叠，因此三者依次排布，把激活区峰值推到 20480 字节。而 logits 与 res 无关，得以回填到起始偏移。这组对照直观说明了复用率由拓扑而非张量大小决定。

第三，复用收益仅有 9.1%。小图上残差支路的长生命周期主导了峰值，压缩余地有限。若把残差相加的时机前移，或让支路只保留加法所需的部分，复用率会明显提升；这也解释了为什么编译器在调度融合顺序时会把缩短长生命周期张量作为首要目标。

接入运行时只需两步：把 assign 返回的峰值作为 arena 的申请长度，把每个张量的 offset 作为基址偏移量传给内核。推理路径上不再出现任何分配调用。

---

9. 调优判据与常见问题

调优判据。 若 arena 峰值显著高于权重总体积，瓶颈在中间张量，应优先做融合与复用；若峰值接近权重体积而端到端时延仍不达标，瓶颈在带宽，应转向量化与布局优化；若两者均正常但吞吐不稳，通常指向分块粒度与缓存层级的匹配问题。

常见问题。

• 融合后精度异常：多见反量化与归约算子融合时累加器精度不足，症状是特定输入的输出偏差而非全局错误，需检查中间累加类型是否满足动态范围；
• 内存占用随输入长度线性增长：通常是无上界的分块策略所致，应改为按可用内存反推分块尺寸；
• 复用率异常偏低：先检查生命周期标注是否有误，尤其是多消费者张量的 last_use 是否取到真正的最后一个消费者；
• 分块后性能反而下降：halo 占比过高或分块尺寸未对齐向量寄存器位宽，可通过增大块尺寸或调整分块维度验证。