Skip to content

InfiniTensor/RefactorGraph

Folders and files

NameName
Last commit message
Last commit date

Latest commit

9eddc91 · Nov 27, 2024
Nov 6, 2023
Jan 26, 2024
Nov 6, 2023
Dec 25, 2023
Feb 19, 2024
Dec 6, 2023
Jul 14, 2023
Aug 22, 2024
Jan 26, 2024
Aug 22, 2024
Oct 31, 2023
Jan 4, 2024
Jan 26, 2024

Repository files navigation

重构图表示

Build issue license

目录

安装

使用 make install-python 编译并安装 Python 前端到全局,安装的包名为 refactor_grpah

环境变量:

  • 添加 TYPE=DebugTYPE=Release 以启用指定的优化级别,默认为 Debug
  • 添加 CUDA=OFFCUDA=ON 以打开或关闭英伟达显卡支持,默认为 OFF

使用前端

import sys
import numpy as np
from onnx import load
from refactor_graph.onnx import make_compiler, find_device
from onnxruntime import InferenceSession

model = load(sys.argv[1])  # ------------------------------------ 加载模型
input = np.random.random((10, 3, 224, 224)).astype(np.float32)  # 加载测试样本

compiler = make_compiler(model)  # ------------------------------ 模型导入到编译器
compiler.substitute("N", 10)  # --------------------------------- 代换输入中的变量
find_device("nvidia", 0)  # ------------------------------------- 初始化指定加速硬件
executor = compiler.compile("cuda", "default", [])  # ----------- 编译模型(选择平台、分配器和优化选项)
executor.set_input(0, input)  # --------------------------------- 设置输入
executor.dispatch(find_device("nvidia", 1), "default")  # ------- 执行器可以随时调度到另一个硬件
executor.run()  # ----------------------------------------------- 推理

session = InferenceSession(model.SerializeToString())  # -------- 与 onnxruntime 对比结果以验证推理
answer = session.run(None, {session.get_inputs()[0].name: input})
print([(executor.get_output(i) - answer[i]).flatten() for i in range(len(answer))])

对于使用外部数据的模型,支持直接加载以减少一次拷贝:

import sys
from pathlib import Path
from onnx import load
from refactor_graph.onnx import make_compiler

model_path = Path(sys.argv[1])  # ---------------------------- 假设模型和数据保存在相同路径
model = load(model_path.as_uri(), load_external_data=False)  # 不直接加载外部数据以避免额外拷贝
compiler = make_compiler(model, model_path.parent.as_uri())  # 导入时直接加载外部数据
executor = compiler.compile("cuda", "default", [])  # -------- 编译模型

# 下同

调试功能

项目现已依托前端提供多种调试功能。

  1. 列出算子信息

    executor.dbg()
  2. 保存运行中间结果

    executor.trace("path_to_store_files", "data_file_format")

    调用这个方法将启动一次模型推理,并在每个算子推理完成后将算子的所有输入输出张量保存到 path_to_store_files 参数指示的目录中。 如果目录不存在,将创建此目录。每个张量保存到一个文件,已存在的同名文件将被删除。 同时,为每个算子创建一个元信息文本文件,命名为 node<N>.meta,其内容具有下述格式:

    <NodeName>\t<N>
    <input/output>\t<K>\t<EdgeName>\t[FileName]
    
    • NodeName: 节点的名字;
    • N: 节点序号;
    • input/output: 张量是节点的输入/输出;
    • K: 输入/输出的序号;
    • EdgeName: 张量的名字;
    • FileName: (optional) 数据文件名。如果张量无效,不会保存数据文件,则文件名为空;
  3. 逐算子计时

    executor.bench(<sync>)

    对每次推理计时。sync 是一个指示是否在每次推理后插入同步的布尔参数,若设置为 False,则计时可能是推理异步启动的时间。

项目结构

构建系统

项目构建系统采用 CMake,方便集成一些第三方库。所有第三方库以 git submodule 的形式导入,公共的位于根目录下的 3rd-party 目录下,另有专用于 python 前端的 pybind11 位于 src/09python_ffi 目录下。

整个项目的源码以子项目的形式解耦,放在 src 的子目录中,每个子项目有自己的 CMakeLists.txt,并由根目录的 CMakeLists.txt 调用。src 的每个子目录带有一个编号,其中编号大的可以依赖编号小的,方便维护子项目之间的依赖关系,避免循环依赖。当前已有 00-09 共 10 个子项目,它们之间的依赖关系如下图所示:

┌─────────┐ ┌──────────────┐ ┌───────────┐ ┌──────────┐
│ 00cmmon ├←┤ 01graph_topo ├←┤ 03runtime ├←┤ 04kernel │
└───┬─────┘ └──────────────┘ └─────┬─────┘ └─────┬────┘
    ↑                              │             ↑
    │       ┌───────────────┐      │     ┌───────┴───────┐
    └───────┤ 02mem_manager ├←─────┘     │ 05computation │
            └───────────────┘            └───────┬───────┘
