cann组织链接：https://atomgit.com/cann
graph-autofusion仓库链接：https://atomgit.com/cann/graph-autofusion

前言

在深度学习模型部署中，算子融合（Operator Fusion）是提升执行效率、降低内存占用的关键优化手段。CANN（Compute Architecture for Neural Networks）开源项目中的 graph-autofusion 仓库（https://atomgit.com/cann/graph-autofusion）提供了一套轻量级、解耦式的自动融合框架，其核心在于一套基于规则与成本模型的融合决策算法。该算法不仅支持传统 element-wise 融合，还能处理复杂的控制流、动态 shape 与自定义 Kernel 场景。

1. 自动融合的整体架构与工作流程

graph-autofusion 采用“图分析 → 规则匹配 → 成本评估 → Kernel 生成”的四阶段流水线：

该流程由 graph_fusion_manager.cpp 驱动，支持多次迭代融合。

关键组件

• FusionAnalyzer：遍历图，识别可融合子图；
• RuleMatcher：基于 YAML 规则匹配算子模式；
• CostModel：估算融合前后性能收益；
• SuperKernel：JIT 编译融合 Kernel。

---

2. 融合规则定义：YAML 驱动的模式匹配

graph-autofusion 使用 classify_rule.yaml 定义融合模式，支持拓扑结构与属性约束双重匹配。

2.1 规则语法示例

# classify_rule.yaml
fusion_rules:
  - name: "MatMulAddRelu"
    pattern:
      root: MatMul
      consumers:
        - op: Add
          inputs: [root.output, ?]  # ? 表示任意输入
          consumers:
            - op: Relu
    constraints:
      - MatMul.transpose_a == false
      - MatMul.transpose_b == true
    fusion_type: "super_kernel"

✅ 特性：

• 支持嵌套消费者（consumers of consumers）；
• ? 通配符匹配任意张量；
• 属性约束确保数值正确性。

2.2 规则加载与匹配引擎

规则由 RuleParser 加载为内部表示：

// src/rule_engine/rule_parser.cpp
struct FusionPattern {
    std::string root_op;
    std::vector<ConsumerPattern> consumers;
    std::vector<std::string> constraints;
};

class RuleMatcher {
public:
    bool Match(const Node* root, const FusionPattern& pattern) {
        if (root->op_type != pattern.root_op) return false;
        
        // 递归匹配消费者
        for (const auto& consumer_pat : pattern.consumers) {
            bool found = false;
            for (const auto& consumer : root->outputs[0]->consumers) {
                if (MatchConsumer(consumer, consumer_pat)) {
                    found = true;
                    break;
                }
            }
            if (!found) return false;
        }
        return CheckConstraints(root, pattern.constraints);
    }
};

该引擎支持多规则并发匹配，并记录所有候选融合子图。

---

3. 融合决策算法：成本模型驱动的选择

并非所有匹配成功的子图都值得融合。graph-autofusion 引入成本模型（Cost Model）进行收益评估。

3.1 成本模型设计

成本函数定义为：
$$\text{Benefit}={\text{Cost}}_{\text{unfused}}-{\text{Cost}}_{\text{fused}}$$
其中：

• ${\text{Cost}}_{\text{unfused}}=\sum \text{KernelLaunchOverhead}+\sum \text{MemoryAccess}$
• ${\text{Cost}}_{\text{fused}}=\text{FusedKernelCost}+\text{FusedMemoryAccess}$

关键指标包括：

• Kernel 启动次数：每次启动约 1–5 μs 开销；
• 中间张量大小：决定 DRAM 读写量；
• 计算强度：FLOPs / Bytes，影响带宽瓶颈。

3.2 成本估算实现

// src/cost_model/fusion_cost_evaluator.cpp
float FusionCostEvaluator::EstimateUnfusedCost(
    const std::vector<Node*>& subgraph) {
    float total_cost = 0.0f;
    for (auto* node : subgraph) {
        total_cost += kKernelLaunchOverhead; // ~2us
        total_cost += EstimateMemoryCost(node); // 基于 tensor size
    }
    return total_cost;
}

