TVM Relay源码深度解读

---

一 、从Constant看Relay表达式的设计哲学

  在TVM的Relay IR中，即使是看似简单的常量表达式relay.const(1)，其背后也隐藏着整个类型系统的精妙设计。让我们从include/tvm/relay/expr.h中的Constant类入手，逐步拆解…"

1. 类定义概述

类名	继承关系	角色	关键特性
ConstantNode	public ExprNode	常量表达式的实际数据存储	包含常量数据（NDArray）、类型信息，并实现属性访问、哈希和相等比较逻辑。
Constant	public RelayExpr	常量表达式的智能指针封装	提供用户友好的构造函数和访问方法，隐藏内存管理细节。

---

2. ConstantNode 详解

class ConstantNode : public ExprNode {public:/*! \brief The data of the tensor */runtime::NDArray data;/*! \return The corresponding tensor type of the data */TensorType tensor_type() const;/*! \return Whether it is scalar(rank-0 tensor) */bool is_scalar() const { return data->ndim == 0; }void VisitAttrs(tvm::AttrVisitor* v) {v->Visit("data", &data);v->Visit("span", &span);v->Visit("mdata", &mdata);v->Visit("_checked_type_", &checked_type_);}bool SEqualReduce(const ConstantNode* other, SEqualReducer equal) const {return equal(data, other->data);}void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(data); }static constexpr const char* _type_key = "relay.Constant";TVM_DECLARE_FINAL_OBJECT_INFO(ConstantNode, ExprNode);
};

1. 核心成员

• data (runtime::NDArray)

  • 存储常量张量的实际数据（如权重、偏置等），TVM 使用 NDArray 统一表示多维数组。
  • 示例：卷积层的权重矩阵会被存储在这里。

• tensor_type()

  • 根据 data 的维度（shape）和数据类型（dtype）自动生成对应的 TensorType。
  • 用途：类型推断时确定常量的类型。

• is_scalar()

  • 判断常量是否为标量（0维张量），如 data->ndim == 0。

2. 关键方法

• VisitAttrs

  • 实现属性的序列化/反序列化，支持以下字段：

    v->Visit("data", &data);          // 张量数据
    v->Visit("span", &span);         // 源码位置信息
    v->Visit("mdata", &mdata);       // 元数据（如调试信息）
    v->Visit("_checked_type_", &checked_type_);  // 类型检查后的类型
    

• SEqualReduce 和 SHashReduce

  • 结构化相等比较：比较两个 ConstantNode 的 data 是否相同（用于优化中的常量折叠）。
  • 哈希计算：基于 data 生成哈希值（用于快速查找重复常量）。

3. 类型系统注册

TVM_DECLARE_FINAL_OBJECT_INFO(ConstantNode, ExprNode);

• _type_key = "relay.Constant"：唯一标识常量节点类型。
• FINAL：禁止继承，确保常量节点的行为不可被修改。

---

3. Constant 详解

class Constant : public Expr {public:/*!* \brief The constructor* \param data The data of the constant tensor.* \param span The source span of the expression.*/TVM_DLL explicit Constant(runtime::NDArray data, Span span = Span(), MetaData mdata = MetaData());TVM_DEFINE_OBJECT_REF_METHODS(Constant, RelayExpr, ConstantNode);
};

1. 核心功能

• 构造函数

  explicit Constant(runtime::NDArray data, Span span = Span(), MetaData mdata = MetaData());
  

  • 接收 NDArray 数据，构造一个常量表达式。
  • 示例：

    # Python 前端等价代码
    data = np.array([1, 2, 3], dtype="float32")
    const_expr = relay.Constant(tvm.nd.array(data))
    

• 智能指针方法

  TVM_DEFINE_OBJECT_REF_METHODS(Constant, RelayExpr, ConstantNode);
  

  展开后提供：

  • operator->()：直接访问 ConstantNode 成员（如 const_expr->data）。
  • get()：获取底层 ConstantNode 指针。
  • 自动内存管理（通过 ObjectRef 的引用计数）。

---

