Relay传递基础架构

最新推荐文章于 2024-06-14 09:32:16 发布

牛牛存

最新推荐文章于 2024-06-14 09:32:16 发布

阅读量1.6k

点赞数 1

分类专栏： tvm TVM Relay

本文链接：https://blog.csdn.net/weixin_42164269/article/details/104291538

版权

TVM Relay 同时被 2 个专栏收录

25 篇文章 21 订阅

订阅专栏

tvm

14 篇文章 6 订阅

订阅专栏

Relay传递基础架构

Relay具有一系列优化过程，可改进模型的性能指标，例如特定设备的平均推断，内存占用量或功耗。有一套标准优化以及特定于机器学习的优化，包括常量折叠，消除死代码，更改操作符布局和操作符融合等。这些步骤中的每一个都被构造为抽象语法树（AST）的上Relay到Relay转换，使用遍历期间和/或之前收集的分析结果。

但是，随着Relay的快速发展，对管理这些通道的更加系统和有效的方法的需求日益明显。本文档介绍了此类基础架构的设计，该基础架构利用了生产编译器的方式用于管理优化过程以及采用现代深度学习框架的样式来构建层。

例如，许多现有的生产编译器，例如GCC和LLVM，都使用过程管理器来有效地管理转换的执行。由于转换次数很少，因此最初管理转换非常简单，但是成熟的编译器将包含数百个单独的转换。通常，外部用户会希望正确安排自定义转换，而不必手工修改单个转换顺序。

类似地，诸如Pytorch和MXNet Gluon之类的现代深度学习框架也倾向于分别通过【Sequential】和【Block】来实现转换样式的层构建方案。通过这种结构，这些现代框架能够方便地将模块/层添加到其容器中，并轻松地构建神经网络。

Relay转换基础架构的设计在很大程度上受到LLVM中使用的转换管理器以及流行的深度学习框架中使用的块式容器的启发。转换基础架构的主要目标包括：

实现更好的优化程序编排。这使用户可以灵活地自定义和构建自己的优化管道。
提供一种用户友好的方式来调试优化过程。
减轻开发人员手动和分别解决转换之间的依赖性的麻烦。
为开发人员简化了新转换的实施。例如，我们允许用户在Python中实现转换，并让该转换基础架构在下面操纵其执行。

设计

我们专注于为用户扩展的便利性，使用户可以快速添加新的转换而不会失去向后兼容性。该设计同时包含后端和前端。前者实现了转换基础架构的主要逻辑。后者为用户提供了与之交互的简单API，即允许用户快速创建自己的优化管道。

C ++后端

我们提供了一个【PassInfo】对象包含转换所需基本信息的对象。【name】是转换名称，【opt_level】指示要在哪个优化级别启用过程，【required】表示执行特定转换所需的转换（有关更多详细信息，请参见【include / tvm / ir / transform.h】）。例如，在转换注册期间（稍后将介绍），转换开发人员可以指定转换的名称，将在其上执行转换所需的转换。【opt_level】在用户提供的优化级别下运行时，可用于帮助通过转换基础架构识别是否需要执行某个过程。【required】转换基础结构可以使用该字段来解析传递依赖性。

class PassInfoNode : public RelayNode {
  std::string name;
  int opt_level;
  std::vector<std::string> required;
};

PassContext

【PassContext】携带有用的信息以进行优化。例如，它包含错误报告系统，因此优化作者可以提供有关优化失败原因的诊断。【PassContext】还可以用来代替旧版【BuildConfig】，后者用于帮助用户配置编译选项，包括优化级别和必需/禁用的转换。例如，我们可能有一个由【PassContext】提供的配置，在【opt_level=3】下执行所有转换，使用【disabled_pass=xx】提供禁用的转换。现在我们可以遍历【opt_level=3】的所有转换，排除禁用列表中的转换。

这个类用于让用户方便地编写Python【with 】语法以在特定配置下执行优化。另外，PassContext::Current()由于线程本地存储【RelayPassContextThreadLocalStore】用于保存创建的转换上下文对象，因此用户可以通过线程安全的方式获得在某个程序范围内可用的上下文。稍后将提供示例，以展示如何使用C ++和Python API通过转换上下文创建编译管道。

