前言

昇腾CANN作为昇腾异构计算架构的核心软件栈,其图引擎是连接上层框架与底层算子执行的关键组件。图引擎负责将框架定义的计算图转换为可在昇腾NPU上高效执行的图结构,并通过算子融合、内存优化、并行调度等技术手段,最大化NPU的利用率。在昇腾CANN五层架构中,图引擎位于第3层,承担计算图优化与编译的核心职责。本文深入剖析图引擎的架构设计、算子融合策略、调度机制,以及在昇腾NPU算子生态中的实际价值。

图引擎在CANN五层架构中的位置

根据CANN五层架构,图引擎位于第3层的核心位置。具体来说,第1层的AscendCL接收来自PyTorch、TensorFlow等框架的计算图描述。AscendCL将计算图传递给图引擎进行处理。图引擎对计算图进行一系列优化变换,包括算子融合、内存分配优化、并行性提取、算子调度等。

优化后的计算图被传递给第3层的图编译器和第4层的图执行器,最终在昇腾NPU上执行。图引擎的输出是优化后的计算图,该图不仅包含算子执行顺序,还包含每个算子的内存布局、数据分块策略、计算单元分配等底层执行细节。

这些细节对上层框架透明,框架开发者只需关注计算语义的正确性,性能优化由图引擎自动完成。这种设计使得框架开发者可以专注于模型本身的创新,而无需深入掌握底层硬件的优化技巧。图引擎充当了框架开发者与硬件优化之间的桥梁,大幅降低了在昇腾NPU上进行高性能计算的门槛。

图引擎架构设计

图引擎的架构分为四层:图解析层、图优化层、图编译层、图执行接口层。这种分层设计使得各层可以独立演进,同时保持层间接口的稳定性和向后兼容性。

图解析层负责将不同框架的计算图统一表示为图引擎的内部图结构。该层支持多种输入格式,包括ONNX、MindSpore IR、AscendCL Graph API等。不同框架的计算图在语义上存在差异,图解析层通过算子映射表将这些差异统一起来。

算子映射表定义了框架算子到昇腾算子库的映射关系,包括算子名称映射、参数顺序调整、数据类型转换等。这个映射表是图引擎支持新框架或新算子的核心维护点。当新增一个框架支持时,只需在映射表中添加相应的映射规则,无需修改图引擎的核心代码。

图优化层是图引擎的核心,包含数十种图优化过程。每个优化过程对计算图进行一次遍历和变换,多个优化过程按特定顺序组合成优化流水线。主要的优化过程包括算子融合、内存优化、死代码消除、算子替换等。

// WHY: 使用Visitor模式遍历计算图,而非手动递归遍历
// 因为计算图可能包含多种类型的节点(算子节点、数据节点、控制节点)
// Visitor模式将节点处理逻辑与节点类型解耦,新增节点类型时无需修改遍历代码
// 这种可扩展性对于支持新算子、新控制流非常关键

class GraphOptimizationPass {
public:
    virtual ~GraphOptimizationPass() = default;
    
    // 执行图优化过程
    // 输入:原始计算图
    // 输出:优化后的计算图
    virtual Status Run(Graph& graph) = 0;
    
    // WHY: 返回优化过程的名称,而非使用typeid或__func__
    // 因为名称用于日志输出、性能分析、调试信息
    // 人类可读的名称比编译器生成的名称更有用
    virtual std::string Name() const = 0;
};

// 算子融合过程
class OpFusionPass : public GraphOptimizationPass {
public:
    std::string Name() const override {
        return "OpFusionPass";
    }
    
