Apache TVM 是一個深度的深度學習編譯框架，適用于 CPU、GPU 和各種機器學習加速芯片。更多 TVM 中文文檔可訪問 →https://tvm.hyper.ai/

Compiler Pass 是擴展 Relay 功能集及優化 Relay 程序的主要接口。通過編寫 compiler pass，用戶可以基于最終目標，修改 AST 或收集 AST 相關信息。事實上，Relay 內置的一些重要特性（如自動微分和類型推斷）都“標準”的 compiler pass。

整體來看，編寫 pass 包括兩個關鍵組成部分：

• 創建一個或多個遍歷程序的 C++ 類
• 將遍歷實現及其在 pass manager API 中的元數據包裝，從而方便與?Pass Infrastructure?輕松交互

首先，我們將概述編寫 compiler pass 的關鍵機制。然后通過 Relay 中常量折疊 pass 的具體示例進行演示。

AST 遍歷器（Traversers）

用于遍歷 Relay 程序的基類是?ExprFunctor。它提供的公共接口是一個?VisitExpr?方法，該方法接收一個表達式以及零個或多個參數，并返回某種類型的實例。擴展此類時，可以通過覆蓋每種表達式類型的?VisitExpr_?實現，來定義 AST 遍歷模式。

VisitExpr?和?VisitExpr_?之間的關系與調度有關。每個?VisitExpr_?定義都針對特定類型的表達式，但用戶無法每次都得知要訪問的節點類型。為了解決這個問題，ExprFunctor?提供了一個?VisitExpr?函數，將給定表達式路由轉換為?VisitExpr_?實例進而解決問題。盡管 C++ 已經提供了動態調度，但?ExprFunctor?定義了自己的虛表供?VisitExp?使用。通過定義虛表可以更好地控制調度。例如，定義一個在每次訪問之前都打印 “Here” 的?PrintVisitor?遍歷器，可以覆蓋?VisitExpr：

void PrintVisitor::VisitExpr(const Expr& expr) {
  std::cout << "Here" << std::endl;
  ExprFunctor::VisitExpr(expr);
}

ExprFunctor?本身是一個非常通用的類，這就是為什么更多時候你會擴展?ExprVisitor?或?ExprMutator。這些類擴展了?ExprFunctor，并提供了?VisitExpr_?的默認實現，這些實現捕獲了每種表達式類型的常見遍歷模式。有了這些默認的實現，開發者只需針對想要不同行為的表達式類型，提供覆蓋的實現。后續章節將針對每個子類進行詳細描述。

表達式訪問器（Expression Visitors）

ExprVisitor?不用于修改程序的pass，而是用于實施程序分析和收集信息的 pass。使用這個類，VisitExpr?和私有 counterparts 不會返回任何內容。此類提供的?VisitExpr_?實現只是訪問表達式的所有表達式字段。?IfNode?的默認實現如下所示：

void ExprVisitor::VisitExpr_(const IfNode* op) {
  this->VisitExpr(op->cond);
  this->VisitExpr(op->true_branch);
  this->VisitExpr(op->false_branch);
}

注意，這里調用的是?VisitExpr?而非?VisitExpr_，因此用戶可以使用?ExprFunctor?中的虛表進行路由。

如果要編寫一個?CallChecker?類來檢查程序中是否出現函數調用，只需擴展?ExprVisitor?并定義以下?VisitExpr_?方法：

void VisitExpr_(const CallNode* n) final {
  result_ = true;
}

其中?result_?是一個字段。在該示例中，無需在?CallNode?字段上進一步遞歸，因為?result_?已經為 true，原始表達式中包含一個調用。為了使該訪問器可用，可以采用以下方法：

bool Check(const Expr& expr) final {
  result_ = false;
  VisitExpr(expr);
  return result_;
}

以上就是全部操作。在調用 top-level 的遞歸之前，定義一個執行一些記錄的公有接口是很常見的操作。用戶也可以通過創建一個生成?CallChecker?實例，并在其上調用?Check?的獨立程序來進一步包裝 API，重要的是用盡可能少的資源用實現目標。

表達式修改器（Expression Mutators）

ExprMutator?用于以某種方式轉換程序的 pass。通過這個類，VisitExpr?及其對應的私有部分返回?Expr。此類提供的默認?VisitExpr_?實現訪問表達式的所有表達式字段，并將字段設置為訪問它們的結果。TupleGetItemNode?的默認實現如下所示：