class PassContextNode : public RelayNode {
 public:
  ErrorReporter err_reporter;
  int opt_level{2};
  int fallback_device{static_cast<int>(kDLCPU)};
  tvm::Array<tvm::Expr> required_pass;
  tvm::Array<tvm::Expr> disabled_pass;
};

class PassContext : public NodeRef {
 public:
  TVM_DLL static PassContext Create();
  TVM_DLL static PassContext Current();
  /* 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 RelayPassContextThreadLocalEntry {
  /*! \brief The default pass context. */
  PassContext default_context;
  /*! \brief The current pass context. */
  std::stack<PassContext> context_stack;
  RelayPassContextThreadLocalEntry() {
    default_context = PassContext(make_node<PassContextNode>());
  }
};

/*! \brief The thread-local store to hold the pass context. */
typedef dmlc::ThreadLocalStore<RelayPassContextThreadLocalEntry>
     RelayPassContextThreadLocalStore;

转换构造

转换基础架构以分层方式设计，并且可以在Relay程序的不同粒度下工作。引入纯虚类【PassNode】作为不同优化过程的基础。此类包含一些虚函数，这些虚函数必须在子类以模块，函数或转换序列级别上实现。

class PassNode : RelayNode {
  virtual PassInfo Info() const = 0;
  virtual Module operator()(const Module& mod
                            const PassContext& pass_ctx) const = 0;
};

函子显示必须如何实现转换，即，它始终在特定上下文下在Relay模块上起作用。所有转换均以【 Module】到【Module 】方式设计。因此，由转换基础架构进行的优化将始终更新整个模块。

已经创建了几个子类来实现不同类型的优化过程，例如，函数级别的转换，模块级别的转换和顺序转换。每个子类本身都可以充当转换管理器。例如，他们可以收集所需的转换并执行它们，或基于给定的元数据构建依赖关系图。可以在【src / relay / ir / transform.cc】和【src / ir / transform.cc】中找到它们的完整定义。

模块级转换

模块级别转换主要用于全局和过程间优化（IPO），与LLVM中使用的模块转换类似。Relay中一些需要模块全局性的典型转换，例如A-normal格式转换和lambda提升等，都属于此组。在此级别上，用户甚至可以在模块中添加和/或删除函数。

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_func】勾画出真正的优化。例如，我们可能需要在模块上执行无效代码消除。我们可以在【pass_func】中实现该算法，并使其在模块上运行。然后它将删除无效代码，包括模块中未使用的函数。请注意，此字段被设计为打包函数，从而可以在C ++和Python中实现优化。

函数级转换

函数级别转换用于为给定的Relay模块实现各种函数内级别优化。它一次从模块的函数列表中获取一个函数以进行优化，并生成重写的Relay函数。Relay的大多数转换都可以归为此类，例如常见的子表达式消除和推理简化等。

请注意，此级别的通过范围是中继功能。因此，我们无法通过这些传递来添加或删除函数，因为它们不了解全局信息。

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_func】具有用于优化的函数，还需要一个模块，因为我们可能会使用它来报告错误。可以用“ SkipOptimization”注释一个函数，以便在优化过程中将其忽略。

顺序转换

【SequentialPass】与Pytorch中的【nn.Sequential】类似，包含许多执行通道。

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中的几个转换被放入该组。例如，【FoldScaleAxis】要求派遣【ForwardFoldScaleAxis】和【 BackwardFoldScaleAxis】内部。另外，【BackwardFoldScaleAxis】建议先实现。因此，此【SequentialPass】转换是的理想候选人。

以下代码显示了如何调用顺序转换中的各个转换。本质上，我们使用附加到转换列表的顺序依次执行每个过程。

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;
    for (const auto& it : pass_info->required) {
      const auto* name = it.as<tvm::ir::StringImm>();
      CHECK(name);
      mod = GetPass(name->value)(mod, pass_ctx);
    }
    mod = pass(mod, pass_ctx);
  }
  return mod;
}

调用转换后，我们首先检查此转换是否已启用。首先检查用户是否明确禁用了转换，然后检查用户是否将其指定为必需转换。如果仍不确定此转换是否启用，将检查【opt_level】。此转换将被启用并执行，仅当其优化级别不小于在通道上下文中配置的优化级别时。