二. 核心内容概述

  在TVM源码中，include/tvm/relay/expr.h 是 Relay IR（中间表示）的核心头文件，定义了所有Relay表达式的基础数据结构和类型系统。它是实现TVM高层计算图表示的关键组成部分。以下是该文件的详细解析：
相关重要文件

文件路径	关联内容
include/tvm/relay/type.h	类型系统（TensorType等）
include/tvm/relay/op.h	运算符定义
include/tvm/relay/adt.h	代数数据类型支持
src/relay/ir/expr.cc	表达式方法的实现

---

include/tvm/relay/expr.h文件主要包含：

• (1) Relay表达式基类（RelayExpr/RelayExprNode）
• (2) 所有具体表达式类型的声明（如变量、常量、函数调用等）
• (3) 表达式类型的遍历和转换接口
• (4) 类型系统和属性访问的支持

---

(1) Relay表达式基类

class RelayExprNode : public BaseExprNode { /*...*/ };
class RelayExpr : public BaseExpr { /*...*/ };

• 角色：所有Relay表达式的公共基类
• 功能：
  • 提供类型系统支持（通过checked_type_字段）
  • 实现属性访问（VisitAttrs）
  • 支持结构化相等比较（SEqualReduce）

1. RelayExprNode 和 RelayExpr 的区别与用法

  RelayExprNode 是 Relay 表达式的实际实现类，是一个 C++ 类，包含了表达式的所有数据和功能实现。它是所有 Relay 表达式类型的基类。
  RelayExpr 是一个智能指针（relay::Expr），它指向 RelayExprNode 或其子类的实例。它提供了对 RelayExprNode 的安全访问和管理。

2. 主要区别

特性	RelayExprNode	RelayExpr
类型	C++ 类	智能指针（std::shared_ptr 的封装）
生命周期管理	需要手动管理	自动管理
使用方式	通常不直接使用，作为实现细节	用户主要交互的接口
继承关系	作为基类定义表达式结构	作为访问接口

3. 使用模式

在 TVM 中，通常的模式是：

1. 定义一个继承自 RelayExprNode 的具体表达式节点类
2. 使用 RelayExpr 作为这些节点的引用

例子1：常量表达式

// 创建一个常量表达式
auto const_node = relay::ConstantNode::make(tvm::runtime::NDArray::Zeros(...));
RelayExpr const_expr = const_node;// 通常更简洁的写法
RelayExpr const_expr = relay::Constant(tvm::runtime::NDArray::Zeros(...));

例子2：变量表达式

// 创建一个变量表达式
auto var_node = relay::VarNode::make("x", relay::Type());
RelayExpr var_expr = var_node;// 或者更简洁地
RelayExpr var_expr = relay::Var("x", relay::Type());

例子3：函数应用

// 创建函数应用表达式
RelayExpr func = ...; // 某个函数
RelayExpr arg = ...;  // 某个参数
auto call_node = relay::CallNode::make(func, {arg});
RelayExpr call_expr = call_node;// 或者
RelayExpr call_expr = relay::Call(func, {arg});

4. 实际使用建议

1. 用户代码：在大多数情况下，你应该使用 RelayExpr 而不是直接操作 RelayExprNode。

2. 扩展 Relay：如果你想定义新的表达式类型，需要继承 RelayExprNode 并实现相应接口。

3. 类型转换：可以使用 as<T> 方法将 RelayExpr 向下转换为特定类型的节点指针：

RelayExpr expr = ...;
if (const auto* call = expr.as<CallNode>()) {// 现在可以访问 CallNode 的特定成员call->op;call->args;
}

4. 创建新表达式：TVM 提供了辅助函数来创建表达式，通常以节点类型名去掉 “Node” 命名（如 relay::Var() 创建 VarNode 的 RelayExpr）。

这种分离设计使得 Relay IR 既灵活又安全，同时保持了良好的性能特性

(2) 具体表达式类型

