Ch21: 图优化——编译器思维在推理引擎中的应用

// ncnn/tools/ncnnoptimize.cpp（简化版关键优化列表）
void NetOptimize::optimize() {
    // 1. Conv + BN 融合
    fuse_convolution_batchnorm();
    // 把 BatchNorm 的参数合并到 Convolution 的权重中
    // weight_new = weight * (gamma / sqrt(var + eps))
    // bias_new = (bias - mean) * gamma / sqrt(var + eps) + beta

    // 2. Conv + ReLU 融合
    fuse_convolution_activation();
    // Convolution 自带 activation_type 参数
    // 设置 activation_type = 1 (ReLU) 就等于融合了后面的 ReLU

    // 3. 消除冗余操作
    eliminate_noop();           // 删除空操作
    eliminate_flatten_after_global_pooling();
    eliminate_reshape_after_global_pooling();
    eliminate_dropout();        // 推理时 Dropout 是恒等操作

    // 4. 常量折叠（简单版）
    // 如果某个算子的所有输入都是常量，提前计算结果
}

// ncnn 加载模型后直接执行，不做运行时优化
int Net::load_param(const char* protopath) {
    // 读取算子参数
    // 不做图变换、不做算子融合
    // 直接按原始顺序存储 Layer
}

// ncnn/src/layer/arm/convolution_arm.cpp（简化版）
int Convolution_arm::forward(...) const {
    // ... 卷积计算 ...

    // 内置激活融合
    if (activation_type == 1) {
        // ReLU：在卷积结果上直接执行，无需额外的内存访问
        for (int i = 0; i < output_size; i++) {
            output[i] = std::max(output[i], 0.f);
        }
    } else if (activation_type == 2) {
        // ReLU6
        for (int i = 0; i < output_size; i++) {
            output[i] = std::min(std::max(output[i], 0.f), 6.f);
        }
    }
    // ... 更多激活类型 ...
}

// 例：Permute (Transpose) 的 Geometry 实现
class GeometryPermute : public GeometryComputer {
    bool onCompute(const Op* op,
                   const std::vector<Tensor*>& inputs,
                   const std::vector<Tensor*>& outputs,
                   Context& context,
                   CommandBuffer& buffer) const override {
        auto input = inputs[0];
        auto output = outputs[0];
        auto perm = op->main_as_Permute()->dims();

        // 输出标记为 MEMORY_VIRTUAL
        auto des = TensorUtils::getDescribe(output);
        des->memoryType = InsideDescribe::MEMORY_VIRTUAL;

        // 创建 Region 描述映射关系
        InsideDescribe::Region region;
        region.origin = input;
        region.offset = 0;

        // 根据 perm 参数设置 stride 映射
        // 这样 Permute 就变成了"零拷贝的视图变换"
        for (int i = 0; i < perm->size(); i++) {
            region.size[i] = input->length(perm->Get(i));
            region.src_stride[i] = input->stride(perm->Get(i));
            region.dst_stride[i] = output->stride(i);
        }

        des->regions.push_back(region);
        return true;
    }
};

// onnxruntime/core/optimizer/graph_transformer_level.h
enum class TransformerLevel {
    Level1,  // 基础优化：常量折叠、死代码消除、基本融合
    Level2,  // 扩展优化：更激进的融合、布局优化
    Level3,  // 后端特定优化：EP 自定义的优化
};

Ort::SessionOptions options;
// 设置优化级别
options.SetGraphOptimizationLevel(
    GraphOptimizationLevel::ORT_ENABLE_EXTENDED  // Level 1 + 2
);

// 或者只启用基础优化
options.SetGraphOptimizationLevel(
    GraphOptimizationLevel::ORT_ENABLE_BASIC     // Level 1 only
);

// onnxruntime/core/optimizer/graph_transformer.h（简化版）
class GraphTransformer {
public:
    virtual ~GraphTransformer() = default;

    // 对整个图执行优化
    virtual common::Status Apply(Graph& graph,
                                  bool& modified) const = 0;

    // 优化器名称（用于日志和调试）
    const std::string& Name() const { return name_; }

protected:
    std::string name_;
};

// 基于规则的优化器（匹配模式 + 替换）
class RuleBasedGraphTransformer : public GraphTransformer {
    // 注册优化规则
    void Register(std::unique_ptr<RewriteRule>&& rule);

    // 遍历图中每个节点，尝试应用所有规则
    Status Apply(Graph& graph, bool& modified) const override;
};

// Level 1 优化
class ConstantFolding : public GraphTransformer {
    // 常量折叠：如果一个算子的所有输入都是常量，
    // 提前计算结果并替换为常量节点
    // 例：Shape(X) 在形状已知时替换为常量 [1,3,224,224]
};

class CommonSubexpressionElimination : public GraphTransformer {
    // 公共子表达式消除：如果两个算子有相同的输入和参数，
    // 只保留一个，另一个引用同一个输出
};

class ConvBNFusion : public RewriteRule {
    // Conv + BN 融合：数学等价合并
    bool SatisfyCondition(const Graph& graph,
                           const Node& node) const override {
        // 检查：node 是 Conv，唯一下游是 BN，BN 参数是常量
        return node.OpType() == "Conv" &&
               OnlyConsumer(graph, node) &&
               graph.GetNode(node.OutputDef(0)->Name())->OpType()
                   == "BatchNormalization";
    }

    Status Apply(Graph& graph, Node& conv_node,
                  bool& modified) const override {
        // 数学合并：把 BN 参数融合到 Conv 权重中
        // 删除 BN 节点
        modified = true;
        return Status::OK();
    }
};

// Level 2 优化
class MatMulAddFusion : public RewriteRule {
    // MatMul + Add -> Gemm 融合
};

