MLIR-Code Documentation-Pass Infrastructure(Pass基础设施)

王棋

2023-12-01

Pass代表转换和优化的基础设施。本文档提供MLIR基础设施-Pass的概要说明已经使用方式。
阅读MLIR specification章节可以获得跟多MLIR的核心内容，例如IR结构与操作。
阅读MLIR Rewrites章节可以快速入门MLIR的图重写，如果一个转换涉及到匹配操作DAGs，这将是一个好的开始。

操作Pass(Operation Pass)

操作是MLIR中抽象和转换的主要单元。因此，Pass管理器被设计为能在操作实例的不同嵌套层次工作。Pass管理器的结构和嵌套层次的概念，将在下面进一步详细说明。MLIR中的所有Passes(Pass的实例)都来自于OperationPass，且遵守以下限制。不遵守任意一条都会导致在多线程或高级场景中出现有问题的行为。

能够修改当前正在操作的操作的外部引用或依赖的任意状态。包括在父块中增加或者删除操作、修改当前操作的属性(根据当前操作的订约)/操作数(operands)/结果/继任者(successors,子块？)。
修改未嵌套在当前正在操作的操作中的另一个操作的状态。
- 可能有其他线程同时对这些操作进行操作。
检查兄弟操作(应该是同一嵌套层次中的操作)的状态。
- 可能有线程正在并行的修改这些操作(指兄弟操作)。
- 允许检查祖先/父操作的状态。
跨调用runOnOperation时Pass的状态可变(就相当于每次Pass调用runOnOperation时，上一个操作的Pass状态都不会再被保留)。一个Pass可能会在许多不同的操作中运行且不保证执行顺序。
- 当使用多线程时，一个特定的Pass实例甚至可能不会执行IR中的所有操作。因此，一个Pass不应该依赖于运行所有操作(意思应该是Pass设计的时候不应该对操作存在耦合的关系，比方说，某个Pass必须依赖某个操作)。
维护任何全局可变状态。例如源代码中的静态变量，所有可变状态都应该由Pass的一个实例来维护。
必须是副本可构成的(用C++的概念解释的话，应该是必须有拷贝复制函数)。
- Pass管理器可能会创建一个pass的多个实例去并行的处理操作。创建一个操作Pass时，根据使用场景的不同，有两种不同的类型可供选择。

操作Pass:具体操作(OperationPass:Op-Specific)

一个Op-Specific的操作Pass明确的工作在给定的操作类型上。操作类型必须遵循Pass管理器设置的pass执行限制。
在定义一个Op-Specific Pass时，派生类必须遵循以下原则：

要继承CRTP类-OperationPass，且提供操作类型作为额外的模板参数。
覆盖虚函数void runOnOperation()

一个简单的Op-Specific Pass如下所示：

namespace {
struct MyFunctionPass : public PassWrapper<OperationPass<FuncOp>,
  MyFunctionPass> {
  void runOnOperation() override {
    FuncOp f = getOperation();
    f.walk([](Operation *inst) {
      ...
    });
  }
};
}

void registerMyPass() {
PassRegistration<MyFunctionPass>(
“flag-name-to-invoke-pass-via-mlir-opt”, ”Pass description here”);
}

操作Pass:操作无关(OperationPass:Op-Agnostic)

与操作无关的Pass被添加到Pass管理器之后，会在Pass管理器的操作类型上进行操作。这意味着这种Pass可以操作各种不同的操作类型。这种Pass通常是用于操作的接口和特性。例如：Common Sub-Expression Elimination和Inlining(内联).
创建这种Pass，派生类必须遵循以下限制：

继承于CRTP类-OperationPass
覆盖虚函数void runOnOperation()

一个简单的与操作无关的Pass例子如下：

struct MyOperationPass : public PassWrapper<OperationPass<>,
  MyOperationPass> {
  void runOnOperation() override {
    Operation *op = getOperation();
    ...
  }
};

依赖方言(Dependent Dialects)

方言必须在创建方言实体(操作、类型、属性)前加载到MLIRContext。方言必须在开始执行多线程Pass管道之前加载。因此，如果一个Pass在不能确保方言已经加载的情况下创建了实体，必须覆盖getDependentDialects()方法且在方言中明确的列出这些Pass的清单。

初始化(Initialization)

在某些情况下，一个Pass可能包含动态构造的状态，在Pass连续运行的过程中重新计算可能需要高昂的代价。比如，当使用PDL-based patterns时，在运行时期间被编译为字节码。在这样的情况下，一个pass可能通过重写以下钩子函数去初始化这个重要的状态：