    Status Run(Graph& graph) override {
        // 步骤1:识别可融合的算子模式
        // WHY: 先识别模式再执行融合,而非边识别边融合
        // 因为某些融合模式可能重叠(如A+B+C可以融合为A+B和B+C)
        // 先识别所有模式,再通过代价模型选择最优的融合方案
        std::vector<FusionPattern> patterns = IdentifyFusionPatterns(graph);
        
        // 步骤2:通过代价模型筛选有益的融合
        std::vector<FusionPattern> beneficial_patterns;
        for (const auto& pattern : patterns) {
            // WHY: 使用代价模型而非硬编码规则来决定是否融合
            // 因为融合的收益依赖于具体输入形状、硬件特征、内存带宽
            // 硬编码规则只能覆盖少数场景,代价模型可以自适应不同场景
            if (ShouldFuse(pattern, graph)) {
                beneficial_patterns.push_back(pattern);
            }
        }
        
        // 步骤3:执行融合
        for (const auto& pattern : beneficial_patterns) {
            Status ret = FusePattern(graph, pattern);
            if (ret != SUCCESS) {
                // WHY: 融合失败时不终止整个优化流水线,只跳过当前模式
                // 因为某些融合可能因为内存不足、算子不支持等原因失败
                // 跳过失败的融合,继续处理其他模式,提高鲁棒性
                LOG(WARNING) << "Failed to fuse pattern: " << pattern.ToString();
                continue;
            }
        }
        
        return SUCCESS;
    }
    
private:
    // 判断融合是否有益
    bool ShouldFuse(const FusionPattern& pattern, const Graph& graph) {
        // 估算融合前后的显存访问量
        size_t mem_access_before = EstimateMemoryAccess(graph, pattern.GetOps());
        size_t mem_access_after = EstimateFusedMemoryAccess(graph, pattern);
        
        // WHY: 只有当显存访问量减少一定比例时才进行融合
        // 因为融合会增加kernel编译时间、降低kernel的可读性
        // 只有当性能收益大于这些代价时才值得融合
        // 阈值(0.7)通过大量实验确定,可以根据硬件特性调整
        if (mem_access_after < mem_access_before * 0.7) {
            return true;
        }
        
        // 检查融合后的kernel大小是否超出限制
        // WHY: 过大的kernel会导致寄存器压力增大、缓存命中率降低
        // 需要限制融合后的kernel大小,通常不超过100个算子
        if (pattern.GetOpCount() > 100) {
            return false;
        }
        
        return false;
    }
};

内存优化过程分析计算图中每个张量的生命周期,识别可以复用显存的张量对,生成显存复用方案。该过程可以显著降低大模型推理时的显存峰值占用,使得更大的模型可以在同一张NPU上运行。内存优化的核心思想是存活期不重叠的张量可以共享同一块显存,通过静态分析识别这些张量对,生成显存复用计划。

图编译层将优化后的计算图转换为底层可执行形式。该层调用图编译器生成每个算子的底层指令流,并生成整个计算图的内存分配方案和调度计划。图编译的核心挑战是如何在保持性能的同时,支持动态shape和动态控制流。

// WHY: 使用基于区间的寄存器分配算法,而非简单的栈式分配
// 因为基于区间的算法可以考虑张量的生命周期,最大化显存复用
// 栈式分配会导致显存碎片,降低显存利用率
// 基于区间的算法虽然复杂度更高,但显存节省的收益远超编译时间开销

