Mini-Infer (24): 动态形状支持 — 运行时形状推理引擎

1. 为什么需要运行时推理？

在传统的静态图中，形状推导（Shape Inference）通常只在模型加载时做一次。但在动态场景下，形状推导必须变成一个运行时 (Runtime) 行为。

想象一个简单的网络：Input -> Conv -> ReLU -> Output。

如果 Input 变了，Conv 的输出形状 $H_{out} = (H_{in} + 2P - K)/S + 1$ 也必须跟着变。

我们不能每次都重新构建整个 Graph，那样太慢了。我们需要一个轻量级的引擎，快速遍历一遍图，只更新 Shape，不碰 Data。

2. 核心架构：拓扑顺序传播

ShapeInferenceEngine 的工作原理就像推多米诺骨牌：

设置起点：用户提供新的输入形状（例如 Input: [1, 3, 512, 512]）。
拓扑遍历：按照依赖顺序访问每个节点。
收集输入：对于节点 $N$ $N$ ，收集它所有输入张量的当前形状。
- 动态输入：来自上游节点的推导结果。
- 静态输入：来自权重（Weights/Bias），形状永远不变。
执行推导：调用算子的 op->infer_shape() 方法，计算输出形状。
更新状态：将输出形状存入缓存，供下游节点使用。

3. 代码实现剖析

A. 引擎入口与缓存机制

为了追求极致性能，我们引入了缓存机制。如果输入形状没变，直接返回成功，零开销。

// mini_infer/runtime/shape_inference_engine.cpp

core::Status ShapeInferenceEngine::infer_shapes(
    const std::unordered_map<std::string, core::Shape>& input_shapes
) {
    // 0. 拓扑排序 (Lazy Initialization)
    ensure_sorted();

    // 1. 缓存输入形状
    inferred_shapes_.clear();
    for (const auto& [name, shape] : input_shapes) {
        inferred_shapes_[name] = shape;
    }
    
    // ... 开始遍历 ...
}

B. 核心循环：处理混合输入

这是实现中最棘手的部分。一个算子（如 Conv2D）的输入来源通常是混合的：

Input 0 (Data): 来自上一个算子（动态，形状未知）。
Input 1 (Weight): 来自模型权重（静态，形状已知）。

我们需要把它们拼凑成一个完整的 input_shapes_vec 传给 op->infer_shape。

// 遍历所有节点
for (auto& node : sorted_nodes_) {
    std::vector<core::Shape> input_shapes_vec;
    
    // Step 1: 收集来自图连接的动态输入
    // 这些形状必须在之前的循环中已经计算出来 (inferred_shapes_)
    const auto& input_nodes = node->inputs();
    for (const auto& input_node : input_nodes) {
        auto it = inferred_shapes_.find(input_node->name());
        if (it != inferred_shapes_.end()) {
            input_shapes_vec.push_back(it->second);
        } else {
            return core::Status::ERROR_RUNTIME; // 依赖未就绪，逻辑错误
        }
    }
    
    // Step 2: 追加来自权重的静态输入
    // Conv 的 inputs() 只包含上一层节点，不包含权重 tensor
    // 权重 tensor 存储在 node->input_tensors() 中
    const auto& imported_tensors = node->input_tensors();
    for (size_t i = input_shapes_vec.size(); i < imported_tensors.size(); ++i) {
        if (imported_tensors[i]) {
            input_shapes_vec.push_back(imported_tensors[i]->shape());
        }
    }
    
    // Step 3: 调用算子推导逻辑
    std::vector<core::Shape> output_shapes;
    auto status = op->infer_shape(input_shapes_vec, output_shapes);
    
    // Step 4: 缓存结果供下游使用
    inferred_shapes_[node->name()] = output_shapes[0];
}

C. 变更检测与重分配

推导完成后，我们需要告诉内存管理器哪些 Tensor 变大了，需要重新 malloc。

std::vector<std::string> ShapeInferenceEngine::get_tensors_needing_reallocation() const {
    std::vector<std::string> tensors_to_reallocate;
    
    for (const auto& node : sorted_nodes_) {
        // 获取当前持有的 Tensor (旧内存)
        auto output_tensor = node->output_tensors()[0];
        // 获取刚刚推导出的 Shape (新需求)
        auto inferred_shape = get_inferred_shape(node->name());
        
        // 如果形状不一致，说明需要重分配
        if (output_tensor->shape() != *inferred_shape) {
            tensors_to_reallocate.push_back(node->name());
        }
    }
    return tensors_to_reallocate;
}

4. 与 MemoryPlanner 的协同

动态形状支持实际上是两种策略的结合：

首次运行 / 形状未变：使用 MemoryPlanner 计算出的静态复用方案，零碎片，低占用。
形状改变：触发 ShapeInferenceEngine，识别出尺寸变化的 Tensor，对其进行单独的 reallocate（此时可能会暂时打破最佳的内存复用，但保证了程序的正确运行）。

在 TensorRT 中，这对应于 IExecutionContext::setInputShape 接口。当调用它时，引擎内部就会执行上述的传播逻辑。