Mini-Infer 架构深潜 (3): Operator 抽象与“自注册”工厂

Mini-Infer 架构深潜 (3): Operator 抽象与“自注册”工厂

引言：缺失的“计算”拼图

在前两篇文章中，我们构建了 Mini-Infer 的数据 (Tensor) 和执行上下文 (Backend)。我们已经有了坚实的地基，但大厦至今仍是“空”的——它没有任何“功能”。

我们如何定义一个“卷积”操作？如何定义一个 “ReLU” 激活？最重要的是，我们的 Net (计算图) 如何在不“写死”依赖的情况下，动态地加载和执行这些操作？

本篇，我们将构建 Mini-Infer 的计算核心：Operator 抽象层。我们将使用 C++ 中一个极其精妙的模式——工厂 (Factory) + 自动注册 (Self-Registration)——来实现一个真正可插拔、可扩展的算子系统。

---

1. Operator 接口：“计算”的合约 (operator.h)

首先，我们必须定义一个标准“合约”，所有计算单元（卷积、激活、池化等）都必须遵守这个合约。这就是抽象基类 Operator 的职责。

// mini_infer/operators/operator.h
#pragma once

#include "mini_infer/core/tensor.h"
#include "mini_infer/core/types.h"
#include <vector>
#include <string>
#include <memory>
#include <unordered_map>

namespace mini_infer {
namespace operators {

/**
 * @brief Operator parameter base class
 */
struct OpParam {
    virtual ~OpParam() = default;
};

/**
 * @brief Operator interface abstract class
 */
class Operator {
public:
    explicit Operator(const std::string& name) : name_(name) {}
    virtual ~Operator() = default;
    
    /**
     * @brief Forward inference
     */
    virtual core::Status forward(
        const std::vector<std::shared_ptr<core::Tensor>>& inputs,
        std::vector<std::shared_ptr<core::Tensor>>& outputs
    ) = 0;
    
    /**
     * @brief Infer the output shape
     */
    virtual core::Status infer_shape(
        const std::vector<core::Shape>& input_shapes,
        std::vector<core::Shape>& output_shapes
    ) = 0;
    
    const std::string& name() const { return name_; }
    
    virtual void set_param(std::shared_ptr<OpParam> param) { param_ = param; }
    
protected:
    std::string name_; //< The name of the operator
    std::shared_ptr<OpParam> param_; //< The parameter of the operator
};
// ... Factory code follows ...

这个接口的设计有三个关键点：

1. OpParam 基类： 这是一个轻量级的 struct，用作所有算子参数（如 stride, kernel_size）的基类。这允许 Convolution 定义一个 ConvParam 结构体，并将其作为 std::shared_ptr<OpParam> 泛型地存入 Operator 基类中。
2. forward(...) (执行): 这是算子的核心执行函数。它接收 Tensor 列表，执行计算，并填充输出 Tensor 列表。
3. infer_shape(...) (塑形): 这是整个框架的性能核心之一。infer_shape 只操作元数据 (Shape)，它不执行任何实际计算。
   • 为什么分离？ 这允许 Net 在运行 forward 之前调用所有算子的 infer_shape。Net 可以预先知道图中每一个中间 Tensor 的确切形状，从而一次性分配好所有内存。这避免了在 forward 过程中动态分配内存的巨大开销。

---

2. OperatorFactory：解耦的“总管” (operator.h & .cpp)

我们如何根据模型文件中的一个字符串（如 “Convolution”）来创建对应的 C++ Convolution 对象？

我们决不能在 Net 类中写 if (op_type == "Convolution") { ... }。这种“硬编码”会产生一个“巨无霸”类，每添加一个新算子，都必须修改 Net 类。

解决方案是工厂模式。OperatorFactory 是一个全局唯一的总管，它是唯一一个“知道”如何创建所有算子的地方。

// mini_infer/operators/operator.h (接上文)

class OperatorFactory {
public:
    /**
     * 定义一个"函数指针"类型，名为 CreateFunc
     * 该指针指向"不带参数"且"返回 std::shared_ptr<Operator>"的函数
     */
    using CreateFunc = std::shared_ptr<Operator>(*)();
    
    // 注册一个算子：将 字符串(op_type) 映射到 创建函数(func)
    static void register_operator(const std::string& op_type, CreateFunc func);
    
