Mini-Infer (35): 插件架构实战 — 从旧架构到新架构的迁移

1. 迁移策略概述

从旧的 Operator + Kernel 架构迁移到新的 Plugin 架构，我们采用以下策略：

A. 保留底层原语

底层计算原语（im2col、gemm、bias、transpose）保持不变，它们是设备无关的数学操作：

1
2
3
4
5
6
7
保留：
├── kernels/cpu/gemm.cpp
├── kernels/cpu/im2col.cpp
├── kernels/cpu/bias.cpp
├── kernels/cuda/gemm.cu
├── kernels/cuda/im2col.cu
└── kernels/cuda/bias.cu

B. 删除旧 Kernel 实现

旧的算子级 Kernel（如 Conv2DKernel、ReLUKernel）被删除，其逻辑移入 Plugin：

1
2
3
4
5
删除：
├── kernels/cpu/conv2d_kernel.cpp  → 移入 Conv2DCPUPlugin
├── kernels/cpu/relu_kernel.cpp   → 移入 ReLUCPUPlugin
├── kernels/cuda/conv2d_kernel.cu → 移入 Conv2DCUDAPlugin
└── kernels/cuda/relu_kernel.cu   → 移入 ReLUCUDAPlugin

C. 删除旧 Operator 实现

旧的 Operator 类被 GenericOperator + Plugin 替代：

1
2
3
4
5
6
7
删除：
├── operators/conv2d.cpp
├── operators/relu.cpp
└── operators/pooling.cpp

保留：
└── operators/generic_operator.cpp  // 通用算子容器

2. PluginOperatorAdapter 适配器

为了让 Plugin 能够与现有的 Node 系统协作，我们实现了一个适配器：

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
// mini_infer/operators/plugin_operator_adapter.h

class PluginOperatorAdapter : public Operator {
public:
    explicit PluginOperatorAdapter(std::unique_ptr<IPlugin> plugin)
        : plugin_(std::move(plugin)) {}

    core::OpType type() const override {
        return plugin_->get_op_type();
    }

    core::Status forward(
        const std::vector<std::shared_ptr<core::Tensor>>& inputs,
        std::vector<std::shared_ptr<core::Tensor>>& outputs,
        backends::DeviceContext* context) override {

        PluginContext plugin_ctx;
        plugin_ctx.device_context = context;
        return plugin_->enqueue(inputs, outputs, plugin_ctx);
    }

    core::Status infer_shape(
        const std::vector<core::Shape>& input_shapes,
        std::vector<core::Shape>& output_shapes) const override {
        return plugin_->infer_output_shapes(input_shapes, output_shapes);
    }

private:
    std::unique_ptr<IPlugin> plugin_;
};

作用：将 IPlugin 接口适配为 Operator 接口，保持与 Node 系统的兼容性。

3. 运行时适配

A. ExecutionContext::execute_node 的 Plugin 调用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// mini_infer/runtime/execution_context.cpp

core::Status ExecutionContext::execute_node(const std::shared_ptr<graph::Node>& node) {
    // ... 收集输入 ...

    // 使用缓存的 Plugin 执行
    auto* cached_plugin = node->get_operator()->cached_plugin();
    if (!cached_plugin) {
        MI_LOG_ERROR("[ExecutionContext] No plugin available for node: " + node->name());
        return core::Status::ERROR_NOT_IMPLEMENTED;
    }

    operators::PluginContext plugin_ctx;
    plugin_ctx.device_context = context.get();
    return cached_plugin->enqueue(merged_inputs, output_tensors, plugin_ctx);
}

B. cached_plugin 机制

在 InferencePlan::infer_shapes 中，我们为每个节点创建并缓存 Plugin：

1
2
3
4
5
6
7
8
9
10
11
12
// 创建 Plugin
auto plugin = PluginRegistry::instance().create_plugin(
    node->type(), config_.device_type);

// 设置参数
auto* generic_op = dynamic_cast<GenericOperator*>(node->get_operator().get());
if (generic_op && generic_op->plugin_param()) {
    plugin->set_param(generic_op->plugin_param());
}

