前言

昇腾CANN作为昇腾异构计算架构，昇腾CANN作为昇腾异构计算架构，昇腾CANN作为昇腾异构计算架构，图引擎GE可能是最核心但最容易被忽视的组件。很多人用PyTorch写模型，直接调用forward，以为这就是深度学习的全部。但实际上，你的PyTorch代码在执行之前，会经历一个复杂的编译过程：Python代码先转换成计算图，计算图经过一系列优化，最后生成NPU能执行的任务序列。GE就是这个编译过程的核心引擎。

GE的作用类似于编译器的后端优化器。它接收计算图（来自PyTorch、TensorFlow或ONNX），分析图中的算子和数据依赖，做算子融合、内存规划、布局优化、常量折叠等变换，最后生成高效的执行计划。这些优化对性能的影响巨大，可能带来2-10倍的性能提升。

一、GE在CANN架构中的位置

1.1 编译流水线

理解GE之前，先看完整的编译流水线。以PyTorch模型为例，你的Python代码会经历这些阶段：

第一阶段是前端编译。PyTorch把Python代码转换成计算图，每个nn.Module对应一个子图，forward方法的调用对应图中的边。

第二阶段是图优化，这就是GE的工作。GE分析计算图，找到可以融合的算子，规划内存布局，消除冗余计算。

第三阶段是后端编译。优化后的计算图被转换成NPU能执行的形式，包括CCE代码生成、任务调度、内存分配等。

# 编译流水线示例

compile_pipeline = """
PyTorch模型到NPU执行的完整流程：

原始代码（Python）：
    class Model(nn.Module):
        def forward(self, x):
            x = self.conv1(x)
            x = self.bn1(x)
            x = self.relu(x)
            return x

    ↓ 前端编译（PyTorch）

计算图（TorchScript IR）：
    input → Conv2d → BatchNorm → ReLU → output

    ↓ 图优化（GE）

优化后计算图：
    input → FusedConvBNReLU → output
    （内存布局：NCHW → NC1HWC0）

    ↓ 后端编译（TBE/CCE）

NPU任务序列：
    Task 1: FusedConvBNReLU kernel
    - 输入地址：0x10000000
    - 输出地址：0x20000000
    - 参数：...

    ↓ 执行（Runtime）

NPU执行：
    - 加载任务到设备
    - 执行计算
    - 返回结果
"""

# 为什么需要这么多阶段？
# 因为Python代码和NPU执行之间差距巨大。
# Python是解释执行的动态语言，
# NPU需要静态的、高度优化的任务序列。
# 编译流水线就是把动态的Python代码，
# 转换成静态的NPU任务。
# GE在其中负责最关键的优化阶段。

1.2 GE的核心职责

GE有四个核心职责：算子融合、内存规划、布局优化、常量折叠。每个职责都对性能有重要影响。

算子融合是把多个小算子合并成一个大算子。比如Conv+BN+ReLU三个算子，融合成一个算子执行，减少内存访问次数。

内存规划是决定每个张量在内存中的位置。NPU的内存层次复杂，有HBM、片上缓存、寄存器，不同位置访问速度差几十倍。好的内存规划可以减少内存访问延迟。

布局优化是决定张量的数据排列方式。NPU对某些布局有硬件加速，比如NC1HWC0格式。GE会把输入数据转换成最优布局。

常量折叠是把可以在编译期计算的表达式提前算好。比如矩阵乘法中有一个常数矩阵，可以在编译时就算出部分结果，运行时就少算一些。

# GE优化示例

import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(3, 64, 3, padding=1)
        self.bn = torch.nn.BatchNorm2d(64)
        self.relu = torch.nn.ReLU()
    
    def forward(self, x):
        # 未优化的执行：
        # x1 = self.conv(x)    # 算子1
        # x2 = self.bn(x1)     # 算子2
        # x3 = self.relu(x2)   # 算子3
        # return x3
        
        # GE优化后的执行：
        # x = fused_conv_bn_relu(x, ...)  # 融合算子
        # return x
        
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

# 融合为什么能提升性能？
# 未融合版本需要三次内存读写：
# Conv输出写入HBM → BN从HBM读取 → BN输出写入HBM → ReLU从HBM读取
# 融合版本只需要一次内存读写：
# 融合算子输出直接写入HBM
# 内存访问是性能瓶颈，减少访问次数能显著提升性能。