    // 创建一个算子：用 字符串(op_type) 查找 创建函数 并调用它
    static std::shared_ptr<Operator> create_operator(const std::string& op_type);
    
private:
    // 获取全局唯一的"注册表"
    static std::unordered_map<std::string, CreateFunc>& get_registry();
};

OperatorFactory 的实现隐藏在 .cpp 文件中，这非常关键：

// mini_infer/operators/operator.cpp
#include "mini_infer/operators/operator.h"
#include <unordered_map>

namespace mini_infer {
namespace operators {

// 【关键】"注册表"本体被定义为 get_registry() 内部的 static 变量
// 这确保了：
// 1. 它是全局唯一的 (Meyers' Singleton)
// 2. 它被"隐藏"在 .cpp 中，外部无法访问
// 3. 它是线程安全的 (C++11 保证)
std::unordered_map<std::string, OperatorFactory::CreateFunc>& 
OperatorFactory::get_registry() {
    static std::unordered_map<std::string, CreateFunc> registry;
    return registry;
}

// 注册：向"注册表"中添加一个条目
void OperatorFactory::register_operator(const std::string& op_type, CreateFunc func) {
    get_registry()[op_type] = func;
}

// 创建：从"注册表"中查找并执行
std::shared_ptr<Operator> OperatorFactory::create_operator(const std::string& op_type) {
    auto& registry = get_registry();
    auto it = registry.find(op_type);
    if (it != registry.end()) {
        // 找到了！it->second 就是我们注册的 CreateFunc
        return it->second(); // 调用函数指针，创建实例
    }
    return nullptr; // 未找到
}

} // namespace operators
} // namespace mini_infer

现在，Net 类只需要和 OperatorFactory::create_operator("Convolution") 这一个函数打交道，它完全不需要知道 Convolution 类的存在。解耦初步达成！

---

3. REGISTER_OPERATOR：C++ 宏的“自动注册”魔法

我们还剩最后一个问题：OperatorFactory 如何“知道”所有算子？谁去调用 register_operator？

我们不希望在 main 函数中手动调用 register_operator("Convolution", ...)，这同样是“硬编码”。

我们希望算子自己“主动”去工厂注册。这就是“自动注册”模式，通过 C++ 宏和“静态初始化”的特性巧妙实现。

// mini_infer/operators/operator.h (接上文)

// 【C++ 宏魔法】
#define REGISTER_OPERATOR(op_type, op_class) \
    namespace { \
        /* 1. 定义一个局部、静态的"创建函数" */ \
        std::shared_ptr<Operator> create_##op_class() { \
            return std::make_shared<op_class>(); \
        } \
        /* 2. 定义一个"注册器"结构体 */ \
        struct op_class##_register { \
            op_class##_register() { \
                /* 3. 在构造函数中，调用工厂的注册方法 */ \
                OperatorFactory::register_operator(#op_type, create_##op_class); \
            } \
        }; \
        /* 4. 【关键】定义一个该结构体的全局静态变量 */ \
        static op_class##_register g_##op_class##_register; \
    }

这是整个框架最精妙的部分，我们来拆解它的执行过程：

假设你创建了一个 relu_operator.cpp 文件，并在文件末尾添加了这一行： REGISTER_OPERATOR(ReLU, ReLUOperator)

在 C++ 程序启动时（main 函数执行之前）：

1. 编译器在 relu_operator.cpp 中看到了第 4 步的 static ReLUOperator_register g_ReLUOperator_register;。
2. C++ 运行时需要初始化这个全局静态变量，因此它会调用该变量的构造函数（第 3 步的 ReLUOperator_register()）。
3. 这个构造函数执行，它调用 OperatorFactory::register_operator("ReLU", create_ReLUOperator)。
4. OperatorFactory 的全局 registry (注册表) 中，被添加了一个条目：{"ReLU", [指向 create_ReLUOperator 的函数指针]}。
5. … 这个过程在所有包含 REGISTER_OPERATOR 宏的 .cpp 文件中都会发生。

最终结果： 当你的 main 函数开始执行时，OperatorFactory 的 registry (注册表) 已经自动被所有算子填满了。Net 类可以立即通过字符串创建任何已注册的算子，而无需知道这些算子类的任何细节。

总结与展望

Mini-Infer 的核心架构已经成型！我们通过三层抽象，实现了彻底的解耦：

• Tensor (数据层): 管理内存和形状。
• Backend (执行层): 抽象硬件操作（memcpy, allocate）。
• Operator (计算层): 抽象计算逻辑（forward, infer_shape）。

OperatorFactory 和 REGISTER_OPERATOR 宏则像“神经中枢”，将这些松散的组件动态地“编织”在一起。