LogicalResult initialize(MLIRContext *context)

一个完整Pass管道的每一次运行都会执行这个钩子函数，这意味着它不能访问在runOnOperation函数调用期间可用的状态。更具体地说，所有对MLIRContext的必要访问都应该通过所提供的contxt参数和利用“每次运行”状态来驱动的方法，比如getContext/getOperation/getAnalysis/等等，都不能用。如果在初始化过程中出现错误，Pass会发出一个错误诊断并返回一个会中断pass管道执行的failure()。

分析管理器(Analysis Management)

与转换Pass一样，分析也是一个重要的概念。它的概念和转换Pass相似，不同的是它计算特定操作的信息却不更改该它。分析(analyses)不是pass，而是独立的类，按需延迟计算并缓存结果，以避免不必要的重新计算。MLIR中的分析必须遵循以下限制：

提供一个接受Operation或者Operation和AnalysisManager &的构造函数。
- 提供的AnalysisManager &将被用作查询任何必要信息时的分析依赖项。
不能修改给定的操作。

分析可以提供额外的钩子函数，用于控制各种各样的行为：

bool isInvalidated(const AnalysisManager::PreservedAnalyses &)

给定一个保留分析集合，如果真的被失效这个分析将会返回true。这允许一个分析还没有被显示的标记为保留状态前，对失效案例进行更多的微调，也有可能因为分析集合的其他属性而被保留(或被失效)。如果有分析将其他分析作为依赖，当依赖分析被失效时，该分析也会被检查。

查询分析(Querying Analyses)

基础类OperationPass提供了工具用于查询和保存正在处理的当前操作的分析。

OperationPass类自动提供以下工具用于查询分析：
getAnalysis<>
- 获得一个当前操作的分析，如有必要，可以进行构造。
getCachedAnalysis<>
- 获得一个当前操作的分析，如果这个分析存在的话。
getCachedParentAnalysis<>
- 获得一个当前操作的父操作的分析，如果这个分析存在的话。
getCachedChildAnalysis<>
- 获得一个当前操作的子操作的分析，如果这个分析存在的话。
getChildAnalysis<>
- 获得一个当前操作的子操作的分析，如有必要，可以进行构造。

使用上面定义的Pass示例，让我们看一些示例：

struct MyOperationAnalysis {
  MyOperationAnalysis(Operation *op);
};

struct MyOperationAnalysisWithDependency {
  MyOperationAnalysisWithDependency(Operation *op, 
    AnalysisManager &am) {
    MyOperationAnalysis &otherAnalysis = 
      am.getAnalysis<MyOperationAnalysis>();
    ...
  }

  bool IsInvalidated(const AnalysisManager::PreservedAnalysis &pa) {
    return !pa.isPreserved<MyOperationAnalysisWithDependency>()
      || !pa.isPreserved<MyOperationAnalysis>();
  }
};

void MyOperationPass::runOnOperation() {
  MyOperationAnalysis &myAnalysis = getAnalysis<MyOperationAnalysis>();

  auto operationAnalysis = getCacheAnalysis<MyOperationAnalysis>();
  If(operationAnalysis)
    ...
  auto operationAnalysis =getCachedParentAnalysis<MyOperationAnalysis>();
  If(operationAnalysis)
    ...
}

保存分析(Preserving Analyses)

当分析被一个Pass请求时，这个分析会被构造。为了避免这个分析被再次请求时重复计算，这个分析还会被缓存。为了避免老旧的分析，所有的分析都被Pass假定为失效的。为了避免失效，Pass必须明确的标记分析，这样才能知道这个分析是需要被保留的。
所有的Pass类都自动的提供了以下保存分析的实用程序：

markAllAnalysesPreserved
markAnalysesPreserved<>

实例：

void MyOperationPass::runOnOperation() {
  markAllAnalysesPreserved();
  markAnalysesPreserved<MyAnalysis, MyAnalyses...>
}

Pass管理器

上面的章节介绍了不同类型的Pass和它们的不变量。本章节介绍Pass管理器的概念，与怎样使用它配置和调度一个Pass管道。与Pass管理器相关的主要有两个类，PassManager与OpPassManager。PassManager类充当最上层的入口点，包含整个Pass管道要用的各种各样的配置。OpPassManager类用于调度Pass运行在指定的嵌套等级。最上层的PassManager也可以作为OpPassManager。