float FusionCostEvaluator::EstimateFusedCost(
    const std::vector<Node*>& subgraph) {
    // 假设融合后仅一次 Kernel 启动
    float cost = kKernelLaunchOverhead;
    
    // 内存成本：仅输入/输出张量
    TensorSet io_tensors = ExtractIOTensors(subgraph);
    for (auto& tensor : io_tensors) {
        cost += tensor.size * kDRAMBandwidthCost;
    }
    
    // 计算成本：基于 FLOPs
    int64_t flops = CalculateFLOPs(subgraph);
    cost += flops * kComputeUnitCost;
    
    return cost;
}

bool FusionCostEvaluator::ShouldFuse(
    const std::vector<Node*>& subgraph) {
    float benefit = EstimateUnfusedCost(subgraph) - 
                   EstimateFusedCost(subgraph);
    return benefit > kMinFusionBenefitThreshold; // 默认 5us
}

📌 阈值可配置：通过环境变量 GRAPH_AUTO_FUSION_MIN_BENEFIT 调整。

---

4. SuperKernel：JIT 编译与运行时集成

对于决策融合的子图，graph-autofusion 调用 SuperKernel 组件生成融合 Kernel。

4.1 Kernel 描述生成

首先将子图转换为 Kernel Description（KD）：

{
  "kernel_name": "MatMulAddRelu_128x128",
  "inputs": ["A", "B", "bias"],
  "outputs": ["C"],
  "ops": [
    {"type": "matmul", "inputs": ["A", "B"], "output": "tmp1"},
    {"type": "add", "inputs": ["tmp1", "bias"], "output": "tmp2"},
    {"type": "relu", "inputs": ["tmp2"], "output": "C"}
  ]
}

4.2 JIT 编译流程

SuperKernel 使用 codegen 模板引擎 生成 C++ Kernel：

// super_kernel/src/codegen/kernel_generator.cpp
std::string KernelGenerator::GenerateKernelCode(
    const KernelDesc& kd) {
    std::string code = R"(
__global__ void KERNEL_NAME(...) {
)";
    
    for (const auto& op : kd.ops) {
        if (op.type == "matmul") {
            code += GenerateMatMulCode(op);
        } else if (op.type == "add") {
            code += GenerateEWiseAddCode(op);
        } else if (op.type == "relu") {
            code += GenerateReluCode(op);
        }
    }
    
    code += "}\n";
    return code;
}

生成的代码通过 运行时编译器（如 NVCC 或 CANN 自研编译器）编译为二进制，并缓存至磁盘（~/.cache/super_kernel/）。

4.3 运行时注册

融合 Kernel 被注册为新算子：

// graph-autofusion/src/fusion_executor.cpp
void FusionExecutor::RegisterFusedKernel(
    const std::string& kernel_name,
    const void* kernel_func) {
    // 注册到全局算子表
    OpRegistry::Instance().RegisterOp(
        kernel_name,
        [kernel_func](const Context& ctx, const Args& args) {
            launch_kernel(kernel_func, ctx.stream, args);
        }
    );
}

后续图执行直接调用该融合算子。

---

5. 高级特性：动态 Shape 与控制流支持

5.1 动态 Shape 处理

对于输入 shape 在运行时确定的模型（如 NLP），graph-autofusion 采用 Shape Specialization：

• 首次遇到新 shape 时，触发 JIT 编译；
• 缓存按 (op_type, input_shapes) 索引；
• 支持 shape 通配（如 batch 维度任意）。

// super_kernel/src/cache/kernel_cache.cpp
std::string KernelCache::GetCacheKey(
    const std::string& pattern_name,
    const std::vector<Shape>& input_shapes) {
    std::string key = pattern_name;
    for (const auto& s : input_shapes) {
        // 将动态维度标记为 -1
        key += "_" + s.ToStringWithDynamic();
    }
    return key;
}

5.2 控制流融合限制

当前版本（v1.2.0）不融合跨越控制流边界（如 if/while）的算子，以确保语义正确性。该限制由 FusionAnalyzer 在图遍历时强制实施。

---

结语

CANN graph-autofusion 通过规则驱动的模式匹配与成本模型引导的决策算法，实现了高效、安全的自动算子融合。其轻量级、解耦式设计使得 SuperKernel 等组件可独立演进，而 YAML 规则系统为社区贡献提供了低门槛入口。随着对稀疏计算、量化融合等场景的支持（见 Roadmap），graph-autofusion 将成为 CANN 图优化栈中不可或缺的一环，持续推动模型执行效率的边界。

cann组织链接：https://atomgit.com/cann
graph-autofusion仓库链接：https://atomgit.com/cann/graph-autofusion