编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载Relay 是 Apache TVM 中负责表达深度学习模型的计算图 IR其类型系统决定了编译器能在编译期对张量形状、数据类型和数据结构做多深入的静态推理。本篇基于 docs/reference/langref/relay_type.rst 系统讲解 Relay 的完整类型语言从TensorType、TupleType、TypeVar与FuncType等基础类型到基于type relations类型关系的约束求解式类型推断再到支持递归定义的代数数据类型ADT。读完你将掌握 Relay 类型节点的语义、类型推断的工作原理以及类型系统如何支撑fuse_ops等编译优化 pass 的实现。为什么 Relay 需要一套独立的类型系统Relay 是一门静态类型、支持类型推断的语言程序可以只写少量显式类型注解编译器即可为整个程序推导出完整类型。静态类型对编译器优化至关重要因为类型本身携带了程序所操作数据的属性——运行时形状shape、数据布局layout与存储方式storage这些信息不需要真正运行程序就能获得。更关键的是Relay 的类型系统实现了针对形状的依赖类型dependent typing它把张量的形状当作类型的一部分来跟踪。把“形状推断”转化为“类型推断”问题之后Relay 就能在编译期静态推理那些输出形状随输入形状复杂变化的算子包括包含分支branching与函数调用function calls的程序从而对所有张量在编译期完成形状推断。这种静态形状推理带来了两个直接收益支持 ahead-of-time 编译程序在运行前即可确定所有张量的具体形状为后续优化提供更充分的张量信息优化以passRelay-to-Relay 的 AST 变换形式实现可以读取推断出的类型例如形状信息来决定如何变换程序。文档中给出的典型实例是算子融合 passsrc/relay/transforms/fuse_ops.cc 就利用推断出的张量形状把 Relay 程序中分散的算子调用替换为融合后的算子实现。而这一切推理能力的落点正是下文要展开的类型关系type relationsRelay 中大部分类型检查本质上是一个**约束求解constraint solving**过程——确保所有类型关系在调用点call sites得到满足。类型关系用相对简单灵活的方式把依赖类型的强大能力引入 Relay同时又不显著增加类型系统的复杂度。Type所有类型的基类Type是 Relay 中所有类型的基类一切 Relay 类型都是它的子类。在 include/tvm/ir/type.h 中可以看到其 C 定义TypeNode直接继承自 TVM 的Object带有一个用于记录源码位置的span字段并声明了结构相等SEqualReduce与结构哈希SHashReduce能力。Python 侧通过tvm.relay.ty暴露相关定义见 python/tvm/relay/ty.py。值得注意的一点是TVM 把 Relay 的类型系统作为整个软件栈统一使用的类型系统见 include/tvm/ir/type.h 的文件注释Relay 的Type与底层 TIR 中的runtime::DataType是两层不同的信息。DataType是粗粒度的、在低层PrimExpr构造时即时构建的类型例如知道某个表达式 dtype 是int32而Type提供更细粒度的信息例如一个表达式 dtype 是Handle()类型却是MemRef[float32]且通常由类型检查惰性构造。统一的Type成为跨 IR 方言的公共桥梁——所有函数都必须有类型签名从而支持跨方言函数调用。TensorType张量的形状与数据类型TensorType是 Relay 中最常用、最具体的张量类型。张量按照**数据类型dtype和形状shape**来标注。当前实现直接使用 TVM 的 DataType 与形状体系文档注明未来 Relay 可能引入独立的形状 AST。常见的数据类型包括bool、float32、int8以及各种位宽和通道数lanes的组合。形状是维度组成的元组TVMIndexExpr例如(5, 5)标量scalar也用元组类型表示其形状为()形状中还可以出现变量和含变量的算术表达式因此 Relay 的约束求解阶段会在运行程序前尝试为所有形状变量找到赋值确保所有形状具体化。例如一个 10×10 的 32 位浮点张量类型写作Tensor[(10, 10), float32]在 include/tvm/ir/tensor_type.h 中TensorTypeNode由shapeArrayPrimExpr元素可以是常量IntImm也可以是任意符号整数表达式从而支持泛化的形状推断和dtypeDataType两个字段构成并提供了Size()方法返回形状各维乘积。Python 侧还提供了两个便捷工具见 python/tvm/relay/ty.pyis_dynamic(tensor_type)检查类型是否含有Any或符号变量形状即是否为动态形状scalar_type(dtype)等价于TensorType((), dtype)用于快速构造标量类型。TupleType定长元组类型元组在 Relay 中是长度静态已知的值序列因此元组类型就是各成员类型的序列。由于元组类型大小静态已知元组投影tuple projection的类型就是对元组类型按下标索引。文档给出了如下示例%t的类型是(Tensor[(), bool], Tensor[(10, 10), float32])而%c %t.1的类型是Tensor[(10, 10), float32]let %t (False, Constant(1, (10, 10), float32)); let %c %t.1; %c从源码看TupleTypeNodeinclude/tvm/ir/type.h只有一个fieldsArrayType字段并提供TupleType::Empty()构造空元组。TVM 中还用空元组类型来表示void见VoidType()与IsVoidType()include/tvm/ir/type.h。TypeVar 与类型参数多态与 Kind类型参数type parameter是用于函数多态polymorphism的占位类型用TypeVar表示。类型参数按kind种类说明它允许替换哪些类型Type对应顶层 Relay 类型如张量类型、元组类型、函数类型BaseType对应张量的基础类型如float32、boolShape对应一个张量形状ShapeVar对应张量形状内部的变量。Relay 类型系统强制类型参数只能出现在其 kind 允许的位置。例如若类型变量t的 kind 是Type那么Tensor[t, float32]就不是合法类型。与普通参数一样调用点必须为类型参数提供具体实参。文档同时注明两点现状目前仅支持 kind 为Type的类型参数类型参数语法在文本格式text format中尚不支持。下面的例子中s是 kind 为Shape的类型参数在调用点被替换为(10, 10)def pluss : Shape(%t1 : Tensor[s, float32], %t2 : Tensor[s, float32]) { add(%t1, %t2) } plus(10, 10)(%a, %b)在 include/tvm/ir/type.h 中TypeKind枚举实际定义了六种 kindkType、kShapeVar、kBaseType、kConstraint、kAdtHandle、kTypeData并有TypeKind2String辅助函数。TypeVarNode由name_hint仅作提示不参与相等性比较与kind组成SEqualReduce时按 kind 与自由变量身份比较。Python 侧ShapeVar(name)就是TypeVar(name, kindTypeKind.ShapeVar)的便捷包装。FuncType函数的类型FuncType是 Relay 中分配给函数的类型。一个函数类型由四部分组成类型参数列表type_params类型约束集合type_constraints参数类型序列arg_types返回类型ret_type。非形式化地写作fntype_params(arg_types) - ret_type where type_constraints类型参数可以出现在参数类型或返回类型中每个类型约束必须在函数的每个调用点成立。约束通常以函数的参数类型和返回类型为参数也可能只取其中一部分。在 include/tvm/ir/type.h 中FuncTypeNode的四个字段与上述结构一一对应。注意普通函数非模板函数的type_params与type_constraints为空FuncType从结构上可以被看作 C 模板函数的类型级对应物。TypeRelation类型推断的核心机制类型关系type relation是 Relay 类型系统中最复杂的特性它允许用户用新规则扩展类型推断。广播算子broadcasting和flatten这类“以复杂方式依赖张量形状”的算子正是通过类型关系来定义类型的从而让 Relay 能够静态推理这些场景下的形状。关系的数学含义类型关系R描述了 Relay 函数输入与输出类型之间的一种关系R是一个作用于类型的函数关系成立输出true不成立输出false。交给关系的类型可能是不完整的或包含形状变量的因此类型推断必须为不完整类型和形状变量赋合适的值使得必要的关系成立如果这样的赋值存在。例如恒等关系可以定义为Identity(I, I) :- true通常的做法是为每个算子定义一个专属关系把参数类型与返回类型的全部必要约束编码进去。例如flatten的关系Flatten(Tensor(sh, bt), O) :- O Tensor(sh[0], prod(sh[1:]))有了Broadcast这样的关系就可以为add这类算子定型add : fnt1 : Type, t2 : Type, t3 : Type(t1, t2) - t3 where Broadcast上面where Broadcast的含义是参数类型与返回类型必须是张量且t3的形状是t1与t2形状广播broadcast的结果。只要满足Broadcast类型系统就接受任意参数类型与返回类型。关系在 C/Python 中的实际实现文档特别强调上述 Prolog 风格写法只是示意当前关系必须由用户在 C 或 Python 中实现。Relay 类型系统对类型关系使用一个ad hoc求解器关系被实现为检查关系是否成立的 C/Python 函数这些函数会命令式地更新形状变量或不完整类型。当前实现约定关系失败返回False关系成立、或信息不足无法判断时返回True。所有关系的函数按需运行只要某个输入被更新就重新触发直到满足以下三个条件之一所有关系成立且无剩余不完整类型——类型检查成功某个关系不成立——产生类型错误到达不动点仍剩余形状变量或不完整类型——要么是类型错误要么需要补充更多类型注解。目前 Relay 中使用的全部关系都在 C 中实现示例见 src/relay/op 目录下的各个算子实现文件。求解器与推断器的源码印证从源码可以完整看到这条推断链路。类型推断器实现在 src/relay/transforms/type_infer.cc其注释约 L94-L96概括了整体流程遍历表达式调用solver_.AddConstraint与solver_.Unify收集必要的约束调用solver_.Solve()求解约束用Resolver以求解得到的checked_type重建表达式。具体到算子的调用点推断器对每个算子调用都会执行solver_.AddConstraint(TypeRelation(rel-func, arg_types, arg_types.size() - 1, attrs))见 src/relay/transforms/type_infer.cc把算子自带的关系函数作为约束加入求解器遇到函数节点中已有的显式关系约束时同样以TypeRelation形式加入src/relay/transforms/type_infer.cc。所有约束收集完毕后调用Solve()最后Resolve()得到最终类型。而求解器本身实现在 src/relay/analysis/type_solver.cc 与 src/relay/analysis/type_solver.h。在 C 侧TypeRelation由四部分组成include/tvm/ir/type_relation.hfunc类型为TypeRelationFn的关系函数args传给关系函数的类型参数num_inputs输入数量attrs附加的算子属性。TypeRelationFn是TypedEnvFuncbool(const ArrayType args, int num_inputs, const Attrs attrs, const TypeReporter reporter)。关系函数通过TypeReporterinclude/tvm/ir/type_relation.h回报求解结果调用Assign(dst, src)创建类型相等约束dst是提示性赋值方向求解器也可能反向求解、调用Assert/AssertEQ对符号形状做断言任一条件含符号时就应使用断言返回false表示可证明断言失败。TypeRelation继承自TypeConstraint因为关系本质上就是函数上的一种约束。TypeConstraint 与 IncompleteTypeTypeConstraint类型约束是一个抽象类文档注明将在未来版本中进一步细化目前类型关系是唯一提供的约束形式。IncompleteType不完整类型是类型推断过程中“尚未知道”的类型或类型的一部分仅在类型推断期间使用任何省略的类型注解都会先被替换为一个不完整类型稍后会被另一个类型替换。编程语言文献中通常称其为“类型变量type variables”或“类型洞type holes”Relay 用“不完整类型”这一名称是为了与类型参数明确区分类型参数必须绑定到某个函数并在调用点用具体类型实参实例化不完整类型可以出现在程序中任意位置在类型推断期间被填充。从实现上看include/tvm/ir/type.h 的注释给出一个直观类比如果把类型关系看作“类型的计算图”那么IncompleteType是图中的中间值TypeVar是图的输入。IncompleteTypeNode只含一个kind字段并且用自由变量FreeVarEqualImpl方式进行相等比较这意味着不同的不完整类型互相区分但可以互相统一。代数数据类型ADT可组合、可递归的结构类型代数数据类型Algebraic Data TypesADT允许轻松灵活地组合类型构造出可以用归纳法推理、并能编写递归函数的数据结构。文档注明 ADT 目前同样不支持文本格式text format其完整概览见 docs/reference/langref/relay_adt.rst。从类型系统角度看ADT 有两个最关键的属性可以带类型参数构造器参数可以是类型参数且不同类型参数的 ADT 实例必须视为不同类型可以递归ADT 的构造器可以接收该 ADT 自身的实例因此树、列表这类结构可以被归纳地构建。GlobalTypeVarADT 的句柄为了紧凑地表示 ADT 并方便支持递归定义每个 ADT 定义都获得一个全局类型变量global type variable作为唯一句柄每次定义都会分配一个新的全局类型变量因此可以用指针相等性来区分不同的 ADT 名称。从类型系统角度看ADT 按名称区分两个 ADT 即便所有构造器结构完全相同只要句柄不同就被视为不同类型。递归的实现方式与全局函数递归完全一致构造器在其定义中直接引用 ADT 句柄全局类型变量。Python 侧对应tvm.relay.ty.GlobalTypeVar源码见 include/tvm/ir/type.h。TypeDataADT 定义类型数据除了名字ADT 还需要存储构造器以及构造器中用到的类型参数。这些信息存放在module中与全局函数定义的存放方式类似参见 docs/reference/langref/relay_expr.rst 中的 module 描述。在对 ADT 使用进行类型检查时类型系统有时需要按 ADT 名称在 module 中索引查询构造器信息。例如当某个构造器在match表达式子句中被模式匹配时类型检查器必须检查构造器签名以确保被绑定的变量被赋予正确的类型。C 实现中include/tvm/ir/adt.hTypeDataNode由三部分构成headerGlobalTypeVar即 ADT 的名字ADT 采用名义类型nominal typing——构造器相同但名字不同的 ADT 是不同类型type_vars类型变量支持多态constructors构造器数组。每个Constructorinclude/tvm/ir/adt.h又由name_hint、inputs输入类型、belong_to所属的GlobalTypeVar与tag注册时分配的构造器表索引组成。TypeCall类型级函数应用由于 ADT 定义可以带类型参数Relay 类型系统把 ADT 定义视为一个类型级函数type-level function定义接收类型参数返回一个带这些类型参数的 ADT 实例的类型。因此任何 ADT 实例都用**类型调用type call**来定型显式列出传给 ADT 定义的类型参数。显式列出类型参数非常重要因为用不同构造器构建但类型参数相同的两个 ADT 实例是相同类型而类型参数不同的两个实例不应被视为同一类型例如整数列表与浮点张量二元组列表就不是同一类型。类型调用中的“函数”是 ADT 句柄且每个类型参数必须有一个实参无类型参数的 ADT 定义其实例的类型调用也不传任何类型实参。源码见 include/tvm/ir/type_relation.hTypeCallNode由func类型级函数与args实参构成。文档还指出TypeRelation相比TypeCall更通用因为关系允许对输入和输出同时进行推断。实例内建的 List ADTRelay 内建了简单的 list ADT 作为默认 ADT用于演示上述全部概念。其定义如下data Lista { Nil : () - List Cons : (a, List[a]) - List }全局类型变量List是 ADT 的句柄module 中的类型数据记录了List接受一个类型参数并有两个构造器Nil签名fna() - List[a]与Cons签名fna(a, List[a]) - List[a]Cons中对List的递归引用正是通过在其构造器定义中使用全局类型变量List实现的。下面两个列表实例及各自类型用类型调用给出Cons(1, Cons(2, Nil())) # List[Tensor[(), int32]] Cons((1, 1), Cons((2, 2), Nil())) # List[(Tensor[(), int32], Tensor[(), int32])]注意Nil()可以是任意列表的实例因为它没有任何使用类型参数的参数尽管如此对任何特定的Nil()实例其类型参数仍必须被指明。以下两个列表会被类型系统拒绝因为类型参数不匹配# 试图把整数放进 int * int 元组的列表 Cons(1, Cons((1, 1), Nil())) # 试图把整数列表放进 int * int 元组列表的列表 Cons(Cons(1, Cons(2, Nil())), Cons(Cons((1, 1), Cons((2, 2), Nil())), Nil()))结语类型系统如何服务编译优化回顾整条链路Relay 用TensorType把张量形状纳入类型用TypeVar与FuncType表达多态函数签名用TypeRelation把形状传播规则编码为可求解的约束再由 src/relay/analysis/type_solver.cc 完成约束求解、填充IncompleteType最终由 src/relay/transforms/type_infer.cc 为每个表达式挂上完整类型。ADT 与GlobalTypeVar/TypeCall机制则为结构化数据列表、树提供可递归、可多态的类型描述。这些编译期获得的类型尤其是形状信息正是后续优化 pass 的决策依据——例如 src/relay/transforms/fuse_ops.cc 利用推断出的张量形状实现算子融合。理解 Relay 类型系统就理解了 TVM 编译器前端如何从“一段有类型的计算图”走向“可高效执行的优化代码”。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐Taichi 类型系统完全指南静态类型、原始类型、复合类型与类型转换实战Taichi 类型系统完全指南静态类型、原始类型、复合类型与类型转换实战 Taichi 是一门 静态类型 的嵌入式语言变量在 Taichi 作用域 ti编程语言编译器高性能计算Darklang类型系统静态类型检查与类型推导Darklang类型系统静态类型检查与类型推导 引言为什么需要强大的类型系统 在现代编程语言设计中类型系统Type System扮演着至关重要的角色Python类型系统完全指南gh_mirrors/ma/materials中的类型提示与静态类型检查Python类型系统完全指南gh_mirrors/ma/materials中的类型提示与静态类型检查 掌握Python类型系统是提升代码质量和开发效率的关键技示例工程教程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考