Relay IR的核心数据结构

Apache TVMRelay IR 中,基础节点(VarConstCallFunctionExpr)是构建计算图的核心数据结构。以下是对它们的详细解析,包括定义、作用、内部组成及相互关系:


1. Expr(表达式基类)

作用

  • 所有 Relay IR 节点的基类,提供统一的类型系统和遍历接口。
  • 支持递归访问、变换和类型检查。

关键组成

字段/方法 说明
checked_type_ 表达式的推断类型(如 TensorTypeTupleType)。
span 源代码位置信息(用于调试和错误报告)。
VisitAttrs(visitor) 递归访问所有属性和子节点(用于序列化、优化等)。
Mutate() 生成表达式的副本(用于变换和优化)。

子类关系

Expr Var Const Call Function Tuple Let If


2. Var(变量)

作用

  • 表示计算图中的 输入变量中间变量(符号化张量)。
  • 类似于深度学习模型中的输入占位符或中间激活值。

关键组成

字段 说明
name_hint 变量名称(字符串标识符,如 "x")。
type_annotation 变量的显式类型注解(可选,如 TensorType({1,3}, float32))。
vid 内部唯一 ID(用于优化和去重)。

示例

cpp 复制代码
// 定义一个浮点型张量变量
Var x("x", TensorType({1, 3}, DataType::Float(32)));

3. Const(常量)

作用

  • 表示 不可变的数据(如模型权重、超参数)。
  • 在计算图中作为叶子节点存在。

关键组成

字段 说明
data 存储的常量值(runtime::NDArrayrelay::ConstantNode)。
checked_type_ 常量的类型(通常从 data 自动推断)。

示例

cpp 复制代码
// 定义一个常量张量
NDArray weight = NDArray::Empty({3, 3}, DataType::Float(32), {kDLCPU, 0});
Const weight_const(weight);

4. Call(函数调用)

作用

  • 表示对 算子(Operator)函数(Function) 的调用。
  • 是构建计算图的核心节点(如 addconv2d)。

关键组成

字段 说明
op 调用的目标(OpFunctionGlobalVar)。
args 参数列表(Array<Expr>,可以是 VarConst 或其他 Call)。
attrs 调用的属性(如卷积的 stridespadding)。

示例

cpp 复制代码
// 调用加法算子
Expr a = Var("a", TensorType({1}, DataType::Float(32)));
Expr b = Var("b", TensorType({1}, DataType::Float(32)));
Expr add_call = Call(Op::Get("add"), {a, b});

5. Function(函数定义)

作用

  • 封装可复用的计算单元(类似 Lambda 表达式)。
  • 用于表示模型中的子图或复合算子(如 conv2d + relu 融合)。

关键组成

字段 说明
params 输入参数列表(Array<Var>)。
body 函数体的表达式(Expr)。
ret_type 返回值的类型(如 TensorType)。
type_params 泛型类型参数(支持多态,类似 C++ 模板)。

示例

cpp 复制代码
// 定义一个简单的加法函数
Var x("x", TensorType({1}, DataType::Float(32)));
Var y("y", TensorType({1}, DataType::Float(32)));
Expr body = Call(Op::Get("add"), {x, y});
Function add_func({x, y}, body, TensorType({1}, DataType::Float(32)));

6. 节点间的协作关系

计算图示例

复制代码
z = (x + y) * 2

对应的 Relay IR 结构:

  1. 变量xyVar 节点)。
  2. 常量2Const 节点)。
  3. 调用add(x, y)multiply(add_result, 2)Call 节点)。
  4. 函数 :封装整个计算(Function 节点)。

代码实现

cpp 复制代码
Var x("x", TensorType({1}, DataType::Float(32)));
Var y("y", TensorType({1}, DataType::Float(32)));
Expr add = Call(Op::Get("add"), {x, y});
Expr two = Const(NDArray::FromVector({2.0f}));
Expr mul = Call(Op::Get("multiply"), {add, two});
Function func({x, y}, mul, TensorType({1}, DataType::Float(32)));

7. 类型系统支持

所有 Expr 节点都关联类型信息:

  • Var/Const :通过 type_annotationchecked_type_ 指定张量类型。
  • Call :根据算子的类型规则推断返回类型(如 add(Tensor, Tensor) -> Tensor)。
  • Function :通过 ret_type 声明返回值类型。

总结

节点 角色 关键特性
Expr 所有节点的基类 提供类型检查和遍历接口。
Var 输入/中间变量 符号化表示,支持类型注解。
Const 常量数据 存储不可变值(如权重)。
Call 算子或函数调用 构建计算图的核心节点,依赖 opargs
Function 可复用的计算单元 封装参数、计算体和返回类型,支持多态。

这些基础节点共同构成了 Relay IR 的 静态计算图,通过组合它们可以表示复杂的深度学习模型,并为后续优化和代码生成提供基础。

相关推荐
巴里巴气1 小时前
安装GPU版本的Pytorch
人工智能·pytorch·python
「、皓子~1 小时前
后台管理系统的诞生 - 利用AI 1天完成整个后台管理系统的微服务后端+前端
前端·人工智能·微服务·小程序·go·ai编程·ai写作
说私域2 小时前
基于开源AI智能名片链动2+1模式S2B2C商城小程序的抖音渠道力拓展与多渠道利润增长研究
人工智能·小程序·开源
笑衬人心。2 小时前
初学Spring AI 笔记
人工智能·笔记·spring
luofeiju2 小时前
RGB下的色彩变换:用线性代数解构色彩世界
图像处理·人工智能·opencv·线性代数
测试者家园2 小时前
基于DeepSeek和crewAI构建测试用例脚本生成器
人工智能·python·测试用例·智能体·智能化测试·crewai
张较瘦_2 小时前
[论文阅读] 人工智能 + 软件工程 | Call Me Maybe:用图神经网络增强JavaScript调用图构建
论文阅读·人工智能·软件工程
大模型真好玩2 小时前
准确率飙升!Graph RAG如何利用知识图谱提升RAG答案质量(四)——微软GraphRAG代码实战
人工智能·python·mcp
Baihai_IDP2 小时前
vec2text 技术已开源!一定条件下,文本嵌入向量可“近乎完美地”还原
人工智能·面试·llm
江太翁2 小时前
Pytorch torch
人工智能·pytorch·python