二、算子融合技术

2.1 融合规则

GE内置了大量融合规则，覆盖常见的算子组合。Conv+BN+ReLU、MatMul+Add、Softmax+Dropout、LayerNorm+Add+Dropout等，都有对应的融合规则。

融合规则的核心思想是：如果两个算子可以合并执行，且合并后不会增加计算量，就应该融合。判断两个算子能否融合，需要考虑数据依赖、形状兼容性、精度要求等因素。

# 融合规则示例

fusion_rules = """
GE内置的常见融合规则：

1. 线性层融合
   - MatMul + Add → FusedLinear
   - MatMul + Add + ReLU → FusedLinearReLU
   
2. 卷积层融合
   - Conv + BN → FusedConvBN
   - Conv + BN + ReLU → FusedConvBNReLU
   - Conv + ReLU → FusedConvReLU
   
3. 注意力融合
   - Softmax + Dropout → FusedSoftmaxDropout
   - Q*K + Softmax + V*Attn → FusedAttention
   
4. 归一化融合
   - LayerNorm + Add + Dropout → FusedLayerNormAddDropout
   - BatchNorm + ReLU → FusedBNReLU

5. 激活函数融合
   - Conv + GELU → FusedConvGELU
   - Linear + GELU → FusedLinearGELU

融合条件：
- 数据依赖：后一个算子的输入正好是前一个算子的输出
- 形状兼容：融合后的形状能正确推导
- 精度要求：融合不会影响数值精度
"""

# 为什么这些组合可以融合？
# 因为它们满足融合的基本条件：
# 数据流是线性的（没有分支），
# 每个算子的计算可以嵌入到前一个算子的末尾，
# 融合后的计算量不增加。
# 比如ReLU，只需要在卷积输出后判断正负，
# 不需要额外的内存访问。

2.2 自动融合流程

GE会自动分析计算图，找到可以融合的算子组合。这个过程不需要用户干预，GE会遍历图中的所有节点，尝试应用融合规则。

自动融合的流程是：首先识别图中的算子类型和连接关系，然后匹配融合规则库，找到所有可以融合的算子组，生成融合后的新算子，最后替换原图中的算子组。

三、内存规划技术

3.1 内存层次

NPU的内存分三个层次：HBM（高带宽内存）、片上缓存、寄存器。HBM容量大（几十GB）但速度慢，片上缓存容量小（几MB）但速度快，寄存器最快但容量最小（几十KB）。

GE的内存规划目标是：尽可能让数据留在片上缓存，减少HBM访问。对于必须访问HBM的数据，要优化访问模式，提高带宽利用率。

# 内存层次示例

memory_hierarchy = """
NPU内存层次（以Ascend 910为例）：

1. HBM（高带宽内存）
   - 容量：32-64 GB
   - 带宽：约1.2 TB/s
   - 延迟：约100-200 ns
   - 用途：存储模型参数、输入输出、中间结果

2. 片上缓存（L1/L2 Buffer）
   - 容量：约8-16 MB
   - 带宽：约10 TB/s
   - 延迟：约10-20 ns
   - 用途：存储当前计算的临时数据

3. 寄存器
   - 容量：约64 KB（Vector）/ 512 KB（Cube）
   - 带宽：约100 TB/s
   - 延迟：约1-2 ns
   - 用途：存储正在计算的元素

GE的内存规划策略：
- 大张量（模型参数）放在HBM
- 中间结果尽可能留在片上缓存
- 当前计算的数据加载到寄存器
- 融合算子减少HBM写入
"""

# 为什么片上缓存这么重要？
# 因为HBM延迟是片上缓存的10倍，
# 如果每次计算都要从HBM读数据，
# 大部分时间都在等内存，而不是在计算。
# GE通过算子融合，让中间结果留在片上缓存，
# 省掉HBM访问，大幅提升性能。

3.2 内存复用

除了优化数据放置，GE还会做内存复用。两个张量如果生命周期不重叠，可以共享同一块内存。比如一个网络的中间层，前向计算完成后就不需要了，它的内存可以给反向计算用。

内存复用可以显著降低显存占用。一个看起来需要10GB显存的模型，经过GE的内存复用优化后，可能只需要6GB。