表达式类型	说明	关键成员/方法
VarNode	变量（输入/中间结果）	String name_hint, Type type_annotation, Id vid
ConstantNode	常量张量（如模型权重）	runtime::NDArray data, tensor_type(), is_scalar()
CallNode	函数/运算符调用	Expr op, Array<Expr> args, Attrs attrs, Array<Type> type_args
LetNode	Let绑定（实现变量作用域）	Var var, Expr value, Expr body
TupleNode	元组结构（多返回值）	Array<Expr> fields
TupleGetItemNode	从元组中获取元素	Expr tuple, int index
IfNode	条件表达式	Expr cond, Expr true_branch, Expr false_branch
OpNode	基本运算符（如add/concat）	通过Op::Get("op_name")获取
FunctionNode	函数定义（在function.h中声明，但属于表达式）	Array<Var> params, Expr body, Type ret_type, Array<TypeVar> type_params
RefCreateNode	创建可变引用（用于状态更新）	Expr value
RefReadNode	读取引用值	Expr ref
RefWriteNode	更新引用值	Expr ref, Expr value
ConstructorNode	代数数据类型（ADT）的构造器（在adt.h中声明）	String tag, Array<Type> inputs
MatchNode	模式匹配（ADT处理）	Expr data, Array<Clause> clauses
TempExprNode	临时表达式（用于优化过程中的中间表示）	通常作为优化Pass的中间载体
GlobalVarNode	全局函数引用（跨模块调用）	String name_hint
SeqExprNode	顺序执行多个表达式（类似语句块）	Array<Binding> bindings, Expr body

---

1. 表达式类型 VarNode举例子

include/tvm/relay/expr.h

class Var;
/*! \brief Container for Var */
class VarNode : public ExprNode {public:/*!* \brief The unique identifier of the Var.** vid will be preserved for the same Var during type inference* and other rewritings, while the VarNode might be recreated* to attach additional information.* This property can be used to keep track of parameter Var* information across passes.*/Id vid;/*!* \brief type annotaion of the variable.* This field records user provided type annotation of the Var.* This field is optional and can be None.*/Type type_annotation;/*! \return The name hint of the variable */const String& name_hint() const { return vid->name_hint; }void VisitAttrs(tvm::AttrVisitor* v) {v->Visit("vid", &vid);v->Visit("type_annotation", &type_annotation);v->Visit("span", &span);v->Visit("mdata", &mdata);v->Visit("_checked_type_", &checked_type_);}bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {return equal(type_annotation, other->type_annotation) && equal.FreeVarEqualImpl(this, other);}void SHashReduce(SHashReducer hash_reduce) const {hash_reduce(type_annotation);hash_reduce.FreeVarHashImpl(this);}static constexpr const char* _type_key = "relay.Var";TVM_DECLARE_FINAL_OBJECT_INFO(VarNode, ExprNode);
};class Var : public Expr {public:/*!* \brief The constructor* \param name_hint The name hint of a variable.* \param type_annotation The type annotation of a variable.* \param span The source span of the expression.*/TVM_DLL Var(String name_hint, Type type_annotation, Span span = Span(), MetaData mdata = MetaData()): Var(Id(name_hint), type_annotation, span, mdata) {}/*!* \brief The constructor* \param vid The unique id of a variable.* \param type_annotation The type annotation of a variable.* \param span The source span of the expression.*/TVM_DLL Var(Id vid, Type type_annotation, Span span = Span(), MetaData mdata = MetaData());TVM_DEFINE_OBJECT_REF_METHODS(Var, RelayExpr, VarNode);
};

1. 核心设计理念

VarNode 和 Var 共同实现了 Relay IR 的变量系统，采用 TVM 标准的 Object-ObjectRef 设计模式：

• VarNode：存储实际数据的节点类（继承自 ExprNode）
• Var：管理 VarNode 的智能指针包装类（继承自 Expr）

---

2. 关键成员解析

(1) 核心字段

成员	类型	作用
vid	Id	唯一标识符，跨 Pass 保持不变（即使节点被重建）
type_annotation	Type	用户显式指定的类型注解（可空）
name_hint()	String	通过 vid->name_hint 获取的可读名称（非唯一）
span	Span	源码位置信息（用于错误定位）
mdata	MetaData	扩展元数据

(2) 特殊方法

方法	功能
SEqualReduce	结构化相等比较（用于优化 Pass 的重复检测）
SHashReduce	哈希计算（支持快速查找）
VisitAttrs	属性序列化/反序列化