Expr ExprMutator::VisitExpr_(const TupleGetItemNode* g) {
  auto t = this->Mutate(g->tuple);
  if (g->tuple == t) {
    return GetRef<Expr>(g);
  } else {
    return TupleGetItem(t, g->index);
  }
}

這里有幾點需要注意。首先，Mutate?是?ExprMutator?中?VisitExpr?的別名。其次，如果對?Mutate?的調用修改了?tuple?字段，則只返回一個新節點。這種更新的方法稱為功能更新，這樣做可以避免不必要的分配。

ExprMutator?有、而?ExprVisitor?沒有的一個功能，是用于緩存結果的內置?memo_?字段。ExprMutator?有一個記憶器（memoizer）這是合理的，因為用戶知道正在緩存哪些類型的結果（即?Expr），而?ExprVisitor?的訪問方法不返回任何內容。通常，當用戶要在?ExprVisitor?的子類中緩存結果時，需要自行定義緩存。

如果希望編寫一個?IfCollapser?類，用它的真實分支替換每個 if 語句，用戶將為?IfNode?覆蓋?VisitExpr_：

Expr ExprMutator::VisitExpr_(const IfNode* op) {
  return this->Mutate(op->true_branch);
}

注意：返回的表達式不一定是?IfNode，這是正常的，因為返回類型是?Expr。接下來創建一個公有接口：

Expr CollapseIfs(const Expr& expr) final {
  return this->Mutate(expr);
}

雖然使用這個修改器無需做任何記錄，但仍然鼓勵用戶將描述性方法作為接口。

示例：常量折疊

為了更好地理解編寫 pass 的過程，本部分將以常量折疊 pass（可在?src/relay/transforms/fold_constant.cc?中找到）作為示例進行講解。常量折疊 pass 相對簡單，且包含兩種類型的遍歷。

常量折疊涉及只包含常量的程序評估表達式（evaluating expression），然后用評估它們的結果替換這些表達式。此過程的目的是預加載可以進行的所有計算。為了實現這一點，常量折疊 pass 使用了一個訪問器（ConstantChecker）和一個修改器（ConstantFolder）。

ConstantChecker?訪問器

此訪問器用于檢查表達式是否為常量。在 Relay 中，用戶將?ConstantNode?或者只有常量字段的?TupleNode?的表達式定義為常量。

使用?memo_?字段從節點映射到它們是否為常量，并緩存這些結果。下面是?ConstantChecker?中的?VisitExpr_?定義。

void VisitExpr_(const ConstantNode* n) final {
  memo_[GetRef<Constant>(n)] = true;
}

void VisitExpr_(const TupleNode* n) final {
  bool result = true;
  for (const auto& field : n->fields) {
    if (!Check(field)) {
      result = false;
      break;
    }
  }
  memo_[GetRef<Tuple>(n)] = result;
}

用于協調這些定義的記錄是一個?Check?方法，它返回給定的表達式是否被認定為常量。

bool Check(const Expr& expr) {
  const auto it = memo_.find(expr);
  if (it != memo_.end())
    return it->second;
  VisitExpr(expr);
  return memo_[expr];
}

并不是所有遇到的節點都要修改?memo_；相反，用戶只有在遇到的節點有可能是常數時，才修改?memo_。當?memo_?不包含?expr?時，需要依賴默認的 false 值。

ConstantFolder?修改器

這個修改器執行了大部分的常量折疊過程，并在內部使用?ConstantChecker。在 Relay 中，常量折疊涉及三種節點類型：LetNode、TupleItemGetNode?和?CallNode。后續段落中將進行詳細講解。

Expr VisitExpr_(const LetNode* op) final {
  Expr value = this->Mutate(op->value);
  if (value.as<ConstantNode>()) {
    memo_[op->var] = value;
    return this->Mutate(op->body);
  } else {
    Var var = Downcast<Var>(this->Mutate(op->var));
    Expr body = this->Mutate(op->body);
    if (var.same_as(op->var) &&
        value.same_as(op->value) &&
        body.same_as(op->body)) {
      return GetRef<Expr>(op);
    } else {
      return Let(var, value, body);
    }
  }
}

在?LetNode?示例里，首先嘗試常量折疊綁定在表達式的值。如果可以，填充?memo_?并返回訪問主體的結果——本質上是將綁定的值傳到主體中的使用點。如果無法常量折疊綁定的值，可以參照默認的實現方法：

