- [翻译]Relay Pass基础结构— tvm 0.6.0文档
- ">Relay Pass基础结构
[翻译]Relay Pass基础结构— tvm 0.6.0文档
Relay Pass基础结构
Relay具有一系列优化Pass，它们可改善模型的性能指标，例如平均推断，内存占用量或特定设备的功耗。有一套标准优化以及专为机器学习的优化，包括常量折叠（constant folding）、死代码的消除、算子布局的更改和算子融合（operator fusion）等。这些Pass中的每一个都使用遍历期间和/或之前收集的分析结果,从而构造为抽象语法树（AST）上的Relay到Relay转换。但是，随着Relay的快速发展，对管理这些Pass的更加系统和有效的方法的需求日益明显。本文档描述了此类基础结构的设计，该基础结构利用了生产编译器用于管理优化Pass的方式以及采用现代深度学习框架来构建层的样式。例如，许多现有的生产编译器，如GCC和LLVM，都使用Pass管理器来有效地管理Pass的执行。由于Pass数量很少，因此最初管理Pass非常简单，但是成熟的编译器将包含数百个单独的Pass。通常，外部用户会希望合理安排自定义的Pass，无需修改单个手工Pass顺序。类似地，诸如Pytorch和MXNet Gluon之类的现代深度学习框架也倾向于分别通过Sequential和Block来实现Pass样式的层构建方案。通过这种构造，这些现代框架能够方便地将模块/层添加到其容器中，并轻松地构建神经网络。Relay Pass基础结构的设计在很大程度上受到LLVM中使用的分层Pass管理器和流行的深度学习框架中使用的块式容器的启发。Pass基础结构的主要目标包括：
- 实现更好的优化程序编排。这让用户可以灵活地自定义和构建自己的优化pipeline。
- 提供一种用户友好的方式来调试优化Pass。
- 减轻了开发人员手动地分别解决Pass之间的依赖性的麻烦。
- 为开发人员简化了新Pass的实施。例如，我们允许用户在Python中实现Pass，并让该Pass基础结构控制其执行。
设计
我们专注于为用户扩展的便利性，使用户可以快速添加新Pass而不会失去向后兼容性。该设计包含后端(C++)和前端(Python)。前者实现了Pass的主要逻辑，后者为用户提供了与之交互的简单API，即允许用户快速创建自己的优化pipeline。### C++后端 我们提供了一个PassInfo对象，其中包含Pass所需的基本信息。name是Pass名称，opt_level表示要在哪个优化级别启用Pass，required表示执行该Pass所需的其它的pass（有关更多详细信息，请参见include/tvm/relay/transform.h）。例如，在注册Pass期间（稍后将介绍），Pass开发人员可以指定Pass的名称，将在其上执行的优化级别，以及/或所需的其它pass。opt_level在用户提供的优化级别下运行时，可用于帮助Pass基础结构识别是否需要执行某个Pass。Pass基础结构可使用字段required来解析Pass依赖性。class PassInfoNode : public RelayNode{ std::string name;//pass name int optlevel;//which optimization level the pass will be enabled std::vectorPassContext携带有用的信息用来优化Pass。例如，它包含错误报告系统，因此优化作者可以提供有关优化失败原因的诊断。PassContext的设计也可以用来代替旧版BuildConfig，后者用于帮助用户配置编译选项，包括优化级别和必需/禁用的Pass等。例如，我们可能有一个配置，使用PassContext提供的disabledPass=xx禁用一些Pass，来执行optlevel=3所有Pass。现在，我们可以遍历opt_level=3的所有Pass，并排除禁用Pass列表中的那些Pass。此类旨在让用户方便地编写Pythonwith语法，以在特定配置下执行优化。另外，由于线程本地存储RelayPassContextThreadLocalStore用于保存创建的Pass上下文对象，因此用户可以用Pass线程安全的方式，通过PassContext::Current()获得在某个程序范围内可用的上下文。稍后将提供示例，以展示如何使用C++和Python API 通过Pass上下文创建编译pipeline。class PassContextNode : public RelayNode{public: ErrorReporter errreporter; int optlevel{2}; int fallbackdevice{staticcast(kDLCPU)}; tvm::Array (self, ptype, value, trace): _transform.ExitPassContext(self)
class PassContext : public NodeRef{public: TVMDLL static PassContext Create(); TVMDLL static PassContext Current(); / Other fields are omitted. /
private: // The entry of a pass context scope. TVMDLL void EnterWithScope(); // The exit of a pass context scope. TVMDLL void ExitWithScope();
// Classes to get the Pythonwithlike syntax. friend class tvm::With;}; 和__exit，为PassContext提供一个包装器，以启用
struct RelayPassContextThreadLocalEntry{ /! \brief The default pass context. / PassContext defaultcontext; /! \brief The current pass context. / std::stackcontext_stack; RelayPassContextThreadLocalEntry() { default_context = PassContext(make_node ()); }};
/! \brief The thread-local store to hold the pass context. /typedef dmlc::ThreadLocalStoreRelayPassContextThreadLocalStore;#### Pass构造 Pass基础以分层方式设计，并且可以在Relay程序的不同粒度下工作。 PassNode引入纯虚拟类作为不同优化Pass的基础。此类包含几个虚函数，这些虚函数必须由子类在模块、函数或连续的Pass上实现。class PassNode : RelayNode{ virtual PassInfo Info() const = 0; virtual Module operator()(const Module &mod, const PassContext &pass_ctx) const = 0;};该功能显示必须如何实现Pass，即，它始终在特定上下文下在Relay模块上起作用。所有Pass均以Module到Module方式设计。因此，由Pass基础结构控制的优化将始终更新整个模块。我们已经创建了几个子类来实现不同类型的优化Pass，例如，函数级Pass，模块级Pass和连续的Pass。每个子类本身都可以充当Pass管理器。例如，它们可以收集所需的Pass并执行它们，或者基于给定的元数据构建依赖关系图。可以在src/relay/Pass/pass_manager.cc中找到它们的完整定义。#### 模块级Pass 模块级别Pass将整个程序视作一个单元来处理的Pass，主要用于全局和Pass间优化（IPO），这与LLVM中使用的模块Pass类似。Relay中一些需要全局模块图的典型Pass，例如A-normal格式转换和lambda提升等，都属于此类。在此级别上，用户甚至可以在模块中添加和/或删除函数。class ModulePassNode : PassNode{ PassInfo pass_info; runtime::TypedPackedFuncpass_info维护模块级别Pass所需的信息，pass_func勾画出真正的优化。例如，我们可能需要在Module上删除无效代码。我们可以在pass_func中实现该算法，并使其在Module上运行。然后它将删除无效代码，包括模块中未使用的函数。请注意，pass_func被设计为包装函数（packed function），从而可以在C++和Python中实现该优化。#### 函数级Pass 函数级别Pass是以单个函数为作用域的Pass, 每个函数间是相互独立的，用于为给定的relay模块实现各种函数内级别的优化。它一次从模块的函数列表中获取一个函数以进行优化，并生成重写的Relay函数。Relay的大多数Pass都可以归为此类，例如常见的子表达式消除和推理简化等。请注意，此级别Pass的作用域是Relay函数。因此，我们无法通过这些Pass来添加或删除函数，因为它们不了解全局信息。class FunctionPassNode : PassNode{ PassInfo pass_info; runtime::TypedPackedFuncpass_info与我们刚才在模块Pass中描述的相同。pass_func采用一个Function进行优化，它还需要一个Module，因为我们可能会使用它来报告错误。可以用“SkipOptimization”来标注某个Function，以便在优化时忽略该Pass。#### 连续的PassSequentialPass与Pytorchnn.Sequential相似，其中包含许多要执行的Pass。class SequentialPassNode : PassNode{ PassInfo pass_info; // Passes need to be executed. Arraypasses; 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) { CHECK(pass.defined()) << “Found undefined pass for optimization.”; const PassInfo &pass_info = pass->Info(); if (!PassEnabled(pass_info)) continue; //没有启用该pass for (const auto &it : pass_info->required) { const auto name = it.asopt_level。此Pass 将被启用，因此仅当其优化级别不小于在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
ModulePass CreateModulePass(std::string name, int opt_level, PassFunc pass_func);
SequentialPass CreateSequentialPass(std::string name, int opt_level, ArrayPasses, Array SequentialPass上工作。出于说明目的，仅提供一个代码段。首先，我们创建一个简单的Relay程序y=f(x)。然后，我们基于该函数构建一个模块。创建模块之后，我们实例化一个连续的Pass对象，该对象包含一些标准的Relay优化Pass，包括类型推断、无效代码消除、公共子表达式消除和布局更改。最后，构建Pass上下文，并且将按顺序执行Pass。在执行这些Pass期间，因为我们已经在注册Pass中对相关Pass进行了编码，所以Pass依赖关系将自动解决。// Create a simple Relay program.auto tensor_type = relay::TensorTypeNode::make({}, tvm::Bool());auto x = relay::VarNode::make(“x”, relay::Type());auto f = relay::FunctionNode::make(tvm::Array
auto y = relay::VarNode::make(“y”, tensor_type);auto call = relay::CallNode::make(f, tvm::Array
// Create a module for optimization.auto mod = relay::ModuleNode::FromExpr(fx);
// Create a sequential pass.tvm::Array
// Create a pass context for the optimization.auto ctx = relay::transform::PassContext::Create();ctx->opt_level = 2;ctx->fallback_device = kDLCPU;
// Use the Python with syntax to execute the sequence of optimizations.tvm::With
// View the updated module.LOG(INFO) << relay::AsText(mod) << std::endl;其他类型的Pass应该直接调用以在模块上执行。例如，用户可以直接在给定模块上应用const fold pass, mod=transform::FoldConstant()(mod)。但是，明确地执行所需的Pass是用户的责任。### Pass注册 我们已经介绍了不同级别的Pass概念以及用于编译的上下文。看看用户如何轻松注册Pass会很有趣。让我们以常量折叠(const folding)为例。在Relay该Pass已实现为Pass FoldConstant（位于src/relay/Pass/fold_constant.cc中）。提供了一个API来执行Expr到Expr转换：ExprFoldConstant(constExpr&expr);(1) 定义Pass函数为了将此Pass注册到Pass基础结构，我们首先需要确定将在哪个级别执行此Pass（是模块级还是函数级）。因为常量折叠(const folding)发生在单独的函数中，我们直接通过CreateFunctionPass创建一个FunctionPass。将pass_func其作为打包函数(packed function)返回，该函数在Relay模块中的每个函数上调用ExprtoExprAPI。{}表示此Pass不需要任何先决条件。否则，Pass开发人员必须识别并列出它们。(2) 注册Pass API用 relay._transform.FoldConstant注册了一个Pass API端点。因此，该Pass成为注册表中的一个入口，可以在需要时从C++（例如，GetPass上述版本）和Python访问该Pass。namespace transform{
Pass FoldConstant(){ runtime::TypedPackedFunc(FoldConstant(f)); }; return CreateFunctionPass(pass_func, 2, “FoldConstant”, {});}
TVM_REGISTER_API(“relay._transform.FoldConstant”) .set_body_typed(FoldConstant);
} // namespace transform(3) 声明Pass函数为了允许其他C++模块使用此Pass，我们在include/tvm/relay/transform.h中声明了该函数，如下所示：### Python前端 前端只需要一些简单的API。例如，我们可以为用户提供以下API以创建和执行Pass（完整实现在python/tvm/relay/transform.py中提供）。后端接收信息，并决定应使用哪个函数来创建Pass对象。#### PassContext Python前端覆盖__enterwith语法。为用户提供了一种current静态方法，用来获取在一定范围内使用的上下文。@register_relay_nodeclass PassContext(RelayNode): def __enter(self): _transform.EnterPassContext(self) return self
def __exit
@staticmethod def current(): “””Return the current pass context.””” return _transform.GetCurrentPassContext()可以通过build_configAPI来实例化一个PassContext对象，该对象被relay用来配置编译选项，包括优化级别，用于异构执行的后备设备以及必需/禁用的Pass。#### Pass对象Pass是所有Pass对象的基类。这里的所有方法只是在后端实现的简单包装器。定义它们是为了使用户可以方便地与Python中的基类进行交互。在Pass基类中仅定义了一个__call ，以使子类成为可调用对象，以便可以轻松执行调用（例如pass_xx(arg)）。@register_relay_nodeclass Pass(RelayNode): def __call(self, mod) return _transform.RunPass(self, mod)提供了一些辅助API，可轻松创建来自Python前端的Pass，并使Pass基础控制执行。例如，将module_Pass，function_Pass和sequential提供给用户，以便他们可以自定义自己的Pass或Passpipeline。(4) 添加Python API对于在C++后端中实现的所有Pass，我们在python/tvm/relay/transform.py中提供了相应的Python API。例如，常量折叠具有类似以下的Python API：
def FoldConstant(): return _transform.FoldConstant()
用户可以通过装饰器来构建Pass，如下所示：@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 = relay.Module({gv: func}) new_mod.update( ) return new_mod
module_pass = transformassert isinstance(module_pass, transform.ModulePass)assert module_pass.info.opt_level == 2该transform函数在此处为输入模块添加了一个函数abs，但也可以是模块级别的任何自定义优化。创建module_Pass之后，用户可以将其应用于任何Relay模块。例如，我们可以构建一个空模块并应用此Pass去添加abs功能。mod = relay.Module()mod = module_pass(mod)相应地，我们还为function_pass提供了此类功能。例如，函数级别Pass示例可以这么写：@relay.transform.function_pass(opt_level=1)class TestReplaceFunc: def __init
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 f1fpass = TestReplaceFunc(f1)# Now every function in input_mod is replaced by f1res_mod = fpass(input_mod)或者，用户也可以不使用装饰器直接注册Pass，然后调用它。让我们用Sequential来演示这种情况。#### Python顺序示例 此示例不仅说明用户如何使用Python API直接创建连续的Pass（这也可以应用于模块级和函数级Pass），还说明了如何使用Sequential与其他类型的Pass关联的方式构建优化pipeline 。# Create a simple Relay program.shape = (1, 2, 3)c_data = np.array(shape).astype(“float32”)tp = relay.TensorType(shape, “float32”)c = relay.const(c_data)x = relay.var(“x”, tp)y = relay.add(c, c)y = relay.multiply(y, relay.const(2, “float32”))y = relay.add(x, y)z = relay.add(y, c)z1 = relay.add(y, c)z2 = relay.add(z, z1)func = relay.Function([x], z2)
# Customize the optimization pipeline.seq = _transform.Sequential([ relay.transform.InferType(), relay.transform.FoldConstant(), relay.transform.EliminateCommonSubexpr(), relay.transform.AlterOpLayout()])
# Create a module to perform optimizations.mod = relay.Module({“main”: func})
# Users can disable any passes that they don’t want to execute by providing# a list, e.g. disabled_pass=[“EliminateCommonSubexpr”].with relay.build_config(opt_level=3): with tvm.target.create(“llvm”): # Perform the optimizations. mod = seq(mod)### 调试 Pass基础提供了特殊的Pass（PrintIR），以便在应用特定Pass之后dump 整个模块的IR。连续的Pass示例的略微修改版本可能类似于以下内容，以启用IR转储以进行FoldConstant优化。seq = _transform.Sequential([ relay.transform.InferType(), relay.transform.FoldConstant(), relay.transform.PrintIR(), relay.transform.EliminateCommonSubexpr(), relay.transform.AlterOpLayout()])在FoldConstant之后插入PrintIR Pass，当FoldConstant完成后，Pass基础结构将转储模块IRFoldConstant。用户可以在要调试的任何Pass之后插入此Pass，以查看优化效果。有关Python和C++中与Pass相关的更多示例，请分别参考tests/python/relay/test_pass_manager.py和tests/cpp/relay_transform_sequential.cc。
若有收获，就点个赞吧
0 人点赞