使用前 vs 使用后：GE优化的效率对比

指标	使用前（无GE优化）	使用后（GE优化）	提升效果
ResNet-50推理延迟	约18ms/张	约5ms/张	约3.6倍加速
BERT-base推理延迟	约45ms/序列	约12ms/序列	约3.8倍加速
显存占用（ResNet-50）	约8GB	约4GB	减少50%
算子执行次数	约50个/前向	约20个/前向	减少60%
HBM访问次数	约200次/前向	约50次/前向	减少75%

GE的性能提升主要来自三个方面。第一，算子融合减少内存访问，这是最大的性能瓶颈。第二，内存规划优化数据放置，让热点数据留在片上缓存。第三，内存复用降低显存占用，允许更大的批量。

四、布局优化技术

4.1 数据布局

数据布局是指张量在内存中的排列方式。最常见的布局是NCHW（batch, channel, height, width），这是PyTorch的默认布局。但NPU对某些布局有硬件加速，比如NC1HWC0格式。

NC1HWC0是把channel维度分成C0=16的小块。对于float16数据，一次内存访问可以读16个元素，正好是一个向量。这种布局可以提高内存带宽利用率。

# 布局转换示例

layout_example = """
数据布局对比：

NCHW格式（PyTorch默认）：
- 形状：(N, C, H, W)
- 内存顺序：先遍历batch，再channel，再height，最后width
- 优点：与PyTorch兼容，直观
- 缺点：连续访问width方向的数据才是连续的

NC1HWC0格式（NPU优化）：
- 形状：(N, C1, H, W, C0)，其中C1 = C / C0，C0 = 16
- 内存顺序：先batch，再channel组，再height，width，最后channel内元素
- 优点：连续访问C0个元素，内存带宽利用率高
- 缺点：需要转换，有额外开销

GE的布局优化：
- 在图编译阶段自动转换布局
- 在NPU执行时使用最优布局
- 在输出时转换回NCHW（如果需要）
"""

# 为什么NC1HWC0格式更快？
# 因为NPU的Vector单元一次处理16个float16，
# NC1HWC0格式下，连续的16个channel正好是一个向量，
# 一次内存访问就能加载完整的向量数据。
# 而NCHW格式下，16个channel分散在内存中，
# 需要多次访问才能凑齐一个向量。

4.2 自动布局转换

GE会自动插入布局转换节点。当PyTorch的NCHW张量传入NPU时，GE会插入一个转置操作，把NCHW转换成NC1HWC0。当计算结果输出给PyTorch时，GE又会转换回NCHW。

这些转换有开销，但通常可以摊薄。如果一个张量被多次使用（比如残差连接），只需要转换一次，后续使用都是优化布局。

五、调试与可视化

5.1 图可视化

GE提供了图可视化工具，可以把计算图导出成图片或JSON格式。这对于调试优化效果很有帮助，你可以看到哪些算子被融合了，内存是怎么规划的。

# 图可视化示例

import torch

model = ...  # 你的模型

# 导出计算图
scripted = torch.jit.script(model)

# 保存GE优化后的图
torch_npu.npu.set_dump_options("./graph_dump")
output = model(input_tensor)

# 在./graph_dump目录下可以找到：
# - 原始图（before_optimization.json）
# - 优化后图（after_optimization.json）
# - 融合关系（fusion_info.txt）

5.2 性能分析

GE还提供了性能分析工具，可以看到每个算子的执行时间、内存访问次数、融合效果等。这些数据对于优化模型很有参考价值。

六、总结

GE是昇腾CANN的核心图引擎，负责计算图的优化和编译。算子融合、内存规划、布局优化、常量折叠，这些优化对性能有巨大影响，可能带来2-10倍的性能提升。算子融合是GE最重要的优化技术，它把多个小算子合并成一个大算子，减少内存访问次数。Conv+BN+ReLU、MatMul+Add这些常见组合，GE都有对应的融合规则。内存规划让热点数据留在片上缓存，减少HBM访问延迟。内存复用让不重叠的张量共享内存，降低显存占用。布局优化把数据转换成NPU友好的格式，提高内存带宽利用率。NC1HWC0格式让连续访问更高效，GE会自动插入转换节点。

---

仓库链接：https://atomgit.com/cann/ge