标签：分析 const transform 基础架构 pass PassInstrument PassContext mod
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。在用户提供的优化级别下Runtime，是否需要执行某个 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 {
/*! \brief The default pass context. */
PassContext default_context;
/*! \brief The current pass context. */
std::stack<PassContext> context_stack;
PassContextThreadLocalEntry() {
default_context=PassContext(make_node<PassContextNode>());
}
};
/*! \brief 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();
passInstruments
Pass Instrument 是一种分析pass本身的机制。例如，可使用基础架构，了解一次pass需要多少时间和内存，或者一次pass，如何转换 IR 模块。
生命周期中的四个Instruments点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重写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每个Instruments都会运行的含义。对于 C++，参阅include/tvm/support/with.h。）
InstrumentAfterPass
- RunAfterPass按instruments传递给 PassContext的顺序执行。
- 当异常发生时，立即抛出。依靠 Python Context Manager 或Withclass( include/tvm/support/with.h)，安全退出PassContext
buildInstruments
有几种内置工具，标有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教程。
passInstruments
可通过在实现以下方法的类上，使用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 中的Instruments
提供了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方法调用。
参考链接：
https://tvm.apache.org/docs/dev/pass_infra.html#pass-infra
标签：分析,const,transform,基础架构,pass,PassInstrument,PassContext,mod 来源： https://www.cnblogs.com/wujianming-110117/p/15665385.html
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。