// 缓存到 Operator
node->get_operator()->set_cached_plugin(std::move(plugin));

优势：

  1. 避免重复创建：Plugin 只在构建时创建一次。
  2. 参数绑定：参数在构建时设置，执行时直接使用。
  3. 类型安全：通过 cached_plugin() 获取正确类型的 Plugin。

4. CUDA 推理端到端示例

A. lenet5_cuda_inference.cpp 解析

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
// examples/lenet5_cuda_inference.cpp

#include "mini_infer/mini_infer.h"

int main() {
    // 1. 解析 ONNX 模型
    importers::OnnxParser parser;
    auto graph = parser.parse_from_file("lenet5.onnx");

    // 2. 配置 CUDA 引擎
    runtime::EngineConfig config;
    config.device_type = core::DeviceType::CUDA;
    config.device_id = 0;
    config.enable_graph_optimization = true;
    config.enable_memory_planning = true;

    // 3. 构建 InferencePlan
    auto plan = std::make_shared<runtime::InferencePlan>(config);
    plan->build(graph);

    // 4. 创建 ExecutionContext
    auto ctx = plan->create_execution_context();

    // 5. 准备输入（CPU 上）
    auto input = core::Tensor::create({1, 1, 28, 28}, core::DataType::FLOAT32);
    // ... 填充输入数据 ...

    // 6. 设置输入并执行
    ctx->set_inputs({{"input", input}});
    plan->execute(ctx.get());

    // 7. 获取输出（自动从 GPU 拷贝回 CPU）
    auto output = ctx->outputs()[0];
    const float* output_data = static_cast<const float*>(output->data());

    // 8. 打印结果
    std::cout << "Prediction: " << argmax(output_data, 10) << std::endl;

    return 0;
}

B. CPU vs GPU 性能对比

1
2
3
4
5
6
7
8
LeNet-5 推理性能 (batch_size=1, 1000 次迭代):

| 设备 | 首次推理 | 平均推理 | 吞吐量 |
|------|----------|----------|--------|
| CPU  | 0.5 ms   | 0.3 ms   | 3333 img/s |
| GPU  | 2.0 ms   | 0.1 ms   | 10000 img/s |

注：GPU 首次推理包含 CUDA 初始化和权重预加载开销。

5. GPU 缓存安全性改进

A. shared_ptr 键的必要性

在早期实现中，我们使用裸指针作为缓存的 Key：

1
2
// 旧实现（有问题）
std::unordered_map<const core::Tensor*, std::shared_ptr<core::Tensor>> cache;

问题：如果原始 Tensor 被释放，Key 变成悬空指针。

新实现：

1
2
3
4
// 新实现（安全）
std::unordered_map<std::shared_ptr<const core::Tensor>,
                   std::shared_ptr<core::Tensor>,
                   TensorPtrHash, TensorPtrEqual> gpu_weight_cache_;

使用 shared_ptr 作为 Key，确保只要缓存存在，原始 Tensor 就不会被释放。

B. CUDA 错误处理增强

1
2
3
4
5
6
7
8
9
10
// 使用 cudaMemcpyAsync 替代 cudaDeviceSynchronize
cudaError_t status = cudaMemcpyAsync(
    outputs[0]->data(), tensor->data(), size_bytes,
    cudaMemcpyHostToDevice, stream);

if (status != cudaSuccess) {
    MI_LOG_ERROR("[InferencePlan] Failed to copy input to GPU: " +
                 std::string(cudaGetErrorString(status)));
    return core::Status::ERROR_RUNTIME;
}

6. 测试与示例更新

A. test_operators.cpp 适配

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
// tests/test_operators.cpp

