Mini-Infer (22): 架构重构 — 链接器的魔法与“副作用”驱动的自动注册

1. 核心思想：从“拉 (Pull)”到“推 (Push)”

在此之前的架构中，我们采用的是 “显式触发” 模式：

  • 代码逻辑：Main 函数 -> 调用 Init() -> Init 调用 RegisterConv(), RegisterRelu()…
  • 缺点：耦合度高。每增加一个算子，都要修改 Init 函数。

现在，我们转向 “副作用驱动 (Side-Effect Driven)” 模式：

  • 代码逻辑：Conv.cpp 定义一个静态全局变量 -> 程序启动 -> 加载 Conv.o -> 触发静态变量构造函数 -> 构造函数执行 Register()。
  • 优点：高度解耦。Main 函数根本不需要知道 Conv 的存在，注册行为是加载库文件带来的“副作用”。

2. 隐形的大坑：静态链接的“死亡剔除” (Dead Code Stripping)

如果你直接在 .cpp 里写一个静态变量，然后把这个文件编译进静态库 (.a 或 .lib)，你会发现：注册代码根本没有执行！

这是 C++ 链接器（Linker）的默认优化行为：

  1. 链接器发现 main 函数没有引用 FusionPass 这个符号。
  2. 链接器判定 fusion_pass.o 是“无用的死代码”。
  3. 链接器丢弃了整个 fusion_pass.o 文件。
  4. 结果：里面的静态变量从未被初始化，注册逻辑从未执行。

3. 解决方案：CMake 接口目标与 --whole-archive

为了解决这个问题，我们需要强制链接器“吞下”所有的目标文件，无论它们是否被引用。

我们将修改 mini_infer/graph/CMakeLists.txt，引入一个特殊的接口目标 (Interface Target)：mini_infer_graph_all。

CMake 实现解析

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 1. 基础库：包含所有源文件
add_library(mini_infer_graph ${GRAPH_SOURCES})
# ... (常规配置) ...

# 2. 【核心】接口目标：强制全量链接
add_library(mini_infer_graph_all INTERFACE)

target_link_libraries(mini_infer_graph_all
    INTERFACE
        # 情况 A: 动态库或 MSVC 环境
        $<$<OR:$<BOOL:${BUILD_SHARED_LIBS}>,$<BOOL:${MSVC}>>:mini_infer_graph>
        
        # 情况 B: Linux/Mac 静态库 -> 使用 --whole-archive
        $<$<AND:$<NOT:$<BOOL:${BUILD_SHARED_LIBS}>>,$<NOT:$<BOOL:${MSVC}>>>:-Wl,--whole-archive>
        $<$<AND:$<NOT:$<BOOL:${BUILD_SHARED_LIBS}>>,$<NOT:$<BOOL:${MSVC}>>>:mini_infer_graph>
        $<$<AND:$<NOT:$<BOOL:${BUILD_SHARED_LIBS}>>,$<NOT:$<BOOL:${MSVC}>>>:-Wl,--no-whole-archive>
)

# 情况 C: MSVC 静态库 -> 使用 /WHOLEARCHIVE
if (NOT BUILD_SHARED_LIBS AND MSVC)
    target_link_options(mini_infer_graph_all INTERFACE "/WHOLEARCHIVE:mini_infer_graph")
endif()

设计精髓：

  • 传递性依赖：任何可执行文件（如 mini_infer_run）只需要链接 mini_infer_graph_all，CMake 就会自动将 --whole-archive 标志传递给链接器。
  • 平台兼容性：自动处理了 GCC/Clang (-Wl,--whole-archive) 和 MSVC (/WHOLEARCHIVE) 的差异。

4. 代码实现：Pass 的自动注册

有了构建系统的支持，我们就可以在 C++ 中放心使用静态构造技巧了。

以 FusionPass 为例，我们在 mini_infer/graph/fusion_pass.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
// mini_infer/graph/fusion_pass.cpp

// 使用匿名命名空间，避免符号污染全局
namespace {

// 工厂函数
std::shared_ptr<mini_infer::graph::OptimizationPass> create_FusionPass() {
    return std::make_shared<mini_infer::graph::FusionPass>();
}

// 注册辅助结构体
struct FusionPass_Register {
    FusionPass_Register() {
        // 获取单例注册表并注册
        // 优先级设为 100 (数值越大越早执行，或根据设计决定)
        mini_infer::graph::OptimizationPassRegistry::instance()
            .register_pass("FusionPass", create_FusionPass, 100);
    }
};

// 【关键】静态全局变量
// 它的构造函数会在 main() 执行前被运行时调用
static FusionPass_Register g_FusionPass_register;

}  // namespace

5. 全面重构：移除 KernelRegistryInitializer

既然我们解决了链接问题，是时候清理技术债务了。

我们不再需要 KernelRegistryInitializer 这个类，也不需要在 main 函数里调用 initialize()。

步骤 A：Kernel 的自动注册

在 gemm_cpu.cpp, im2col_cpu.cpp 等文件中，使用同样的 static struct 技巧：

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

namespace {
struct GEMM_CPU_Register {
    GEMM_CPU_Register() {
        // 这里的逻辑就是原来 KernelRegistryInitializer 里的代码
        GEMMRegistry_NN<float>::instance().register_kernel(
            KernelBackend::CPU,
            mini_infer::kernels::cpu::gemm_nn_impl<float>,
            []() { return true; }
        );
        // ... 注册其他类型 ...
    }
};
static GEMM_CPU_Register g_gemm_register;
}

步骤 B：构建系统的统一

为 kernels 模块也应用同样的 CMake 模式：创建 mini_infer_kernels 和 mini_infer_kernels_all。

用户现在的 CMakeLists.txt 只需要这样写：

1
2
3
4
5
6
7
8
add_executable(my_inference_app main.cpp)
target_link_libraries(my_inference_app 
    PRIVATE 
    mini_infer_core
    mini_infer_kernels_all  # 自动包含所有内核
    mini_infer_operators_all # 自动包含所有算子
    mini_infer_graph_all    # 自动包含所有 Pass
)

6. 总结

通过这次重构，我们完成了从代码级耦合到构建级解耦的转变。

  • 开发者体验：编写新 Kernel 或 Pass 时，只需新建文件，写完注册代码即可，无需修改任何现有文件。
  • 运行时零开销：注册过程发生在 main 函数之前，运行时不需要反复检查 is_initialized 标志。
  • 模块化：我们可以轻松地通过 CMake 决定链接哪些模块（例如，只链接 CPU Kernels 而不链接 CUDA Kernels），而不需要修改 C++ 源代码中的宏定义。

这就是成熟框架（如 TensorRT, PyTorch）背后的工程魔法。