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基础架构的主要目标包括：

1. 实现更好的优化编程调度。允许用户灵活地定制和构建优化pass。
2. 提供一种用户友好的方式来调试优化pass。
3. 减轻开发人员手动和分别解决pass之间的依赖关系。
4. 为开发人员简化新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 {

  /*! \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();

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方法调用。

 

 

参考链接：

https://tvm.apache.org/docs/dev/pass_infra.html#pass-infra