OpPassManager

OpPassManager本质上就是一个在指定操作类型上执行的Pass集合。此操作类型必须符合以下要求：

必须被注册且被标记为IsolatedFromAbove。
- Pass不能修改当前正在处理的操作~~以及当前操作的上层操作？~~ 如果该操作不是独立的，有可能会无意的修改或遍历一个不应该执行的操作的SSA使用列表。

Pass管理器可以通过方法addPass来增加Pass。该Pass要么与OpPassManager拥有相同的操作类型，要么就是一个与操作无关(指的是和操作类型无关)的Pass。

创建一个OpPassManager通常用另一个已存在的OpPassManager使用方法nest<>显示地嵌套一个管道。该方法需要一个Pass管理器即将操作的操作类型。在顶层，一个PassManager可以充当一个OpPassManager。这种意义上的嵌套，对应于IR的Regions的structural嵌套。
例子如下，.mlir文件

module {
  spv.module "Logical" "GLSL450" {
    fun @foo() {
      ...
    }
  }
}

嵌套结构如下：

`module`
  `spv.module`
    `function`

下面是一个构建管道的示例，该管道在上述结构上运行：

PassManager pm(ctx);
PassManager pm(ctx, "module");
pm.addPass(std::make_unique<MyModulePass>());
OpPassManager &nestedModulePM = pm.nest<spirv::ModuleOp>();
nestedModulePM.addPass(std::make_unique<MySPIRVModulePass>());
OpPassManager &nestedFunctionPM = nestedModulePM.nest<FuncOp>();
nestedFunctionPM.addPass(std::make_unique<MyFunctionPass>());
ModuleOp m = ...;
if (failed(pm.run(m)))
  ...

上面的Pass管理器包含以下管道结构

OpPassManager<ModuleOp>
  MyModulePass
  OpPassManager<spirv::ModuleOp>
    MySPIVModulePass
    OpPassManager<FuncOp>
      MyFunctionPass

然后，这些管道在同一时间每次只运行一个操作。例如，在FuncOp上给定一系列连续的Pass，它将会在第一个函数上执行所有操作，然后在第二个函数上执行所有操作，以此类推，直到整个程序在这些Pass上都运行通过。这么做(或者说设计这样的模式)能提供以下几个好处：

这改进了编译器的缓存行为，(不明白)因为它一次只涉及一个函数，而不是遍历整个程序。
这改进了多线程性能，通过减少需要调度的作业数量，也提高每项作业的效率。一个完整的函数管道可以异步的运行在每个函数上。

动态Pass管道(Dynamic Pass Pipelines)

在某些情况下，Pass管道与另一些Pass一起运行可能是很有用的，允许配置或者过滤当前正在操作的操作的不变量。举个例子，Inliner Pass可能希望在内联的同时运行函数中简化Pass，以生成更好的成本模型，并提供更佳的内联。要启用此功能，Pass可以运行在任意一个当前操作正在操作的OpPassManager中，或者是当前操作通过方法LogicalResult Pass::runPipeline(OpPassManager &, Operation *)创建的嵌套。该方法会返回动态管道成功还是失败，返回结果与顶层方法PassManager::run类似。一个简单的例子如下所示：

void MyModulePass::runOnOperation() {
  ModuleOp module = getOperation();
  if (hasSomeSpecificProperty(module)) {
    OpPassManager dynamicPM("moduyle");
    ...;
    if (failed(runPipeline(dynamicPM, module)))
      return signalPassFailure();
  }
}

注意：虽然上面的动态管道是在runOnOperation方法中构造的，但这是不必要的，而且管道应该在可能的情况下缓存，因为OpPassManager类可以安全地复制构造。
当一个Pass管道只能以嵌套的方式运行时，即该嵌套管道不能与主Pass管道的其余部分一起静态调度时，应使用本节中描述的机制。更具体地说，PassManager通常不需要在Pass中构造。使用runPipeline还可以确保所有分析、仪器仪表和其他过程管理器相关组件都与正在执行的动态管道集成。

第一次阶段性整理。

特定实例的Pass选项？(Instance Specific Pass Options)

MLIR为Pass提供了一个内置机制，以指定选项配置其行为。这些选项在Pass构建时独立地被每一个Pass实例解析。选项通过Option<>和ListOption<>类定义，并遵循LLVM command line 标志定义规则。下面看一个小的例子：