要执行转换，我们首先需要使用转换名称在TVM打包函数注册表中检索已注册的转换。这是可能的，因为每次通过都向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);
  CHECK(f != nullptr) << "Cannot find " << fpass_name
                      << "to create the pass " << pass_name;
  return (*f)();
}

提供了一些辅助函数来创建上述转换的每种类型。这些辅助函数还暴露于Python前端，以使用户可以方便地使用Python API创建特定的传递对象。

FunctionPass CreateFunctionPass(std::string name,
                                int opt_level,
                                PassFunc pass_func);

ModulePass CreateModulePass(std::string name,
                            int opt_level,
                            PassFunc pass_func);

SequentialPass CreateSequentialPass(std::string name,
                                    int opt_level,
                                    Array<Pass> passes,
                                    Array<tvm::Expr> disabled);

C ++顺序示例

现在让我们以一个示例来说明转换的基础架构中【SequentialPass】如何工作。出于说明目的，仅提供一个代码段。首先，我们创建一个简单的Relay程序【y = f(x)】。然后，我们基于该函数构建一个模块。创建模块后，我们实例化一个顺序转换对象，该对象包含一些标准的Relay优化转换，包括类型推断，死代码消除，公共子表达式消除和布局更改。

最后，构建转换上下文，并且将按顺序执行转换。在执行这些转换期间，因为我们已经在注册过程中对相关转换进行了编码，所以转换依赖关系将自动解决。

// 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<relay::Var>{ x }, x, relay::Type(), {});

auto y = relay::VarNode::make("y", tensor_type);
auto call = relay::CallNode::make(f, tvm::Array<relay::Expr>{ y });
auto fx = relay::FunctionNode::make(tvm::Array<relay::Var>{ y }, call, relay::Type(), {});

// Create a module for optimization.
auto mod = IRModule::FromExpr(fx);

// Create a sequential pass.
tvm::Array<relay::transform::Pass> pass_seqs{
   relay::transform::InferType(),
   relay::transform::DeadCodeElimination(),
   relay::transform::EliminateCommonSubexpr(),
   relay::transform::AlterOpLayout()
};
relay::transform::Pass seq = relay::transform::Sequential(pass_seqs);

// 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<relay::transform::PassContext> scope(ctx);
mod = seq(mod);

// View the updated module.
LOG(INFO) << relay::AsText(mod) << std::endl;

其他类型的转换应该直接调用以在模块上执行。例如，用户可以将常量折叠遍历直接应用于给定的模块。但是，明确地执行所需的通行证是用户的责任。mod = transform::FoldConstant()(mod)

转换注册

我们已经介绍了不同级别的转换概念以及用于编译的上下文。看看用户如何轻松注册转换会很有趣。让我们以常量折叠为例。此转换已实现为在Relay函数中折叠常量（位于 src / relay / pass / fold_constant.cc中）。

提供了一个API来执行【Exprt】到【Expr】转换。

Expr FoldConstant(const Expr& expr);

为了将此转换注册到通转换基础架构，我们首先需要确定将在哪个级别执行此转换。作为常量折叠发生在个别的函数，我们应该直观地创建【FunctionPass】通过【CreateFunctionPass】。将【pass_func】其作为打包函数返回，该函数在Relay模块中的每个函数上调用【Expr】到【Expr】API。{} 表示此通行证不需要任何先决条件。否则，转换开发人员必须识别并列出他们。

同时，pass API端点已使用name注册 relay._transform.FoldConstant。因此，此转换成为注册表中的一项，可以在需要时由C ++（例如，GetPass上述版本）和Python访问。