---

3. 变量标识系统

(1) vid (Unique ID)

class IdNode : public Object {public:String name_hint;// ... 其他元数据
};

• 核心特性：
  • 通过 Id(name_hint) 构造，但系统会保证其唯一性
  • 即使优化 Pass 重建变量节点，vid 保持不变
  • 用于跨 Pass 跟踪参数变量（如梯度更新时识别同一参数）

(2) name_hint 与 vid 的关系

x = relay.var("input", shape=(1,3))  # 实际创建：# vid = Id("input_0x7f") (自动去重)# name_hint = "input" (用户友好)

---

4. 类型系统整合

(1) 类型注解流程

graph TDA[用户构造] -->|relay.var(..., dtype="float32")| B(type_annotation)B --> C[类型检查]C -->|更新| D(_checked_type_)

(2) 类型推导规则

• 若 type_annotation 存在：必须与实际使用类型兼容
• 若为空：从上下文推断类型

---

5. 内存模型与跨语言交互

(1) C++ 层构造

// 方式1：通过 name_hint
Var x("data", TensorType({1,3}, DataType::Float(32)));// 方式2：直接指定 Id
Var x(Id("data_0x7f"), TensorType({1,3}, DataType::Float(32)));

(2) Python 绑定

# Python 前端接口
x = relay.var(name="input",shape=(1,3),dtype="float32",span=SourceSpan(...)
)

(3) 对象生命周期

sequenceDiagramPython->>C++: relay.var() 创建请求C++->>Heap: 分配 VarNodeC++->>Python: 返回 Var(ObjectRef)Python->>C++: 析构时触发引用计数-1

---

6. 关键应用场景

(1) 函数参数定义

def build_linear():x = relay.var("x", shape=(1,3))w = relay.var("w", shape=(3,2))b = relay.var("b", shape=(2,))y = relay.add(relay.matmul(x, w), b)return relay.Function([x, w, b], y)

(2) 优化 Pass 中的变量处理

