本文基于CANN开源社区的ge仓库进行技术解读

CANN组织地址：https://atomgit.com/cann

ge仓库地址：https://atomgit.com/cann/ge

前言

深度学习框架（PyTorch、TensorFlow）定义的模型，怎么在NPU上高效执行？这中间需要一个"翻译官"——把框架的计算图转换成NPU能理解和执行的指令。

GE（Graph Engine）就是这个"翻译官"，它是CANN的图引擎，负责计算图的编译、优化和执行。

什么是GE

GE是CANN的核心组件，处理从框架到硬件的整个流程：

框架层（PyTorch/MindSpore）
         ↓
计算图（Graph）
         ↓
GE图引擎 ← 这一层
  - 图解析
  - 图优化
  - 算子选择
  - 内存分配
  - 指令生成
         ↓
NPU执行

GE的核心功能

1. 图解析

将框架的计算图转换为GE内部表示：

# 框架定义的模型
import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.bn = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
  
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

# GE会解析成计算图：
# Input → Conv2d → BatchNorm2d → ReLU → Output

2. 图优化

对计算图进行各种优化：

算子融合

优化前：
Conv2d → BatchNorm2d → ReLU

优化后：
ConvBnRelu（融合算子）

好处：
- 减少内存访问
- 减少kernel启动开销
- 提升性能

常量折叠

# 优化前
x = input * 2 + 3

# 优化后（如果2和3是常量）
x = input * 2 + 3  # 在编译时计算好

死代码消除

# 优化前
def forward(x):
    y = x * 2  # y没有被使用
    z = x + 1
    return z

# 优化后
def forward(x):
    z = x + 1
    return z

3. 算子选择

为每个算子选择最优实现：

同一个Conv2d算子可能有多种实现：
- 直接卷积
- im2col + GEMM
- Winograd算法
- FFT卷积

GE会根据：
- 输入shape
- 卷积核大小
- 硬件特性
选择最优实现

4. 内存优化

优化内存分配和复用：

优化前：
Conv1: 分配内存A
BN1:   分配内存B
ReLU1: 分配内存C
Conv2: 分配内存D

优化后（内存复用）：
Conv1: 分配内存A
BN1:   复用内存A
ReLU1: 复用内存A
Conv2: 分配内存B

5. 执行调度

生成高效的执行计划：

串行执行：
Op1 → Op2 → Op3 → Op4

并行执行（如果Op2和Op3无依赖）：
Op1 → Op2 ↘
           → Op4
      Op3 ↗

图优化Pass

GE包含多种优化Pass：

1. 算子融合Pass

// 伪代码
class FusionPass {
    void Run(Graph& graph) {
        // 查找可融合的模式
        for (auto pattern : fusion_patterns) {
            auto matches = FindPattern(graph, pattern);
            for (auto match : matches) {
                // 替换为融合算子
                ReplaceWithFusedOp(graph, match);
            }
        }
    }
};

// 常见融合模式：
// Conv + BN + ReLU
// MatMul + Add + GELU
// Softmax + Dropout

2. 常量折叠Pass

class ConstantFoldingPass {
    void Run(Graph& graph) {
        for (auto node : graph.nodes()) {
            if (AllInputsAreConstant(node)) {
                // 在编译时计算结果
                auto result = EvaluateNode(node);
                // 替换为常量节点
                ReplaceWithConstant(graph, node, result);
            }
        }
    }
};

3. 公共子表达式消除Pass

class CSEPass {
    void Run(Graph& graph) {
        std::map<Expression, Node*> expr_map;
      
        for (auto node : graph.nodes()) {
            auto expr = GetExpression(node);
            if (expr_map.count(expr)) {
                // 找到重复计算，复用结果
                ReplaceWith(node, expr_map[expr]);
            } else {
                expr_map[expr] = node;
            }
        }
    }
};

4. 内存优化Pass

class MemoryOptimizationPass {
    void Run(Graph& graph) {
        // 分析张量生命周期
        auto lifetimes = AnalyzeLifetimes(graph);
      
        // 分配内存
        for (auto tensor : graph.tensors()) {
            // 尝试复用已释放的内存
            auto reused_mem = FindReusableMemory(tensor, lifetimes);
            if (reused_mem) {
                AssignMemory(tensor, reused_mem);
            } else {
                AllocateNewMemory(tensor);
            }
        }
    }
};

图执行模式

1. 静态图模式

编译一次，多次执行：

import torch

# 定义模型
model = MyModel()

# 第一次执行：编译
output = model(input1)  # GE编译图

# 后续执行：直接运行编译好的图
output = model(input2)  # 快
output = model(input3)  # 快

优点：

• 性能好（编译优化充分）
• 内存占用少

缺点：

• 不支持动态shape
• 调试困难

2. 动态图模式

每次执行都重新构建图：

# 动态图模式
for i in range(10):
    if i < 5:
        output = model.path1(input)
    else:
        output = model.path2(input)