Expr VisitExpr_(const TupleGetItemNode* op) final {
  Expr res = ExprMutator::VisitExpr_(op);
  op = res.as<TupleGetItemNode>();
  if (const auto* tuple = op->tuple.as<TupleNode>()) {
    return tuple->fields[op->index];
  } else {
    return res;
  }
}

在?TupleItemGetNode?的例子里，需要檢查?op->tuple?字段是否為?TupleNode。如果是，我們將 get 元組替換為?op->index?指向的元組的字段。這樣做的原因是因為?op->tuple?可能被錯誤評估為一個元組。

Expr VisitExpr_(const CallNode* call) final {
  static auto op_stateful = Op::GetAttrMap<TOpIsStateful>("TOpIsStateful");
  Expr res = ExprMutator::VisitExpr_(call);
  call = res.as<CallNode>();
  // 我們不使用零參數的常量折疊函數。
  // 這是一個很有用的啟發式方法。
  // 例如折疊那些 shape=(4, 5) 是有害的。
  if (call->args.size() == 0) return res;
  const OpNode* op = call->op.as<OpNode>();
  if (op == nullptr) return res;
  // 跳過有狀態的算子。
  if (op_stateful.get(GetRef<Op>(op), false)) return res;
  bool all_const_args = true;
  for (Expr arg : call->args) {
    if (!checker_.Check(arg)) {
      all_const_args = false;
    }
  }
  if (all_const_args) {
    return ConstEvaluate(res);
  } else {
    return res;
  }
}

在?CallNode?示例中，首先使用?ExprMutator?的?VisitExpr_?來訪問調用，它將調用的所有字段都常量折疊了。之所以使用?ExprMutator::VisitExpr_?而不是?VisitExpr，是因為我們想要繞過虛表（以避免死循環）并使用?ExprMutator?提供的默認實現。只有當所有參數都是常量時，才評估調用（使用?ConstantChecker）。評估調用會產生一個值，因此這里使用輔助方法?ValueToExpr?，將評估的表達式放回 AST 中。

現在，我們為常量文件夾構造了一個更方便的接口?FoldConstant。FoldConstant?是?ConstantFolder?類之外的一個獨立函數，它負責接收表達式并在內部創建和使用?ConstantFolder?實例（其完整的定義在?src/relay/transforms/fold_constant.cc?中）。

用 Pass Manager 注冊 Pass

*注意：更多詳情請參閱?Pass Infrastructure?中的文檔。

編寫 AST 遍歷器后，用以下代碼可將 pass 注冊為 TVM API 端點：

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", {});
}

}  // 命名空間轉換

將上述代碼生成的?Pass?對象提供給 pass 基礎架構，可以使得 AST 遍歷應用于給定 Relay 模塊中的所有函數，這是常量折疊過程預期的行為（它應該盡可能折疊所有常量）。

函數?CreateFunctionPass?允許注冊 pass 的優化級別（在本例中為 2），可用于根據 pass 的一般實用性、 pass 名稱和 pass 中的任何依賴項將 pass 組合在一起。pass 的依賴項以列表形式給出，羅列了當前 pass 運行所必需的所有 pass 的結果。FoldConstant?沒有任何依賴，但是很多 Relay pass 確實依賴有類型信息，所以?InferType?是一個常見的依賴；其他的可能依賴于程序為 A-范式，通過?ToANormalForm?pass。

注意，PassContext?對象包含 pass 用于錯誤報告和配置選項的信息；?FoldConstant?不需要此信息，但其他 pass 可能會引用它們的?PassContext?對象。

現在可以通過 pass 基礎結構調用 pass 了，推薦為 pass 添加 Python 綁定，如以下代碼片段所示：

TVM_REGISTER_GLOBAL("relay._transform.FoldConstant")
.set_body_typed(FoldConstant);

通過以上方法定義了?Pass?對象后，就可以用 pass 基礎架構的?Sequential?結構來調用了。?Sequential?接收一個 pass 列表，并將其按順序應用于 Relay 模塊，從而獲得轉換后的模塊。例如，下面的代碼將?FoldConstant?和?ToANormalForm?pass 逐一應用于?mod?中的每個函數，并獲得一個新模塊。

seq = transform.Sequential([
    relay.transform.FoldConstant(),
    relay.transform.ToANormalForm()
])
new_mod = seq(mod)

更多注冊相關的內容，請查看?TVM Runtime 系統；pass 管理器接口相關的更多信息，請查看?Pass 基礎架構； Relay 的標準 pass 列表及實現方式，請分別查看?include/tvm/relay/transform.h?及?src/relay/transforms/。