┌────────────────────────────────────────────┐   ↑
│ ┌──────────────┐ ┌────────┐ ┌────────────┐ │   │
│ │ 09python_ffi ├→┤ 07onnx ├→┤ 06frontend ├─┼───┘
│ └─────┬────────┘ └────────┘ └──────┬─────┘ │
│       │     ┌─────────────────┐    ↑       │
│       └────→┤ 08communication ├────┘       │
│ frontend    └─────────────────┘            │
└────────────────────────────────────────────┘

所有子项目使用 PUBLIC 依赖向下传递自己依赖,并使用 CMake 的 target_include_directories 机制使子项目头文件目录随依赖关系传递。

操作 CMake、构建目录和其他项目管理功能的命令和配置封装在根目录下的 Makefile 中,现有的命令包括:

  • build: 默认命令,构建项目。
  • install-python: 将 Python 前端安装到系统路径。
  • reconfig: 清除 CMake 缓存,以重新配置 CMake。
  • clean: 删除构建目录。
  • clean-log: 清除日志目录。
  • test: 执行单元测试。
  • format: 调用格式化工具。

子项目简介

源码的 10 个子项目的简介如下:

序号 项目 说明
0 common 用于所有子项目的类型和函数定义。
1 graph_topo 与元素类型解耦的图拓扑结构表示,包括存储结构和变换算法。
2 mem_manager 存储空间管理抽象。
3 runtime 运行时,执行模型推理的图层。
4 kernel 核函数层,包含核函数库以及从核函数图层下降到运行时图层的算法。
5 computation 计算图层,包含算子的语义表示定义,以及在算子语义层次进行图变换的规则定义。
6 frontend 前端图层,支持混合存储不同编程框架导入的前端算子,以及基于前端算子的张量形状推导和动态性消解、常量折叠机制。从前端图层下降到计算图层时,图中的动态性(即输出形状的计算还和形状中的变量)必须全部消解。
7 onnx onnx 前端算子库。
8 communication 分布式通信算子库。
9 python_ffi Python 前端项目。

点击项目名可以跳转到各个项目的文档。

第三方依赖

技术要点

多层计算图表示

目前主流的深度学习框架几乎都应用了两种基本的设计,即计算图和多层 IR,本框架也应用了这两种设计。

计算图是 AI 模型计算的一种表示形式。在 AI 模型的高层表示中,通常使用以算子为节点、张量为边的图结构来表示模型,例如,ONNX 模型可视化后通常表现为这样的形式:

onnx model

计算图的本质在逻辑上是一张数据流图,在数据结构上是一张有向无环图。数据流图意味着其中的节点表示一种运算的过程,而边表示在数据在运算之间流动的起点和终点。与经典的数据结构意义上的有向无环图相比,计算图中节点的输入和输出是有序的,不可互换,例如某个节点入度为 3 不是对节点输入的完备描述,必须说明第 0 个、第 1 个、第 2 个入边分别是什么。

产品级的 AI 编译器总是具有较长的工作流程。首先,模型从多种框架的高层表示(Pytorch/ONNX/……)转化到编译器定义的计算图形式,然后应用各种图上的变换,在确定的硬件上选择合适的计算方法(kernel),最后实际执行。在流程的不同阶段,需要关注和维护的信息是不同的。举例来说,Reshape 算子表示改变张量形状的语义,在高层的模型表示中是必要的,但在 AI 应用程序实际执行时,Reshape 算子不会真正改变数据,甚至可以直接从图上删除。因此,最好使用多层 IR 表示模型,在不同的层级改变关注的信息,从而降低每层的复杂度。

图拓扑抽象

灵活的多层 IR 表示要求层与层之间尽量使用无关的节点和边类型,以保证每层的自由设计,同时要尽量复用拓扑结构的表示,因为无论在哪一层,图结构的表示和操作方法是类似的,因此,本框架采用了拓扑结构与节点/边信息解耦的实现方式。拓扑结构以一系列模板类型的形式定义在 graph_topo 子项目中。这些拓扑结构表示被定义成类似容器类型的形式,只包含关键的节点和边的有向无环、出入有序的基本信息,对节点和边具体是什么没有约束,从而支持在不同的计算图中复用这些容器,并能在不与编译器的业务耦合的情况下编写、优化和测试,具有良好的工程特性。

目前,图拓扑表示包含以下主要的类型定义:

  • GraphTopo: 精简、连续存储的拓扑类型,用于持久化存储和、遍历和快速随机访问;
  • Searcher: 基于 GraphTopo 引用建立的查询缓存结构,用于包含图上一些冗余但常用的信息,使这些信息不必一直随着拓扑结构移动或拷贝;
  • Builder: 所有字段都公开可访问的结构体类型,可以表示拓扑结构并支持开发者自由操作,用于快速构建拓扑结构,并提供一个方法将拓扑信息压缩,转化成 GraphTopo
  • LinkedGraph: 链式的拓扑表示,在这种表示中修改拓扑连接关系、增删节点具有 O(1) 时间复杂度,但空间复杂度更高、不保证维持拓扑序,访问时也更容易缓存不命中;

参考子项目依赖关系图,几乎所有依赖 graph_topo 的子项目(除了前端算子库和接口层)都定义了一套最适于表示其信息的节点和边定义,并复用上述拓扑结构类型来表示一种特定的计算图。