优点：

• 灵活
• 易于调试

缺点：

• 性能较差
• 内存占用大

图IR（中间表示）

GE使用多层IR表示计算图：

1. 高层IR

接近框架的表示：

Graph {
    Node: Conv2d
        input: [1, 3, 224, 224]
        weight: [64, 3, 3, 3]
        output: [1, 64, 222, 222]
  
    Node: BatchNorm2d
        input: [1, 64, 222, 222]
        output: [1, 64, 222, 222]
  
    Node: ReLU
        input: [1, 64, 222, 222]
        output: [1, 64, 222, 222]
}

2. 中层IR

优化后的表示：

Graph {
    Node: ConvBnRelu (融合算子)
        input: [1, 3, 224, 224]
        weight: [64, 3, 3, 3]
        bn_scale: [64]
        bn_bias: [64]
        output: [1, 64, 222, 222]
}

3. 低层IR

接近硬件的表示：

Instructions {
    LoadData: input → L1Buffer
    LoadWeight: weight → L1Buffer
    Conv: L1Buffer → L1Buffer
    BN: L1Buffer → L1Buffer
    ReLU: L1Buffer → L1Buffer
    StoreData: L1Buffer → output
}

实际案例

案例1：ResNet优化

# 原始ResNet Block
class ResNetBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
  
    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu1(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += identity
        out = self.relu2(out)
        return out

# GE优化后的执行图：
# Conv1+BN1+ReLU1 (融合)
# Conv2+BN2 (融合)
# Add
# ReLU2

案例2：Transformer优化

# 原始Attention
class Attention(nn.Module):
    def forward(self, q, k, v):
        # Q @ K^T
        scores = torch.matmul(q, k.transpose(-2, -1))
        # Scale
        scores = scores / math.sqrt(d_k)
        # Softmax
        attn = torch.softmax(scores, dim=-1)
        # Dropout
        attn = self.dropout(attn)
        # @ V
        output = torch.matmul(attn, v)
        return output

# GE优化：
# 1. MatMul + Scale 融合
# 2. Softmax + Dropout 融合
# 3. 使用FlashAttention算子（如果支持）

调试和分析

1. 查看计算图

# 保存计算图
torch.jit.save(model, "model.pt")

# 可视化计算图
# 使用Netron等工具查看

2. 性能分析

# 开启profiling
import torch_npu
from torch_npu.profiler import Profile

with Profile() as prof:
    output = model(input)

# 查看每个算子的执行时间
print(prof.key_averages().table())

3. 内存分析

# 查看内存使用
print(f"内存占用: {torch.npu.memory_allocated() / 1024**2:.2f} MB")
print(f"峰值内存: {torch.npu.max_memory_allocated() / 1024**2:.2f} MB")

性能优化建议

1. 使用静态shape

# 好：固定shape
input = torch.randn(32, 3, 224, 224)

# 不好：动态shape
for batch_size in [16, 32, 64]:
    input = torch.randn(batch_size, 3, 224, 224)
    # 每次都要重新编译

2. 避免频繁的CPU-NPU数据传输

# 不好
for i in range(100):
    input_npu = input_cpu.npu()  # 频繁传输
    output = model(input_npu)
    result = output.cpu()  # 频繁传输

# 好
input_npu = input_cpu.npu()  # 一次传输
for i in range(100):
    output = model(input_npu)
result = output.cpu()  # 一次传输

3. 使用算子融合

# GE会自动融合，但可以手动提示
# 使用torch.jit.script
@torch.jit.script
def fused_ops(x, y):
    z = x + y
    z = z * 2
    z = torch.relu(z)
    return z

常见问题

问题1：编译时间长

# 第一次执行慢（编译）
output = model(input)  # 慢

# 后续执行快
output = model(input)  # 快

# 解决方案：预编译
model = torch.jit.trace(model, example_input)

问题2：动态shape支持

# 如果必须使用动态shape
# 可以设置几个固定的shape档位
# GE会为每个档位编译一次

问题3：内存占用大

# 使用梯度检查点
from torch.utils.checkpoint import checkpoint

def forward(x):
    x = checkpoint(layer1, x)
    x = checkpoint(layer2, x)
    return x

应用场景

场景一：模型训练

GE优化训练图，提升训练速度。

场景二：模型推理

GE优化推理图，降低推理延迟。

场景三：模型部署

GE生成优化的执行计划，适合生产环境。

场景四：算子开发

基于GE开发自定义算子和优化Pass。

总结

GE是CANN的图引擎，主要功能：

• 计算图解析和转换
• 多种图优化Pass
• 算子选择和融合
• 内存优化和执行调度
• 支持静态图和动态图

对于理解CANN的工作原理，GE是核心组件。

相关链接

ge仓库地址：https://atomgit.com/cann/ge

CANN组织地址：https://atomgit.com/cann

metadef仓库地址：https://atomgit.com/cann/metadef