class MemoryPlanner {
public:
    // 生成显存分配方案
    MemoryPlan Plan(const Graph& graph) {
        // 步骤1:计算每个张量的生命周期区间
        // 生命周期区间 = [第一次使用, 最后一次使用]
        std::vector<LiveRange> live_ranges = ComputeLiveRanges(graph);
        
        // 步骤2:按生命周期区间排序(开始时间升序)
        std::sort(live_ranges.begin(), live_ranges.end(),
                  [](const LiveRange& a, const LiveRange& b) {
                      return a.start < b.start;
                  });
        
        // 步骤3:使用贪心算法分配显存
        // WHY: 使用贪心算法而非最优算法(如整数线性规划)
        // 因为最优算法的复杂度是NP-hard,编译时间不可接受
        // 贪心算法虽然不一定得到最优解,但在实践中效果很好
        // 关键是按照合适的顺序处理张量(如按大小降序)
        std::vector<MemoryBlock> allocated_blocks;
        MemoryPlan plan;
        
        for (const auto& range : live_ranges) {
            // 尝试找到可以复用的显存块
            bool reused = false;
            for (auto& block : allocated_blocks) {
                // WHY: 检查生命周期是否重叠,而非仅检查是否同时存活
                // 因为生命周期重叠意味着同一时刻两个张量都需要显存
                // 只有生命周期不重叠,才能安全复用显存
                if (!IsOverlapped(range, block.live_range)) {
                    // 可以复用
                    plan.Assign(range.tensor_id, block.address);
                    block.live_range = MergeLiveRange(range, block.live_range);
                    reused = true;
                    break;
                }
            }
            
            if (!reused) {
                // 无法复用,分配新的显存块
                size_t alignment = GetAlignment(range.tensor_desc);
                size_t size = GetSize(range.tensor_desc);
                
                // WHY: 使用边界对齐的分配,而非简单递增分配
                // 因为某些算子要求显存地址对齐到特定边界(如128字节)
                // 不对齐会导致算子执行失败或性能下降
                size_t address = AlignUp(next_free_address_, alignment);
                plan.Assign(range.tensor_id, address);
                
                MemoryBlock new_block;
                new_block.address = address;
                new_block.size = size;
                new_block.live_range = range;
                allocated_blocks.push_back(new_block);
                
                next_free_address_ = address + size;
            }
        }
        
        return plan;
    }
    
private:
    size_t next_free_address_ = 0;
    
    // 检查两个生命周期是否重叠
    bool IsOverlapped(const LiveRange& a, const LiveRange& b) {
        return !(a.end <= b.start || b.end <= a.start);
    }
    
    // 合并两个生命周期
    LiveRange MergeLiveRange(const LiveRange& a, const LiveRange& b) {
        LiveRange merged;
        merged.start = std::min(a.start, b.start);
        merged.end = std::max(a.end, b.end);
        merged.tensor_id = -1;  // 无效,仅用于生命周期
        return merged;
    }
};

图执行接口层向运行时提供优化后计算图的执行入口。该层提供执行接口,运行时通过调用这些接口触发计算图在NPU上的执行。接口层还负责输入数据的预处理和输出数据的后处理,确保数据格式符合算子的要求。

算子融合策略深度剖析

算子融合是图引擎最核心的优化技术,其目标是减少计算图中算子之间的显存读写次数。在NPU上,显存带宽通常比计算能力更稀缺,减少显存访问次数往往比减少计算次数更能提升性能。算子融合通过减少显存读写次数,显著提升计算图的执行效率。

图引擎支持以下融合模式。

流水融合

流水融合将计算图中具有生产者-消费者关系的多个算子融合为一个算子。例如,卷积、批归一化、ReLU这三个算子,在未融合情况下,卷积的输出需要写入显存,批归一化再从显存读取该输出;融合后,卷积的输出直接传递给批归一化,无需经过显存。

流水融合的触发条件是:生产者算子的输出只被一个消费者算子使用,且两个算子都支持融合接口。图引擎维护一个融合规则库,描述哪些算子之间可以进行流水融合。融合规则库通过模式匹配识别可融合的算子序列,然后通过代价模型决定是否执行融合。

流水融合的核心挑战是融合后的kernel代码生成。图引擎通过代码模板库实现融合kernel的自动生成。每个算子提供其计算表达式的代码片段,融合引擎将这些代码片段组合成一个完整的kernel函数。代码生成需要考虑输入输出张量的内存布局、数据类型的兼容性、边界条件的处理等复杂问题。

就地融合

就地融合将多个逐元素算子融合为一个算子。逐元素算子的每个输出元素只依赖于对应位置的输入元素,因此可以将其计算合并到同一个kernel中。例如,Transformer模型中的层归一化、残差连接、Dropout这三个逐元素操作,在未融合情况下需要三次显存读写;融合后只需一次显存读写,三次计算在同一个kernel中完成。

就地融合的优势在于显著减少显存读写次数,尤其对于逐元素算子这种内存带宽受限的操作,融合的收益非常明显。就地融合的挑战在于算子融合后的kernel代码复杂度急剧上升,调试困难。图引擎通过代码模板库和自动代码生成技术,降低融合kernel的开发难度。