struct MyPass ... {
  MyPass() = defult;
  MyPass(cong MyPass& pass){}
  Option<int> exampleOption{*this, "flag-name", llvm::cl::desc("...")};
  ListOption<int> exampleListOption{*this, "list-flag-name",
    llvm::cl::desc("...")};
}

PassPipelineRegistration模板提供了一个额外的模板参数–可选的Option结构定义。这个结构需要继承于mlir::PassPipelineOptions，且包含想要的管道选项。使用PassPipelineRegistration时，构造函数接受一个形如void(OpPassManager &pm, const MyPipelineOptions&)的函数签名，该函数根据所提供的选项来构建Pass，并将这些选项传递给pm(OpPass管理器)：

struct MyPipelineOptions : public PassPipelineOptions {
  Option<int> exampleOption{*this, "flag-name", llvm::cl::desc("...")};
  ListOption<int> exampleListOption{*this, "list-flag-name",
                                    llvm::cl::desc("...")};
};

void registerMyPasses() {
  PassPipelineRegistration<MyPipelineOptions>(
    "example-pipeline", "Run an example pipeline.",
     [](OpPassManager &pm, const MyPipelineOptions &pipelineOptions) {
     //Initialize the pass manager.
  });
}

Pass统计信息(Pass Statistics)

统计信息是跟踪编译器正在做什么以及各种转换的效果如何的一种方法。用于查看特定转换对特定输入的影响以及它们触发的频率通常很有用。Pass统计信息特定于每个Pass实例，这允许查看将特定转换放置在Pass管道中的特定位置时的效果。例如，他们帮助回答“如果我在这里再次运行CSE会发生什么？”这样的问题。

统计信息可以通过Pass::Statistic类被加入到Pass。该类构造函数需要的参数：父Pass，名称，描述。该类的作用就像原子无符号整数，可以相应地递增和更新。这些统计信息依赖于与llvm::Statistic相同的基础结构，因此具有类似的使用约束。收集统计信息可以通过Pass管理器使用PassManager::enableStatistics来转储，也可以使用命令行-pass-statistics和-pass-statistics-display。一个简单的例子如下所示：

struct MyPass ... {
  MyPass() = default;
  MyPass(const MyPass& pass) {}
  Statistic exampleStat{this, "exampleStat", "An example statistic"};
  void runOnOperation() {
    ...
    ++exampleStat;
    ...
  }
};

统计信息可以在两种类型的视图中聚合：
管道视图，对pass管理器结构进行建模，也是默认视图：

$ mlir-opt -pass-pipeline='func(my-pass,my-pass)' foo.mlir -pass-statistics

===-------------------------------------------------------------------------===
                         ... Pass statistics report ...
===-------------------------------------------------------------------------===
'func' Pipeline
  MyPass
    (S) 15 exampleStat - An example statistic
  VerifierPass
  MyPass
    (S)  6 exampleStat - An example statistic
  VerifierPass
VerifierPass

列表视图，将同一种Pass类型实例的统计信息聚合在一起：

$ mlir-opt -pass-pipeline='func(my-pass, my-pass)' foo.mlir -pass-statistics -pass-statistics-display=list

===-------------------------------------------------------------------------===
                         ... Pass statistics report ...
===-------------------------------------------------------------------------===
MyPass
  (S) 21 exampleStat - An example statistic

疑问？

Pass注册(Pass Registration)

在各种Pass类型的定义例子中简要的展示了PassRegistration类。此机制允许注册Pass类，以便在文本Pass管道描述(textual pass pipelline description?)中创建它们。一个简单的注册例子如下所示：

void registerMyPass() {
  PassRegistration<MyPass>("argument", "description");
}

Mypass是派生Pass类的名字
"argument"是在文本格式中该Pass的引用参数
"description"是该Pass的简要描述

对于Pass来讲是不能被默认构建的，PassRegistration接受一个可选的第三参数，可提供回调函数来创建Pass:

void registerMyPass() {
  PassRegistration<MyParametricPass>(
    "argument", "description",
    []() -> std::unique_ptr<Pass> {
      std::unique_ptr<Pass> p = std::make_unique<MyParametricPass>
                                (/*options*/);
      return p;
    });
}

这种注册变体是有用的，例如，接受来自命令行参数的Pass配置并将其传递给Pass构造函数。

注意：要确保Pass是可复制构造的，以不共享数据的方式？因为Pass管理器可能会创建Pass副本用于并行。
疑问？