namespace transform {

Pass FoldConstant() {
  runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func =
    [=](Function f, Module 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 ++模块应用此过程，我们在`include / tvm / relay / transform.h`_中声明一个自由函数，如下所示：

TVM_DLL Pass FoldConstant();

Python前端

前端只需要一些简单的API。例如，我们可以向用户提供以下API以创建和执行转换（完整的实现在python / tvm / relay / transform.py中提供）。后端接收信息，并决定应使用哪个函数来创建Pass对象。

PassContext

Python前端为【PassContext】提供的包装器，通过覆盖【__enter__】和【__exit__】来启用【with】语法。【current 】静态方法为用户提供了一种来获取在一定范围内使用的上下文。

@register_relay_node
class PassContext(RelayNode):
    def __enter__(self):
        _transform.EnterPassContext(self)
        return self

    def __exit__(self, ptype, value, trace):
        _transform.ExitPassContext(self)

    @staticmethod
    def current():
        """Return the current pass context."""
        return _transform.GetCurrentPassContext()

【PassContext】对象可以通过【build_config】API被实例化，其用于通过Relay来配置编译选项，包括优化级别，回退装置异构执行，以及需要/禁止转换。

Pass 对象

【Pass】是所有转换对象的基类。这里的所有方法只是在后端实现的简单包装器。定义它们是为了使用户可以方便地与Python中的基类进行交互。只有【__call__】在传递基类中定义了，以使子类成为可调用对象，以便可以轻松地调用它们（例如【pass_xx(arg)】）以执行。

@register_relay_node
class Pass(RelayNode):
   def __call__(self, mod):
       return _transform.RunPass(self, mod)

提供了一些辅助API，可轻松创建来自Python前端的转换，并使转换基础架构控制执行。例如，将【module_pass】，【function_pass】和【sequential】提供给用户，以便他们可以自定义自己的转换或转换管道。

对于在C ++后端中实现的所有转换，我们在python / tvm / relay / transform.py中提供了相应的Python API 。例如，常量折叠具有类似以下的Python API：

def FoldConstant():
    return _transform.FoldConstant()

用户可以通过如下构装饰建转换：

 @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(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模块。例如，我们可以构建一个空模块并给此过程添加abs 函数。

mod = relay.Module()
mod = module_pass(mod)

相应地，我们还为【function_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)

或者，用户也可以不使用装饰器直接注册转换，然后调用它。让我们使用【Sequential】来演示这种情况。

Python顺序示例

该示例不仅说明用户如何使用Python API直接创建顺序转换（这也可以应用于模块级和函数级转换），还说明了如何使用【Sequential】与其他传递类型相关联的方法来构建优化管道。

# 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)

调试

传递基础框架提供了特殊的传递（【PrintIR】），以在应用特定传递之后转储整个模块的IR。顺序传递示例的略微修改版本可能类似于以下内容，以启用IR转储以进行【FoldConstant】优化。

seq = _transform.Sequential([
    relay.transform.InferType(),
    relay.transform.FoldConstant(),
    relay.transform.PrintIR(),
    relay.transform.EliminateCommonSubexpr(),
    relay.transform.AlterOpLayout()
])

通过在【PrintIR】之后插入转换【FoldConstant】，当【FoldConstant】完成后，转换基础架构将转储模块IR 。用户可以在要调试的任何转换之后插入此转换，以查看优化效果。

构建配置对象还公开了更灵活的调试机制。可以传递一个跟踪函数，该函数可用于在每次转换之前和/或之后执行任意代码。一个跟踪函数将收到一个【IRModule】，【PassInfo】以及一个布尔值指示是否正在执行之前或者之后。下面是一个示例。

def print_ir(mod, info, is_before):
    """Print the name of the pass, the IR, only before passes execute."""
    if is_before:
        print(f"Running pass: {}", info)
        print(mod)

with relay.build_config(opt_level=3, trace=print_ir):
        with tvm.target.create("llvm"):
            # Perform the optimizations.
            mod = seq(mod)

有关Python和C ++中与pass相关的更多示例，请分别参考 tests / python / relay / test_pass_manager.py和 tests / cpp / relay_transform_sequential.cc。

牛牛存

关注

1
点赞
踩
2

收藏

觉得还不错? 一键收藏
0
评论
Relay传递基础架构

中继传递基础架构中继具有一系列优化过程，可改进模型的性能指标，例如特定设备的平均推断，内存占用量或功耗。有一套标准优化以及特定于机器学习的优化，包括恒定折叠，消除死代码，更改操作员布局和操作员融合等。这些步骤中的每一个都被构造为抽象上的中继到中继转换语法树（AST），使用遍历期间和/或之前收集的分析结果。但是，随着Relay的快速发展，对管理这些通道的更加系统和有效的方法的需求...
复制链接

扫一扫

专栏目录