AI中pass架构设计优化
Relay 和 TVM IR，包含一系列优化passes，可提高模型的性能指标，例如平均推理，内存占用，或特定设备的功耗。有一套标准优化，及特定机器学习的优化，包括常量折叠，死代码消除，算子布局更改，算子融合，缓冲区处理和循环转换等。这些passes中的每一个都构造为一个 ir-to -ir 转换，使用在遍历期间和/或前收集的分析结果。
随着 TVM 的快速发展，对管理这些pass的更系统，更有效的方法的需求，变得越来越明显。此外，管理跨 TVM 堆栈不同层（例如 Relay 和 tir）pass的通用框架，为开发人员快速构建原型，将实现的pass插入系统，铺平了道路。
本节描述了基础架构设计，利用产品编译器，管理优化pass，及构建层的深度学习框架。
例如，许多现有的产品编译器，如 GCC 和 LLVM，都采用pass管理器，有效管理pass的执行。最初管理 pass 很简单，因为 pass 的数量很少，成熟的编译器，将包含数百个单独的 pass。通常，外部用户希望正确调度自定义pass，无需修改单个手工制作的pass顺序。
现代深度学习框架，如 Pytorch 和 MXNet Gluon，分别通过Sequential和Block，启用pass-style 层构建方案的趋势。有了这样的结构，这些现代框架能够方便将模块/层添加到容器中，轻松地构建神经网络。
Relay pass infra 的设计，很大程度上受到 LLVM 中，使用的分层pass管理器和流行的深度学习框架中，使用的块式容器的启发。pass基础架构的主要目标包括：
- 实现更好的优化编程调度。允许用户灵活地定制和构建优化pass。
- 提供一种用户友好的方式来调试优化pass。
- 减轻开发人员手动和分别解决pass之间的依赖关系。
- 为开发人员简化新pass的实施。例如，允许用户在 Python 中，实现一个 pass，让 pass infra 操纵执行。
设计
专注于为用户提供易于扩展的功能，让用户可以快速添加新pass，不会失去向后兼容性。该设计包含后端和前端。前者实现了 pass infra 的主要逻辑。后者为用户提供了简单的 API 进行交互，允许用户快速创建优化pass。
C++ 后端
提供一个PassInfo对象，包含pass所需的基本信息。name是传递名称，opt_level指示，在哪个优化级别启用pass， required表示执行特定pass，所需的pass（有关更多详细信息，参阅include/tvm/ir/transform.h）。例如，在注册pass期间，pass开发人员，可以指定pass的名称，将执行的优化级别和/或所需的pass。在用户提供的优化级别下运行时，是否需要执行某个 pass， opt_level可用于帮助 pass infra 识别。 required字段，可以被 pass infra 使用，解决 pass 依赖。
class PassInfoNode : public Object {
String name;
int opt_level;
Array<String> required;
};
传递上下文
PassContext携带用于优化pass的有用信息。例如，包含错误报告系统，可以提供有关优化失败原因的诊断。PassContext旨在替换旧的BuildConfig，用于帮助用户配置编译选项，包括优化级别和必需/禁用的pass等。例如，可能有一个配置， opt_level=3使用disabled_pass=xx提供的某些禁用的pass，执行所有PassContext。可以将所有pass，放在opt_level=3，排除禁用pass列表中的那些。PassContext提供了一种检测所有pass的方法。
PassContext包含优化pass的有用信息。例如，包含错误报告系统，可以提供有关优化失败原因的诊断。PassContext设计用于替换旧的BuildConfig，该配置用于帮助用户配置编译选项，包括优化级别和必需/禁用的pass等。例如，可能有一个配置，该配置使用PassContext提供的disabled_pass=xx，在opt_level=3执行所有pass，一些禁用的pass，使用disabled_pass=xx。现在，可以在opt_level=3时，对所有pass，进行全局排序，排除禁用pass列表中的pass。PassContext提供了一种对所有pass，进行检测的方法。
这个类是为用户设计的，使用语法编写Python，在特定配置下执行优化。用户可以通过PassContext::Current()，线程安全的方式，获得特定程序范围内，可用的上下文，线程本地存储PassContextThreadLocalStore，保存创建的pass context对象。将提供示例来说明，如何使用C++和Python API，使用pass context，创建编译pass。
class PassContextNode: public Object {
public:
int opt_level{2};
tvm::Array<tvm::Expr> required_pass;
tvm::Array<tvm::Expr> disabled_pass;
mutable Optional<DiagnosticContext> diag_ctx;
Map<String, ObjectRef> config;
Array<instrument::PassInstrument> instruments;
};
class PassContext : public NodeRef {
public:
TVM_DLL static PassContext Create();
TVM_DLL static PassContext Current();
TVM_DLL void InstrumentEnterPassContext();
TVM_DLL void InstrumentExitPassContext();
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
/* Other fields are omitted. */
private:
// The entry of a pass context scope.
TVM_DLL void EnterWithScope();
// The exit of a pass context scope.
TVM_DLL void ExitWithScope();
// Classes to get the Python `with` like syntax.
friend class tvm::With<PassContext>;
};
struct PassContextThreadLocalEntry {
/*! rief The default pass context. */
PassContext default_context;
/*! rief The current pass context. */
std::stack<PassContext> context_stack;
PassContextThreadLocalEntry() {
default_context = PassContext(make_node<PassContextNode>());
}
};
/*! rief The thread-local store to hold the pass context. */
typedef dmlc::ThreadLocalStore<PassContextThreadLocalEntry>
PassContextThreadLocalStore;
pass构建
pass infra以分层方式设计，可以在 Relay/tir 程序的，不同粒度下工作。PassNode引入了一个纯虚拟类，作为不同优化pass的基础。包含几个必须由子类在模块，函数或pass序列级别实现的虚拟方法。
class PassNode : Object {
virtual PassInfo Info() const = 0;
virtual Module operator()(const IRModule& mod
const PassContext& pass_ctx) const = 0;
};
函子显示必须如何实现pass，始终在 IRModule特定上下文下工作。所有pass都以ModuletoModule方式设计。由 pass infra 控制的优化，将始终更新整个模块。
已经创建了几个子类，实现不同类型的优化pass，例如，函数级pass，模块级pass和顺序pass。每个子类本身都可以充当pass管理器。例如，可以收集所需的pass执行，或基于给定的元数据构建，依赖关系图。完整定义可以在src/relay/ir/transform.cc和src/ir/transform.cc 中找到。
模块级pass
模块级pass主要用于全局和pass间优化 (IPO)，类似于 LLVM 中使用的模块pass。Relay 中一些典型的 pass，需要一个模块的全局图片，比如 A-normal form 转换和 lambda 提升等，都属于这个集合。在此级别，用户甚至可以在模块中，添加和/或删除功能。所有pass
class ModulePassNode : PassNode {
PassInfo pass_info;
runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
// Other members/methods are omitted
};
pass_info维护模块级pass所需的信息。pass_func勾勒出真正的优化。例如，可能需要对模块执行死代码消除。可以在pass_func中实现算法，在模块上运行。删除死代码，包括模块中未使用的函数。该字段被设计为一个打包函数，可以在 C++ 和 Python 中，实现优化。
函数级pass
函数级pass，用于为给定的 Relay/tir 模块，实现各种函数内级优化。一次从模块的函数列表中，获取一个函数，进行优化，生成重写的 Relay Function或 tir PrimFunc。大多数pass可以归入这一类，如Relay中，常见子表达式消除和推理简化，及tir中的矢量化和扁平化存储等。
此级别的pass范围是 Relay 函数，或 tir 原始函数。无法通过pass，添加或删除函数。
class FunctionPassNode : PassNode {
PassInfo pass_info;
runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
bool SkipFunction(const Function& func) const;
// Other members/methods are omitted...
};
pass_info与刚刚在模块pass中描述的相同。pass_func需要一个函数，进行优化，需要一个模块，可能会报告错误。一个函数可以用“SkipOptimization”注释，在优化pass中忽略。
连续passes
SequentialPass与 Pytorch 类似，nn.Sequential包含许多用于执行的pass。
class SequentialPassNode : PassNode {
PassInfo pass_info;
// Passes need to be executed.
Array<Pass> passes;
bool PassEnabled(const PassInfo& info) const;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
};
仅放置了在Relay中的少数pass。例如，FoldScaleAxis要求在内部调度ForwardFoldScaleAxis和BackwardFoldScaleAxis。建议首先完成BackwardFoldScaleAxis。该pass是SequentialPass的理想候选。
下面的代码显示了如何调用顺序过程中的各个pass。使用pass列表中，在一个顺序pass中，执行每个pass。
Module SequentialNode::operator()(const Module& module,
const PassContext& pass_ctx) const {
Module mod = module;
for (const Pass& pass : passes) {
ICHECK(pass.defined()) << "Found undefined pass for optimization.";
const PassInfo& pass_info = pass->Info();
if (!PassEnabled(pass_info)) continue;
for (const auto& it : pass_info->required) {
const auto* name = it.as<tvm::ir::StringImm>();
ICHECK(name);
mod = GetPass(name->value)(mod, pass_ctx);
}
mod = pass(mod, pass_ctx);
}
return mod;
}
在调用pass时，先检查是否启用了pass。检查用户是否明确禁用pass，是否被用户指定为必需pass完成的。如果不确定，是否启用了此pass，opt_level将进行检查。只有当优化级别不低于PassContext中，配置的优化级别时，才会启用执行pass。
要执行pass，先需要使用pass名称，在 TVM 打包函数注册表中，检索已注册的pass。这是可能的，因为每个pass，都注册了一个 API 端点，将在后面展示。
Pass GetPass(const std::string& pass_name) {
using tvm::runtime::Registry;
std::string fpass_name = "relay._transform." + pass_name;
const auto* f = Registry::Get(fpass_name);
ICHECK(f != nullptr) << "Cannot find " << fpass_name
<< "to create the pass " << pass_name;
return (*f)();
}
提供了一些辅助函数，创建上述每种类型的pass。这些帮助程序，暴露给 Python 前端，使用 Python API，创建特定的 pass 对象。
Pass CreateFunctionPass(
const runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)>& pass_func,
int opt_level,
String name,
Array<String> required);
Pass CreatePrimFuncPass(
const runtime::TypedPackedFunc<PrimFunc(PrimFunc, IRModule, PassContext)>& pass_func,
int opt_level,
String name,
Array<String> required);
Pass CreateModulePass(
const runtime::TypedPackedFunc<IRModule(IRModule, PassContext)>& pass_func,
int opt_level,
String name,
Array<String> required);
Pass Sequential(tvm::Array<Pass> passes, PassInfo pass_info);
pass注册
不同级别pass的概念和用于编译的context,可以轻松注册pass，以 const 折叠为例。这个pass已经实现，折叠 Relay 函数中的常量（在 src/relay/transforms/fold_constant.cc 中找到）。
提供了一个 API，执行ExprtoExpr转换。
Expr FoldConstant(const Expr& expr);
为了将这个pass注册到pass infra，先需要决定这个pass，在哪个级别执行。由于常量折叠，发生在单个函数上，应该直观FunctionPass通过 CreateFunctionPass. 将pass_func作为打包函数返回，该函数在IRModule 中的每个函数上调用Exprto ExprAPI。{}表示此pass,不需要先决条件。否则，pass开发人员必须识别列出。
使用名称 relay._transform.FoldConstant,注册一个pass API 端点 。这个pass成为注册表中的一个条目，可以由C++（如GetPass上面的）和Python访问。
namespace transform {
Pass FoldConstant() {
runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
[=](Function f, IRModule m, PassContext pc) {
return Downcast<Function>(FoldConstant(f));
};
return CreateFunctionPass(pass_func, 2, "FoldConstant", {});
}
TVM_REGISTER_GLOBAL("relay._transform.FoldConstant")
.set_body_typed(FoldConstant);
} // namespace transform
为了允许其它 C++ 模块应用这个pass，在include/tvm/relay/transform.h 中声明了一个自由函数， 如下所示：
TVM_DLL Pass FoldConstant();
pass仪器
Pass Instrument 是一种分析pass本身的机制。例如，可以使用基础架构，了解一次pass需要多少时间和内存，或者一次pass，如何转换 IR 模块。
生命周期中的四个仪器点PassContext。
TVM_DLL void InstrumentEnterPassContext();
TVM_DLL void InstrumentExitPassContext();
TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
当输入PassContext实例的范围时，立即调用InstrumentEnterPassContext。
InstrumentExitPassContext在离开PassContext的作用域时被调用，或者在执行过程中发生异常。当tvm.transform.PassContext中的OverrideU instruments重写仪器时，会调用此方法。
在执行前，调用InstrumentBeforePass。如果运行pass，在执行后调用InstrumentAfterPass。这种行为就像：
if (pass_ctx.InstrumentBeforePass(ir_module, pass_info)) {
new_ir_module = run_pass(ir_module, pass_ctx);
pass_ctx.InstrumentAfterPass(new_ir_module, pass_info);
return new_ir_module;
}
该PassInstrument接口允许在上述四种方法中运行任意代码。多个PassInstrument实例，可以注册到一个 PassContext。PassInstrument实例按照instruments传递给参数序列，依次调用 PassContext。
PassInstrument 提供以下接口：
namespace instrument {
class PassInstrumentNode : public Object {
public:
String name;
virtual void EnterPassContext() const = 0;
virtual void ExitPassContext() const = 0;
virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0;
virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0;
virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0;
/* Other fields are omitted. */
};
class PassInstrument : public ObjectRef {
public:
TVM_DEFINE_OBJECT_REF_METHODS(PassInstrument, ObjectRef, PassInstrumentNode);
};
} // namespace instrument
提供Python前端，以PassInstrument快速实现。
在 PassContext中， PassInstrument实例的调用顺序是这样的：
with PassContext(instruments=[pi]) # pi = a PassInstrument implementation.
pi.EnterPassContext()
if pi.ShouldRun(Pass1):
pi.RunBeforePass()
Pass1()
pi.RunAfterPass()
if pi.ShouldRun(Pass2):
pi.RunBeforePass()
Pass2()
pi.RunAfterPass()
pi.ExitPassContext()
介绍一下PassInstrument接口和PassContext方法的关系。有关更多详细信息，参阅 ( src/ir/transform.cc )。
- InstrumentEnterPassContext
- EnterPassContext()按instruments传递给PassContext 的顺序执行。
- 当异常发生时，PassContext通过清除所有注册的PassInstrument实例，禁用pass检测。
- PassContext执行ExitPassContext()，成功完成的每个PassInstrument实例的方法EnterPassContext()
- 例如，如果PassInstrumentA，B，C，注册到 PassContext，A 完成，EnterPassContext()，B 抛出异常， C 永远不会执行；ExitPassContext()A 的执行。
- InstrumentExitPassContext
- 每个PassInstrument的实例ExitPassContext()，执行顺序是instruments传递给PassContext.
- 当异常发生时，instruments被清除。
- PassInstrument抛出异常后，注册的实例，不执行ExitPassContext。
- InstrumentBeforePass
- 如果pass未列为必需pass，执行ShouldRun。
- RunBeforePass如果传球没有被ShouldRun 阻挡，按照instruments的顺序执行。
- InstrumentBeforePass返回一个布尔值，指示是否应该运行传递。
- 当异常发生时，立即抛出。依靠 Python Context ManagerPassContext安全退出（ExitPassContext每个仪器都会运行的含义。对于 C++，参阅include/tvm/support/with.h。）
- InstrumentAfterPass
- RunAfterPass按instruments传递给 PassContext的顺序执行。
- 当异常发生时，立即抛出。依靠 Python Context Manager 或Withclass( include/tvm/support/with.h )，安全退出PassContext
build仪器
有几种内置工具，标有TODO 的，没有实现。
- PassTimingInstrument（参见src/ir/instrument.cc）
- 分析pass的执行时间。
- PrintIRBefore(TODO)
- 在pass转换前，打印 IR 模块。如果在pass周围插入，tvm.transform.PrintIR()可以达到这个目的。但是，使用PassInstrument，不需要修改passes的顺序。
- 打印后（待办事项）
- 在pass转换后，打印 IR 模块。
Python前端
前端只需要一些简单的 API。例如，可以为用户提供以下 API，创建和执行一个 pass（完整的实现在python/tvm/relay/transform/transform.py和 python/tvm/ir/transform.py 中提供）。后端接收信息，决定应该使用哪个函数，创建 Pass 对象。
PassContext
Python 前端为__enter____exit__current 提供了一个包装器，通过覆盖和PassContext，启用with语法。为用户提供了一种静态方法，获取在一定范围内使用的Context。
@tvm._ffi.register_object("transform.PassContext")
class PassContext(tvm.runtime.Object):
def __enter__(self):
_transform.EnterPassContext(self)
return self
def __exit__(self, ptype, value, trace, config):
_transform.ExitPassContext(self)
@staticmethod
def current():
"""Return the current pass context."""
return _transform.GetCurrentPassContext()
PassContext用于配置编译选项，包括优化级别和必需/禁用的pass。可以带一个配置字典，以便不同的pass，可以方便地获取pass的数据，如回退设备信息和循环展开的步骤/深度等。为了能够获取所需的配置，必须通过 TVM_REGISTER_PASS_CONFIG_OPTION注册密钥。例如，使用以下内容，循环展开pass
TVM_REGISTER_PASS_CONFIG_OPTION("tir.UnrollLoop", UnrollLoopConfig);
更多细节，参考src/tir/transforms/unroll_loop.cc。
pass对象
Pass是所有pass对象的基类。这里的所有方法，都只是在后端实现的简单包装器。为了用户方便地与 Python 中的基类，进行交互定义的。在 pass 基类中只定义了__call__，使子类成为可调用对象，可以很容易调用（例如，pass_xx(arg)）执行。
@register_relay_node
class Pass(RelayNode):
def __call__(self, mod):
return _transform.RunPass(self, mod)
提供了一些辅助 API，从 Python 前端，轻松创建pass，让pass基础控制执行。例如module_pass，function_pass和sequential提供给用户，以便可以定制pass。
对于在 C++ 后端实现的所有 pass，分别在python/tvm/ir/transform.py和 python/tvm/relay/transform/transform.py 中，提供了相应的 Python API 。例如，const 折叠有一个 Python API，如下所示：
def FoldConstant():
return _transform.FoldConstant()
可以构建一个pass through装饰：
@relay.transform.module_pass(opt_level=2)
def transform(mod, ctx):
tp = relay.TensorType((10,), "float32")
x = relay.var("x", tp)
gv = relay.GlobalVar("abs")
func = relay.Function([x], relay.abs(x))
new_mod = tvm.IRModule({gv: func})
new_mod.update(mod)
return new_mod
module_pass = transform
assert isinstance(module_pass, transform.ModulePass)
assert module_pass.info.opt_level == 2
在transform功能增加了一个abs与输入模块的功能，可能是在模块级的任何定制的优化。创建module_pass后，应用于任何 Relay 模块。例如，可以构建一个空模块，应用pass添加一个abs 函数。
mod = tvm.IRModule()
mod = module_pass(mod)
提供function_pass功能，一个示例函数级pass，可以写成如下：
@relay.transform.function_pass(opt_level=1)
class TestReplaceFunc:
def __init__(self, new_func):
self.new_func = new_func
def transform_function(self, func, mod, ctx):
# Just for demo purposes
# Transform func to new_func
return self.new_func
x = relay.var("x", shape=(10, 20))
f1 = relay.Function([x], x)
f2 = relay.Function([x], relay.log(x))
# fpass is now a special pass that replaces every
# function to f1
fpass = TestReplaceFunc(f1)
# Now every function in input_mod is replaced by f1
res_mod = fpass(input_mod)
可以不使用装饰器，直接注册pass，调用。有关如何自定义优化pass，调试 Relay 和 tir pass 的更多示例，参阅 use pass infra教程。
pass仪器
可以通过在实现以下方法的类上，使用pass_instrument decorator（python/tvm/ir/instrument.py），实现PassInstrument。建议使用pass_instrument decorator，实现PassInstrument，不是重写或子类化。
- enter_pass_ctx
- 该方法在进入PassContext时运行。
- exit_pass_ctx
- 此方法在退出PassContext时运行。
- should_run
- 此方法在执行pass前运行，返回一个布尔值，指示是否应运行pass。
- run_before_pass
- 如果应该运行一次pass，在pass执行之前运行此方法。
- run_after_pass
- 此方法在执行一次pass后，立即运行。
PassInstrument实例可以通过tvm.transform.PassContext中的参数 instruments注册。
use pass instrument提供了如何使用 Python API，实现PassInstrument的示例。
覆盖当前 PassContext 中的仪器
提供了current PassContext覆盖instruments 的override_instruments方法。例如，如果在没有显式创建 new 的情况下，运行 pass PassContext，可以通过以下方式注册PassInstrument到全局中PassContext：
cur_pass_ctx = tvm.transform.PassContext.current()
# override PassInstrument instances
cur_pass_ctx.override_instruments([pass_inst])
mod = pass_seq(mod)
result = pass_inst.get_result()
当override_instruments调用时，旧PassInstrument实例的方法exit_pass_ctx会调用。然后new PassInstrument的enter_pass_ctx方法调用。
参考链接：