class AttentionFusion : public GraphTransformer {
    // Q*K^T/sqrt(d)*V 模式识别 -> 融合为 MultiHeadAttention 算子
    // 这是 Transformer 模型优化的关键
};

// Level 3 优化（EP 特定）
class CUDAConvActivationFusion : public GraphTransformer {
    // CUDA EP 特有：Conv + ReLU 融合为 cuDNN 的融合 kernel
};

Ort::SessionOptions options;
// 保存优化后的模型
options.SetOptimizedModelFilePath("optimized_model.onnx");
// 创建 session 时会自动保存优化后的图

// TNN/source/tnn/optimizer/net_optimizer.h（简化版）
class NetOptimizer {
public:
    virtual std::string Strategy() = 0;

    virtual Status Optimize(NetStructure* structure,
                              NetResource* resource);
};

// 具体的优化器实现
class NetOptimizerFuseConvBN : public NetOptimizer {
    std::string Strategy() override {
        return "fuse_conv_batchnorm";
    }

    Status Optimize(NetStructure* structure,
                      NetResource* resource) override {
        // 遍历图，查找 Conv -> BN 模式
        for (int i = 0; i < structure->layers.size() - 1; i++) {
            auto& curr = structure->layers[i];
            auto& next = structure->layers[i + 1];

            if (curr->type == LAYER_CONVOLUTION &&
                next->type == LAYER_BATCH_NORM) {
                // 合并 BN 参数到 Conv 权重
                FuseConvBN(curr, next, resource);
                // 删除 BN 层
                structure->layers.erase(
                    structure->layers.begin() + i + 1);
            }
        }
        return TNN_OK;
    }
};

// TNN/source/tnn/optimizer/const_folder.h（简化版）
class ConstFolder {
public:
    Status Fold(NetStructure* structure,
                 NetResource* resource) {
        bool changed = true;
        while (changed) {
            changed = false;
            for (auto& layer : structure->layers) {
                if (AllInputsAreConst(layer, resource)) {
                    // 执行这个算子，得到常量结果
                    auto result = Execute(layer, resource);
                    // 把结果存为新的常量 Blob
                    resource->constant_map[layer->outputs[0]] = result;
                    // 从图中删除这个算子
                    RemoveLayer(structure, layer);
                    changed = true;
                    break;
                }
            }
        }
    }
};

// paddle-lite/lite/core/optimizer/mir/ssa_graph.h（简化版）
class SSAGraph {
public:
    // 节点列表（算子节点 + 数据节点）
    std::list<std::unique_ptr<Node>> nodes_;

    // 添加算子节点
    Node* AddInstruction(const cpp::OpDesc& desc);
    // 添加数据节点（变量）
    Node* AddArgument(const std::string& name);
};

void Optimizer::Run(Program* program,
                     const std::vector<Place>& valid_places) {
    // ======= 第一类：类型推导与转换 =======
    RunPass("type_target_cast_pass");    // 设备间转换（CPU<->GPU）
    RunPass("type_precision_cast_pass"); // 精度转换（fp32<->fp16<->int8）
    RunPass("type_layout_cast_pass");    // 布局转换（NCHW<->NHWC）

    // ======= 第二类：算子融合 =======
    RunPass("lite_conv_bn_fuse_pass");
    RunPass("lite_conv_activation_fuse_pass");
    RunPass("lite_fc_fuse_pass");
    RunPass("lite_shuffle_channel_fuse_pass");
    RunPass("lite_transpose_softmax_transpose_fuse_pass");
    RunPass("lite_interpolate_fuse_pass");
    // ... 20+ 种融合模式

    // ======= 第三类：Kernel 选择 =======
    RunPass("static_kernel_pick_pass");
    // 根据三维 Place (Target, Precision, Layout)
    // 为每个 Op 选择最佳 kernel

    // ======= 第四类：子图优化 =======
    RunPass("nnadapter_subgraph_pass");
    // 把支持 NNAdapter 的连续算子合并为子图

    // ======= 第五类：内存优化 =======
    RunPass("memory_optimize_pass");
    // 分析张量生命周期，复用内存

    // ======= 第六类：量化相关 =======
    RunPass("lite_quant_dequant_fuse_pass");
    RunPass("weight_quantization_preprocess_pass");
}

// paddle-lite/lite/core/optimizer/mir/fusion/conv_bn_fuse_pass.cc
class ConvBNFusePass : public ProgramPass {
public:
    void Apply(const std::unique_ptr<SSAGraph>& graph) override {
        // 模式匹配：找 Conv -> BN 的连续模式
        for (auto& node : graph->nodes_) {
            if (!node->IsStmt()) continue;
            auto& stmt = node->AsStmt();

            if (stmt.op_type() == "conv2d") {
                // 检查输出是否只被一个 BN 消费
                auto out = stmt.op_info()->Output("Output").front();
                auto consumers = FindConsumers(graph, out);

                if (consumers.size() == 1 &&
                    consumers[0]->AsStmt().op_type() == "batch_norm") {
                    // 执行融合
                    FuseConvBN(&stmt, &consumers[0]->AsStmt());
                    // 删除 BN 节点
                    GraphSafeRemoveNode(graph, consumers[0]);
                }
            }
        }
    }
};

// mace/core/flow/common_fp32_flow.cc（简化版）
MaceStatus CommonFp32Flow::Run(TensorMap* input_tensors,
                                TensorMap* output_tensors) {
    // 直接按顺序执行每个 Operation
    // 没有运行时图优化、没有动态融合
    for (auto& op : net_->operators()) {
        MACE_RETURN_IF_ERROR(op->Run());
    }
    return MaceStatus::MACE_SUCCESS;
}