在深入 OptimizeImpl 阶段,也就是 Pass 优化之前,我们先了解一下 Relay 阶段的一些基本概念
OptimizeImpl 阶段最主要的工作就是图优化,图其实就是一种高级的表示,tvm 和图相关的概念好好几个:
- relay ir:这是 tvm 最 high level 的表示,另外 relay ir -> tir 的过程中,又会依赖 topi 和 te 这2种特定抽象的中间表示
- topi:TVM Operator Inventory,TOPI 提供了比 tir 具有更高抽象的 numpy 风格的,通用操作和调度
-
- te:Tensor Expression,张量表达式
- tir:最接近目标代码的中间表示
- ir:relay ir 和 tir 的一些公共基础结构,和上述2种ir不一样,并不是一个完整独立的抽象
本章我们先来了解下 IR 这个 relay ir 和 tir 最公共的基础设施,后续会依次介绍 relay ir、tir、topi、te
代码目录:
- 代码:src/ir
- 头文件:include/tvm/ir
编程语言最基本的核心概念就3个:类型、运算符、表达式,在 IR 这里分别对应 Type, OP, Expr
1.1 Type
Type 相关的定义都在 include/tvm/ir/type.h ,Type 包括基础的整型/浮点型等,也包括函数类型等相对复杂的类型。
这里我们介绍2种基本的类型:
- PrimType:最原始的 Type,可以直接映射到 low-level IR 的基本数据类型
- FuncType:函数类型
PrimType 可以在这上面做一些 Low-level 优化
定义如下:
class PrimTypeNode : public TypeNode { public: /*! * \brief The corresponding dtype field. */ runtime::DataType dtype; }
可以看到 PrimType 就一个数据成员,runtime::DataType,这个是 runtime 最底层的概念,代码在 include/tvm/runtime/data_type.h
/*! * \brief Runtime primitive data type. * * This class is a thin wrapper of DLDataType. * We also make use of DataType in compiler to store quick hint */ class DataType { public: /*! * \brief Type code for the DataType. * * DLPack consistency: * 1) kInt is consistent with kDLInt * 2) kUInt is consistent with kDLUInt * 3) kFloat is consistent with kDLFloat */ enum TypeCode { kInt = kDLInt, kUInt = kDLUInt, kFloat = kDLFloat, kHandle = TVMArgTypeCode::kTVMOpaqueHandle, kBFloat = kDLBfloat, kCustomBegin = 129 }; /*! \brief default constructor */ DataType() { data_ = DataType::Void(); } /*! * \brief Constructor * \param dtype The DLDataType */ explicit DataType(DLDataType dtype) : data_(dtype) {}