跨层融合

跨层融合是最激进的融合模式,将跨越多个网络层的算子融合为一个超大算子。例如,整个Transformer模块可以被融合为一个算子。跨层融合可以最大化显存读写节省,但也带来kernel代码复杂度急剧上升、调试困难、编译时间变长等问题。

图引擎对跨层融合采取保守策略。只在明确知道收益大于代价的场景下触发跨层融合,且融合后的kernel大小有上限。上限通过配置参数控制,通常不超过100个算子。跨层融合的收益主要通过减少显存读写次数来衡量,只有当显存访问量减少一定比例(如百分之三十)以上时,才进行跨层融合。

动态shape融合

动态shape是指输入张量的某些维度在编译期未知。动态shape给算子融合带来额外挑战,因为融合kernel的分块策略和内存分配都依赖于具体shape。图引擎通过shape约束传播解决部分动态shape融合问题。

如果计算图中某些算子的输出shape可以由输入shape推导出来,图引擎可以在编译期为这些算子生成通用的融合kernel。对于完全无法在编译期确定shape的场景,图引擎采用即时编译融合。在运行时根据实际输入shape动态生成融合kernel。即时编译的首次执行有编译开销,但后续相同shape的输入可以直接复用已编译的kernel。

调度机制

图引擎的调度机制决定计算图中的每个算子在哪个计算单元上执行、以什么并行度执行。昇腾NPU包含多种计算单元,包括人工智能核心、向量计算单元等。不同计算单元适合执行不同类型的算子,图引擎的调度机制需要充分考虑这种异构性。

算子到计算单元的映射是调度机制的核心。图引擎根据算子的计算特性将其映射到合适的计算单元。矩阵乘法、卷积等计算密集型算子映射到人工智能核心;逐元素算子、数据重排算子等轻量级算子映射到向量计算单元。

某些算子可以同时利用多种计算单元。这类算子可以被拆分为多个子任务,分别送到不同的计算单元上并行执行。图引擎支持算子的异构执行,通过任务划分和依赖分析,实现跨计算单元的并行执行。

并行度决策是调度机制的另一核心问题。图引擎决定每个算子使用多少个计算核心并行执行。决策依据包括算子的工作量、当前NPU上其他算子的资源占用情况、算子的并行度上限等。并行度决策通过代价模型进行,选择预计最短的执行时间对应的并行度。

流水线调度是调度机制的高级特性。图引擎识别计算图中可以流水线化执行的子图,生成流水线调度计划。流水线调度将一个大的计算图拆分为多个阶段,每个阶段在一个NPU上执行,阶段之间通过显存或专用互联传递中间结果。流水线调度的核心价值是提升设备利用率。

与框架的集成

图引擎通过多种方式与上层框架集成。PyTorch集成通过适配器实现,PyTorch的计算图被转换为中间表示格式,然后传递给图引擎进行优化。TensorFlow集成通过框架适配器实现,TensorFlow的计算图被直接传递给图引擎。框架适配器处理TensorFlow算子到昇腾算子的映射,以及数据类型和格式的转换。

MindSpore作为原生框架,与图引擎的集成最为深入。MindSpore的计算图中间表示可以直接被图引擎解析,无需经过格式转换,减少转换开销和语义损失。这种深度集成使得MindSpore模型在昇腾NPU上的性能表现最优。

结尾

图引擎作为昇腾CANN的图优化与编译核心,在昇腾NPU的性能发挥中扮演关键角色。其图解析、图优化、图编译、图执行接口的四层架构,系统性地解决了异构计算图的优化与执行问题。算子融合、内存优化、并行调度等核心技术,使得上层框架的模型可以在昇腾NPU上高效执行。对于在昇腾NPU上部署模型的开发者,理解图引擎的优化机制和调度策略,是充分释放昇腾硬件算力的必经之路。

ge开源仓库:https://atomgit.com/cann/ge

CANN开源社区:https://atomgit.com/cann

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