// 在 ConstantFolding 中识别变量引用
if (const VarNode* var = expr.as<VarNode>()) {if (var_map.count(var->vid)) {// 替换为已知常量}
}

(3) 类型检查

// 检查变量类型是否匹配
bool CheckType(const VarNode* var, const Type& expected) {return var->checked_type().as<TensorType>()->dtype == expected;
}

---

7. 设计亮点总结

1. 稳定性：vid 保证变量在优化过程中的持久标识
2. 灵活性：type_annotation 支持显式/隐式类型指定
3. 安全性：TVM_DECLARE_FINAL_OBJECT_INFO 防止错误继承
4. 可调试性：span 和 name_hint 增强错误可读性
5. 性能：SEqualReduce/SHashReduce 优化图操作效率

---

8. 典型问题分析

Q: 为什么需要同时存在 vid 和 name_hint？
A: 分工不同：

• name_hint：面向用户，提供可读性（允许重复）
• vid：面向系统，保证唯一性和跨Pass一致性

Q: 何时会重建 VarNode？
A: 典型场景：

• 类型推断后附加 _checked_type_
• 优化 Pass 中克隆表达式时保留原 vid 但新建节点

(3) TVM_DECLARE_BASE_OBJECT_INFO 宏详解

  这个宏是 TVM 类型系统的核心，用于在 C++ 中动态注册和管理对象的类型信息。它的核心作用是： 为每个类自动生成类型注册代码，使其能被 TVM 运行时识别和操作。

---

1. 宏的参数

#define TVM_DECLARE_BASE_OBJECT_INFO(TypeName, ParentType)

• TypeName：当前类名（如 ConstantNode）
• ParentType：父类名（如 ExprNode）

---

2. 静态断言检查（防止非法继承）

static_assert(!ParentType::_type_final, "ParentObj marked as final");

• 作用：如果父类被标记为 final（通过 _type_final），则禁止子类继承。

2. 运行时类型索引（RuntimeTypeIndex）

static uint32_t RuntimeTypeIndex() {// 检查子类槽位配置是否合法static_assert(TypeName::_type_child_slots == 0 || ParentType::_type_child_slots == 0 ||TypeName::_type_child_slots < ParentType::_type_child_slots,"子类槽位数不能超过父类限制");// 如果已预分配类型ID，直接返回if (TypeName::_type_index != ::tvm::runtime::TypeIndex::kDynamic) {return TypeName::_type_index;}// 否则动态分配return _GetOrAllocRuntimeTypeIndex();
}

• 功能：返回类的唯一类型 ID（uint32_t）。
• 优化：优先使用预分配的 _type_index（性能更高），否则动态分配。

3. 动态分配类型索引（_GetOrAllocRuntimeTypeIndex）

static uint32_t _GetOrAllocRuntimeTypeIndex() {static uint32_t tidx = Object::GetOrAllocRuntimeTypeIndex(TypeName::_type_key,         // 类型名称字符串（如 "relay.Constant"）TypeName::_type_index,       // 预分配的类型IDParentType::RuntimeTypeIndex(), // 父类类型IDTypeName::_type_child_slots, // 为子类预留的槽位数TypeName::_type_child_slots_can_overflow // 是否允许超额);return tidx;
}

• 作用：向 TVM 运行时注册类型，并分配唯一 ID。
• 关键参数：
  • _type_child_slots：限制子类数量（防止类型爆炸）。
  • _type_child_slots_can_overflow：为 true 时允许突破限制。

---

通俗版解释：TVM的类型身份证系统

你可以把TVM的类型系统想象成一个学校的学生管理系统，而TVM_DECLARE_BASE_OBJECT_INFO就是给学生（类）办身份证的机器：

---

1. 为什么要办身份证？

• 每个学生（类）需要唯一学号（类型ID）
• 需要知道他的班主任是谁（父类）
• 防止有人冒充转校生（非法继承）

2. 办证过程（宏的作用）

// 给"小明同学"办证，班主任是"李老师"
TVM_DECLARE_BASE_OBJECT_INFO(小明, 李老师)

这个宏会自动做三件事：

1. 检查家世清白

   static_assert(!李老师::是final班, "班主任明确不收新学生！");
   

   • 如果班主任声明"我们班不接收转学生"，就报错

2. 分配学号

   • 优先用预留的VIP学号（_type_index）
   • 没有就现场摇号（_GetOrAllocRuntimeTypeIndex）

3. 登记亲属关系

   学号 = 教务处.登记(姓名："小明",班主任：李老师.学号,可带小弟人数：3  // _type_child_slots
   );
   

3. 特殊班级（FINAL版）

TVM_DECLARE_FINAL_OBJECT_INFO(学霸班, 实验班)

• 相当于在班级门口挂**“禁止转入”**牌子
• 其他班同学想转学过来会直接报错

4. 实际有什么用？

• 查身份证快：obj->IsInstance<小明>() 比查户口本快
• 安全转班：obj.as<小明>() 能安全转换类型
• 防止冒名顶替：禁止随便认爹（错误继承）

---

举个栗子🌰

# Python前端定义一个"汉堡店"类
@register_relay_node("food.HamburgerShop")
class HamburgerShopNode(ExprNode):_type_key = "food.HamburgerShop"_type_child_slots = 2  # 允许开2家分店

C++层通过这个宏：

1. 给汉堡店分配类型ID（比如9527）
2. 记录它的父类是ExprNode
3. 允许最多2个子类（比如CheeseBurgerShop、ChickenBurgerShop）

---

一句话总结

这个宏就是TVM给类发身份证+建家族档案的工具，让系统能：

• ✅ 快速识别"你是谁"（类型检查）
• ✅ 知道"你爸是谁"（继承关系）
• ❌ 防止"乱认亲戚"（非法继承）

---

(4) 遍历接口

  void VisitAttrs(tvm::AttrVisitor* v) {v->Visit("data", &data);v->Visit("span", &span);v->Visit("mdata", &mdata);v->Visit("_checked_type_", &checked_type_);}

  VisitAttrs 是 TVM 中用于统一序列化、反序列化和属性访问的核心接口。以下是 ConstantNode 使用该函数的具体示例，涵盖 C++ 和 Python 场景：

---

1. C++ 场景示例

(1) 模型序列化（保存为JSON）

// 创建常量节点
runtime::NDArray arr = runtime::NDArray::Empty({2, 2}, DLDataType{kDLFloat, 32, 1}, DLContext{kDLCPU, 0});
ConstantNode* const_node = new ConstantNode();
const_node->data = arr;// 序列化为JSON
JSONAttrVisitor visitor;
const_node->VisitAttrs(&visitor);  // 触发以下调用：// visitor.Visit("data", &data)// visitor.Visit("span", &span)...
std::string json = visitor.GetJSON();

输出JSON片段：

{"type_key": "relay.Constant","data": {"b64": "AABAA...", "dtype": "float32", "shape": [2, 2]},"span": null,"_checked_type_": "TensorType([2,2], float32)"
}

(2) 优化Pass中的常量修改

class ConstantMutator : public AttrMutator {public:void VisitAttrs(AttrVisitor* v) override {if (v->IsMutator()) {  // 检查是否为修改模式runtime::NDArray new_data = ...; // 生成新数据v->Visit("data", &new_data);    // 修改data字段}}
};// 调用示例：
ConstantMutator mutator;
const_node->VisitAttrs(&mutator);  // 修改常量数据

(3) 调试打印

class DebugPrinter : public AttrVisitor {public:void Visit(const char* key, runtime::NDArray* data) override {std::cout << key << ": shape=" << data.Shape();}
};DebugPrinter printer;
const_node->VisitAttrs(&printer);  // 输出：data: shape=[2,2]

---

2. Python 场景示例

(1) 直接属性访问

import tvm
from tvm import relay# 创建常量
data = tvm.nd.array(np.zeros((2,2), dtype="float32"))
const = relay.Constant(data)# Python属性访问（背后调用VisitAttrs）
print(const.data)      # 访问NDArray → 触发Visit("data", &data)
print(const.span)      # 访问源码位置 → Visit("span", &span)

输出：

<tvm.nd.NDArray shape=(2, 2), dtype=float32>
None  # 未设置span时的默认值

(2) 模型保存与加载

# 保存模型（触发序列化）
mod = tvm.IRModule.from_expr(const)
mod.save("const.json")  # 内部调用VisitAttrs# 加载模型（触发反序列化）
loaded_mod = tvm.ir.load_json("const.json")
loaded_const = loaded_mod["main"].body
assert isinstance(loaded_const, relay.Constant)

(3) 自定义属性访问器

class MyVisitor(tvm.ir.AttrVisitor):def visit(self, name, value):print(f"Attribute {name} has type {type(value)}")visitor = MyVisitor()
const.visit_attrs(visitor)  # 显式调用VisitAttrs

输出：

Attribute data has type <class 'tvm.runtime.ndarray.NDArray'>
Attribute span has type <class 'tvm.ir.Span'>
...

---

class Constant;
/*!* \brief Constant tensor type.*/
class ConstantNode : public ExprNode {public:/*! \brief The data of the tensor */runtime::NDArray data;/*! \return The corresponding tensor type of the data */TensorType tensor_type() const;/*! \return Whether it is scalar(rank-0 tensor) */bool is_scalar() const { return data->ndim == 0; }void VisitAttrs(tvm::AttrVisitor* v) {v->Visit("data", &data);v->Visit("span", &span);v->Visit("mdata", &mdata);v->Visit("_checked_type_", &checked_type_);}bool SEqualReduce(const ConstantNode* other, SEqualReducer equal) const {return equal(data, other->data);}void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(data); }static constexpr const char* _type_key = "relay.Constant";TVM_DECLARE_FINAL_OBJECT_INFO(ConstantNode, ExprNode);
};class Constant : public Expr {public:/*!* \brief The constructor* \param data The data of the constant tensor.* \param span The source span of the expression.*/TVM_DLL explicit Constant(runtime::NDArray data, Span span = Span(), MetaData mdata = MetaData());TVM_DEFINE_OBJECT_REF_METHODS(Constant, RelayExpr, ConstantNode);
};

以下是关于 ConstantNode 和 Constant 类的详细解释与概括，结合它们在 TVM Relay IR 中的作用和实现设计：

---

---