ge:昇腾CANN的图引擎架构剖析

#前言

刚接触昇腾CANN那会,图编译让我很困惑。那时候想要把一个PyTorch模型跑到昇腾NPU上,结果发现模型转换过来的图性能很差,推理延迟比GPU高了一倍还多。后来在社区里泡了一圈,才发现昇腾CANN其实已经把图引擎(ge)优化得很好了,只是很多人不知道怎么用。

仓库定位:图引擎

ge是昇腾CANN生态中的图引擎。它的定位非常明确:负责图的编译、优化和执行。

在CANN五层架构中,ge位于第3层:昇腾计算编译层。这一层的核心职责是提供图编译器(Graph Compiler)和BiSheng/ATC编译器。ge就是图编译器的核心组件。

为什么需要图引擎?因为深度学习框架(PyTorch、TensorFlow、MindSpore等)训练出来的模型,通常是一个计算图。这个计算图需要先被编译成底层硬件能执行的代码,然后才能执行。图引擎就是来做这个编译、优化、执行的工作的。

整体架构:三层设计

ge的整体架构可以分为三层:接口兼容层、自动调度层、优化实现层。

接口兼容层 负责对接各种深度学习框架。PyTorch、TensorFlow、MindSpore等框架,它们的计算图表示是不一样的。接口兼容层的作用就是:把这些不同格式的计算图,统一转换成ge内部的中间表示(IR)。

自动调度层 负责图的优化和调度。这一层的核心任务是:对计算图做各种优化(算子融合、内存优化、流水线优化等),然后把优化后的图调度到硬件上执行。

优化实现层 负责具体的优化实现。比如算子融合的实现、内存优化的实现、流水线优化的实现等。这一层会深度利用昇腾NPU的硬件特性,生成高效的机器码。

核心模块:四大组件

ge的核心模块可以分为四大组件:图解析器、图优化器、图执行器、内存管理器。

图解析器 负责解析各种格式的计算图。比如,PyTorch的模型可以通过TorchAir适配层转换成ge能处理的格式;TensorFlow的模型可以通过TFA(TensorFlow Adapter)转换成ge能处理的格式。

图优化器 负责对计算图做各种优化。ge内置了上百种图优化Pass,包括:算子融合(把多个小算子融合成一个大算子)、内存优化(减少内存拷贝次数)、流水线优化(让计算和通信重叠)等。

图执行器 负责执行优化后的计算图。ge使用运行时(runtime)提供的执行调度能力,把计算图调度到NPU上执行。

内存管理器 负责管理图执行过程中的内存。ge使用了内存池、内存复用、零拷贝等技术来优化内存管理。

数据流转:从模型到硬件执行

ge的数据流转过程可以分为以下几个步骤:

  1. 模型导入:用户把一个训练好的模型(比如PyTorch模型)导入到ge中。
  2. 图解析:ge的图解析器解析模型的计算图,转换成内部的中间表示(IR)。
  3. 图优化:ge的图优化器对计算图做各种优化(算子融合、内存优化、流水线优化等)。
  4. 图编译:ge把优化后的计算图编译成NPU能执行的机器码。
  5. 图执行:ge的图执行器把编译好的机器码调度到NPU上执行。
  6. 结果返回:ge把NPU执行的结果返回给用户。

设计理念:代码零修改,性能最大化

ge的核心设计理念是:代码零修改,性能最大化。

代码零修改 是指:用户在PyTorch、TensorFlow等框架上写的模型代码,不需要修改就能在昇腾NPU上运行。ge通过接口兼容层,自动把其他框架的模型转换成昇腾NPU能执行的格式。

性能最大化 是指:ge会对计算图做深度优化,让模型在昇腾NPU上的性能达到最优。ge内置了上百种图优化Pass,能充分利用昇腾NPU的硬件特性。

代码示例:使用ge进行模型编译

下面是一个简单的代码示例,展示如何使用ge进行模型编译和执行:

`python

import torch

import numpy as np

假设已经安装了昇腾CANN和ge

这里用伪代码展示调用流程

1. 准备PyTorch模型

model = MyModel() # 用户定义的PyTorch模型

model.eval()

2. 准备输入数据

input_data = torch.randn(1, 3, 224, 224)

3. 使用TorchAir把PyTorch模型转换成ge能处理的格式

from torch_air import torch_to_ge

ge_graph = torch_to_ge(model, input_data)

伪代码:模型转换

def torch_to_ge(model, input_data):

1. 导出ONNX

torch.onnx.export(model, input_data, 'model.onnx')

复制代码
# 2. 使用ge解析ONNX
# ge_graph = ge.OnnxParser('model.onnx').parse()

# 3. 返回ge图
return 'ge_graph'

ge_graph = torch_to_ge(model, input_data)

4. 使用ge优化图

ge_graph = ge.optimize(ge_graph)

伪代码:图优化

def ge_optimize(graph):

1. 算子融合

graph = ge.FusionPass()(graph)

复制代码
# 2. 内存优化
# graph = ge.MemoryOptimizationPass()(graph)

# 3. 流水线优化
# graph = ge.PipelineOptimizationPass()(graph)

return graph

ge_graph = ge_optimize(ge_graph)

5. 使用ge编译图

ge_compiled = ge.compile(ge_graph)

伪代码:图编译

def ge_compile(graph):

1. 生成IR

ir = ge.IRGenerator()(graph)

复制代码
# 2. 生成机器码
# machine_code = ge.CodeGenerator()(ir)

return 'compiled_graph'

ge_compiled = ge_compile(ge_graph)

6. 使用ge执行图

output = ge.execute(ge_compiled, input_data)

伪代码:图执行

def ge_execute(compiled_graph, input_data):

1. 分配内存

memory = ge.MemoryManager().allocate(compiled_graph)

复制代码
# 2. 拷贝输入数据
# ge.CopyData()(memory, input_data)

# 3. 执行计算
# ge.Executor().run(compiled_graph, memory)

# 4. 拷贝输出数据
# output = ge.CopyData()(memory, 'output')

return 'output'

output = ge_execute(ge_compiled, input_data)

print(f'Output shape: {output.shape}')

相关推荐
童谣110 小时前
越华环保集团市级统筹申报数字化中台:政策校验与减排仿真一体化架构
数据库·架构
2601_9567436812 小时前
上海小程序开发公司技术路径全解:从架构选型到工程落地的核心判断
小程序·架构·开发经验·上海
JAVA面经实录91713 小时前
图解23种设计模式完整知识体系(Java后端面试完整版)
java·架构
AI_小站13 小时前
Loop Engineering又是啥?一文讲清企业Agent落地的四层工程进化论
java·人工智能·架构·prompt·大模型开发·智能体·大模型应用
l1t13 小时前
DeepSeek总结的DuckLake 架构深度剖析-2
数据库·架构
2601_9625029013 小时前
双工位蜡镶机控制系统工程实践:并行架构、视觉对齐与实时同步避坑指南
架构
论迹复利14 小时前
多平台移植 BTBLE Stack 的思路与踩坑实战——从 Ceva IP 到厂商适配层的迁移指南
嵌入式硬件·物联网·架构
代码不会写16 小时前
Apache Iceberg:架构原理、读写机制、性能优化与生态
性能优化·架构·apache·iceberg
heimeiyingwang16 小时前
【架构实战】云原生架构:容器化与Kubernetes落地之路
云原生·架构·kubernetes
DogDaoDao16 小时前
【GitHub】WorldMonitor:一个工程极致主义的实时全球情报仪表盘深度解析
python·程序员·架构·github·go语言·worldmonitor·实时全球情报