TEST(PluginTest, ReLUCPU) {
    // 创建 Plugin
    auto plugin = PluginRegistry::instance().create_plugin(
        core::OpType::kRELU, core::DeviceType::CPU);
    ASSERT_NE(plugin, nullptr);

    // 准备输入
    auto input = core::Tensor::create({2, 3}, core::DataType::FLOAT32);
    float* data = static_cast<float*>(input->data());
    data[0] = -1.0f; data[1] = 2.0f; data[2] = -3.0f;
    data[3] = 4.0f; data[4] = -5.0f; data[5] = 6.0f;

    // 准备输出
    auto output = core::Tensor::create({2, 3}, core::DataType::FLOAT32);

    // 执行
    std::vector<std::shared_ptr<core::Tensor>> inputs = {input};
    std::vector<std::shared_ptr<core::Tensor>> outputs = {output};
    PluginContext ctx;
    auto status = plugin->enqueue(inputs, outputs, ctx);
    ASSERT_EQ(status, core::Status::SUCCESS);

    // 验证
    const float* out_data = static_cast<const float*>(output->data());
    EXPECT_FLOAT_EQ(out_data[0], 0.0f);
    EXPECT_FLOAT_EQ(out_data[1], 2.0f);
    EXPECT_FLOAT_EQ(out_data[2], 0.0f);
    EXPECT_FLOAT_EQ(out_data[3], 4.0f);
    EXPECT_FLOAT_EQ(out_data[4], 0.0f);
    EXPECT_FLOAT_EQ(out_data[5], 6.0f);
}

B. 示例代码更新

所有示例都更新为使用新的 Plugin API：

1
2
3
4
5
6
examples/
├── lenet5_inference.cpp           # CPU 推理
├── lenet5_cuda_inference.cpp      # GPU 推理
├── lenet5_optimized_inference.cpp # 优化推理
├── lenet5_dynamic_multi_batch.cpp # 动态形状
└── memory_tracking_example.cpp    # 内存追踪

7. 架构演进总结

A. 从 Operator + Kernel 到 Plugin

1
2
3
4
5
6
7
8
9
10
11
旧架构：
┌─────────────┐     ┌─────────────┐     ┌─────────────┐
│  Operator   │────►│KernelRegistry│────►│   Kernel    │
│ (元数据)    │     │  (查找)     │     │  (计算)     │
└─────────────┘     └─────────────┘     └─────────────┘

新架构：
┌─────────────┐     ┌─────────────┐
│PluginRegistry│────►│   Plugin    │
│  (查找)     │     │(元数据+计算) │
└─────────────┘     └─────────────┘

B. 代码量减少与维护性提升

指标 旧架构 新架构 变化
注册表数量 2 1 -50%
添加新算子步骤 4 2 -50%
代码行数 ~5000 ~3500 -30%
头文件数量 15 10 -33%

8. 总结与展望

本篇我们完成了 Mini-Infer 插件架构的最后一部分——迁移实战：

  • 迁移策略：保留底层原语，删除旧 Kernel 和 Operator。
  • PluginOperatorAdapter：适配 Plugin 到 Node 系统。
  • cached_plugin 机制：构建时创建，执行时复用。
  • CUDA 端到端示例：完整的 GPU 推理流程。
  • 安全性改进：shared_ptr 键、CUDA 错误处理。

至此，Mini-Infer 的插件架构重构完成。回顾整个系列，我们实现了：

  1. 运行时架构重构：InferencePlan + ExecutionContext 分离。
  2. Core 数据结构优化：Storage 与 Tensor 分离。
  3. CUDA 后端支持：CUDAAllocator、CUDADeviceContext、权重预加载。
  4. 插件架构：IPlugin 接口、CRTP 基类、PluginRegistry。

Mini-Infer 现在是一个功能完整的轻量级推理框架，支持：

  • ONNX 模型导入
  • 图优化（算子融合）
  • 静态内存规划
  • 动态形状支持
  • CPU 和 CUDA 后端
  • TensorRT 风格的 Plugin 架构

未来展望：

  • 更多算子支持（Attention、LayerNorm 等）
  • INT8 量化推理
  • 多 GPU 并行
  • 模型序列化/反序列化
  • 性能分析工具