Mini-Infer (7): 高性能“内核注册表” (A TensorRT-Style Kernel Registry)

1. 架构目标：从“静态分发”到“动态注册”

我们的新目标是：

解耦：GEMMKernel（调度器）不应该“知道”任何具体的实现（如 avx2_gemm_impl）。
可扩展：添加一个新的 AVX512 内核，应该不需要修改任何现有的 GEMMKernel 代码。
高性能：系统必须能自动检测硬件能力，并优先选择最快的可用内核（例如，cuBLAS > AVX2 > CPU）。

为了实现这一点，我们将构建一个“内核电话簿”（Registry），每个内核实现（AVX2、CUDA…）都会在启动时自动将其“电话号码”（函数指针）和“能力”注册到这个“电话簿”中。

2. 核心设计：`KernelRegistryBase` (`kernel_registry.h`)

这是我们的“电话簿”模板。它是一个通用的 C++ 模板类，可以为任何类型的内核（GEMM, im2col…）管理一个实现列表。

// mini_infer/kernels/kernel_registry.h
#pragma once
#include "mini_infer/kernels/kernel_types.h"
#include <functional>
#include <vector>
#include <memory>
#include <algorithm>

namespace mini_infer {
namespace kernels {

// “能力检查器”：一个 std::function，用于检查硬件是否支持
using KernelCapabilityChecker = std::function<bool()>;

// 1. “电话簿条目” (KernelEntry)
// 每一个内核实现都对应一个条目
template<typename FuncType>
struct KernelEntry {
    KernelBackend backend;            // 它是 CPU 还是 CUDA?
    FuncType func;                    // 【核心】指向内核实现的函数指针
    KernelCapabilityChecker is_supported; // 检查硬件的函数 (e.g., []{ return supports_avx2(); })
    int priority;                     // 【关键】优先级 (e.g., AVX2=100, CPU=10)
    
    // ... 构造函数 ...
};

// 2. “电话簿” (KernelRegistryBase)
template<typename FuncType>
class KernelRegistryBase {
public:
    using Entry = KernelEntry<FuncType>;
    
    // 注册一个内核实现
    void register_kernel(KernelBackend backend, 
                         FuncType func,
                         KernelCapabilityChecker checker,
                         int priority = 0) {
        entries_.emplace_back(backend, func, checker, priority);
        
        // 【关键】按优先级降序排序
        // 确保最高优先级的内核总是在列表的最前面
        std::sort(entries_.begin(), entries_.end(),
                  [](const Entry& a, const Entry& b) {
                      return a.priority > b.priority;
                  });
    }
    
    // 【核心 API】获取“最佳”内核
    // 它遍历列表，返回第一个“硬件支持”的内核
    FuncType get_best_kernel() const {
        for (const auto& entry : entries_) {
            if (entry.is_supported()) {
                // 因为列表已排序，第一个支持的就是“最佳”的
                return entry.func;
            }
        }
        return nullptr; // 没有可用的内核
    }
    
    // ... (其他辅助函数如 get_kernel(backend), is_backend_available(...)) ...
    
protected:
    std::vector<Entry> entries_;
};

这个设计的核心是 get_best_kernel()。Linear 算子不再需要知道 AVX2 的存在。它只需要在初始化时调用一次 GEMMRegistry::instance().get_best_kernel()，就能自动获得一个（可能是 AVX2 或 cuBLAS）指向最快 GEMM 函数的指针。

3. C++ 魔法：`AutoRegister` 与“静态初始化”

我们如何“填充”这个“电话簿”？我们不希望在 main() 函数中写一长串 register_kernel(...)。

我们使用与 OperatorFactory（Blog 3）完全相同的技巧：自动注册。

// mini_infer/kernels/kernel_registry.h (接上文)

// 自动注册辅助模板
template<typename Registry, typename FuncType>
struct AutoRegister {
    AutoRegister(KernelBackend backend, 
                 FuncType func,
                 KernelCapabilityChecker checker,
                 int priority = 0) {
        // 在构造时，自动调用 Registry 的“单例”
        // 并注册自己
        Registry::instance().register_kernel(backend, func, checker, priority);
    }
};

如何使用它？ 假设我们在 gemm_avx2.cpp 中编写了一个 AVX2 优化的 GEMM。我们只需在该文件的全局命名空间中添加这一行：

// in: mini_infer/kernels/gemm_avx2.cpp
namespace {
    // 1. 我们的 AVX2 内核
    void gemm_nt_avx2_impl(const float* A, ...) { /* ... AVX2 高速代码 ... */ }
    
    // 2. 硬件能力检查
    bool supports_avx2() { /* ... 检查 CPUID ... */ return true; }

    // 3. 【C++ 魔法】
    // 定义一个全局静态变量，其类型为 AutoRegister
    static auto _gemm_avx2_register = AutoRegister<GEMMRegistry>(
        KernelBackend::CPU_AVX2,     // 它的类型
        gemm_nt_avx2_impl,           // 它的函数指针
        supports_avx2,               // 它的硬件检查器
        100                          // 它的优先级 (高于普通 CPU)
    );
}

在程序启动时（main 执行之前），_gemm_avx2_register 变量的构造函数会自动被调用，从而将 gemm_nt_avx2_impl 自动注册到 GEMMRegistry 中！

4. 解决“终极问题”：`KernelRegistryInitializer` 与链接器

AutoRegister 有一个致命的弱点：静态库（Static Libraries）。

如果 gemm_avx2.cpp 被编译进一个静态库（如 libmini_infer_kernels.a），而我们的主程序 Engine 没有显式调用 gemm_avx2.cpp 中的任何函数，那么链接器（Linker）会认为 gemm_avx2.cpp 整个文件都是“无用的”。

链接器会“智能地”丢弃这个文件，导致 _gemm_avx2_register 这个全局变量根本不存在于最终的可执行文件中。自动注册**“无声地”失败**了。

KernelRegistryInitializer 就是我们的解决方案。

// mini_infer/kernels/kernel_registry_init.h
class KernelRegistryInitializer {
public:
    static void initialize(); // 一个必须被调用的"入口"
private:
    static bool initialized_;
};

// mini_infer/kernels/kernel_registry_init.cpp
// 1. 在每个内核文件中（如 gemm_cpu.cpp）定义一个“注册函数”
//    (在 gemm_cpu.cpp 中)
//    namespace cpu { void register_gemm_kernels() { /* 这个函数是空的 */ } }
//    (在 gemm_avx2.cpp 中)
//    namespace cpu_avx2 { void register_gemm_kernels() { /* 空 */ } }

// 2. 在 initialize() 中“显式”调用这些函数
void KernelRegistryInitializer::initialize() {
    if (initialized_) return;
    
    // 强制注册所有 CPU 内核
    cpu::register_gemm_kernels();
    cpu::register_im2col_kernels();
    
    // 强制注册 AVX2 内核
    // cpu_avx2::register_gemm_kernels(); 
    
    initialized_ = true;
}