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

隨著深度學習工作負載所針對的硬件設備數量不斷增加，用戶在各種設備上實現高性能所需的知識也在不斷增加。為了讓數據科學家在開發新模型時不必擔心性能問題，硬件廠商或是基于一些常見的深度學習算子，提供 MKLDNN 或 cuDNN 等庫，或是提供 TensorRT 等框架，讓用戶按照某種方式描述模型，從而提高模型性能。

然而，用戶在嘗試使用新的庫或設備時，必須學習新的編程接口。因此，一個統一的編程接口變得越來越重要：1）讓所有用戶及硬件廠商信息同步，2）提供一個可行的解決方案，讓特定硬件或庫只支持具有極高性能的、廣泛使用的算子，不受支持的算子則回退到 CPU/GPU 等通用設備。

本開發手冊演示了硬件廠商如何輕松實現自己的 Codegen，并將其注冊為 Relay 后端編譯器，從而支持自己的硬件設備/庫。本手冊涵蓋了兩種基于不同計算圖的 codegen：

1. 希望生成 C 代碼。

如果你的硬件已經具備了一個高度優化的 C/C++ 庫，如對于 CPU 而言的 Intel CBLAS/MKL 庫，或針對 GPU 而言的 NVIDIA CUBLAS 庫，那么本節內容非常適合你。幸運的是，C 源代碼模塊與 TVM runtime 模塊完全兼容，這意味著生成的代碼可以由任何具有適當編譯標志的 C/C++ 編譯器編譯，因此用戶只需實現一個能為子圖生成 C 代碼的 codegen，并將 C 源代碼模塊集成到 TVM runtime 模塊中。下一節內容講詳細演示如何為硬件實現 C codegen。

2. 希望生成任意計算圖。

有時候，硬件可能需要其他形式的計算圖如 JSON。這種情況下，用戶不僅要實現一個 codegen，還要實現一個自定義 TVM runtime 模塊，從而使得 TVM runtime 知道如何執行這個計算圖。如果你的硬件已經擁有完整的計算圖執行引擎（graph execution engine），如適用于 GPU 的 TensorRT，那么該解決方案對你而言非常具有參考價值。

完成 codegen 和 runtime 后，可以讓客戶借助你的自定義標簽，對模型進行注釋并加以利用。終端用戶如何注釋和啟動特定 codegen 的教程，將在后續進行補充。

實現 C Codegen

在這一部分中，我們將演示如何借助預實現的算子函數，生成 C 代碼的 codegen。簡單起見，本示例 codegen 不依賴于第三方庫。相反，我們在 C 中手動實現了兩個宏：

#define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)         \
    extern "C" void p_ID_(float* a, float* b, float* out) { \
        for (int64_t i = 0; i < p_DIM1_; ++i) {             \
            out[i] = a[i] p_OP_ b[i];                       \
        }                                                   \
    }

#define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \
    extern "C" void p_ID_(float* a, float* b, float* out) {   \
        for (int64_t i = 0; i < p_DIM1_; ++i) {               \
            for (int64_t j = 0; j < p_DIM2_; ++j) {           \
                int64_t k = i * p_DIM2_ + j;                  \
                out[k] = a[k] p_OP_ b[k];                     \
            }                                                 \
        }                                                     \
    }

使用這兩個宏，可以為一維和二維張量生成二元算子（binary operator）。例如，給定如下所示的子圖，假設所有輸入都是 shape 為（10, 10）的二維張量：

c_compiler_input0
       |
      add <-- c_compiler_input1
       |
    subtract <-- c_compiler_input2
       |
    multiply <-- c_compiler_input3
       |
      out

我們的目標是生成以下可編譯代碼來執行子圖：

#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include <cstdint>
#include <cstring>
#include <iostream>

#define GCC_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)           \
  extern "C" void p_ID_(float* a, float* b, float* out) { \
    for (int64_t i = 0; i < p_DIM1_; ++i) {               \
      out[i] = a[i] p_OP_ b[i];                           \
    }                                                     \
  }

#define GCC_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \
  extern "C" void p_ID_(float* a, float* b, float* out) { \
    for (int64_t i = 0; i < p_DIM1_; ++i) {               \
      for (int64_t j = 0; j < p_DIM2_; ++j) {             \
        int64_t k = i * p_DIM2_ + j;                      \
        out[k] = a[k] p_OP_ b[k];                         \
      }                                                   \
    }                                                     \
  }

// 注 1
GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);
GCC_BINARY_OP_2D(gcc_0_1, -, 10, 10);
GCC_BINARY_OP_2D(gcc_0_2, +, 10, 10);

// 注 2
extern "C" void gcc_0_(float* gcc_input0, float* gcc_input1,
                       float* gcc_input2, float* gcc_input3, float* out) {
  float* buf_0 = (float*)malloc(4 * 100);
  float* buf_1 = (float*)malloc(4 * 100);
  gcc_0_2(gcc_input0, gcc_input1, buf_0);
  gcc_0_1(buf_0, gcc_input2, buf_1);
  gcc_0_0(buf_1, gcc_input3, out);
  free(buf_0);
  free(buf_1);
}

// 注 3
extern "C" int gcc_0_wrapper(DLTensor* arg0, DLTensor* arg1, DLTensor* arg2,
                             DLTensor* arg3, DLTensor* out) {
  gcc_0_(static_cast<float*>(arg0->data), static_cast<float*>(arg1->data),
         static_cast<float*>(arg2->data), static_cast<float*>(arg3->data),
         static_cast<float*>(out->data));
  return 0;
}
TVM_DLL_EXPORT_TYPED_FUNC(gcc_0, gcc_0_wrapper);

這里詳細介紹一下上面代碼里的注釋：

• 注1：子圖中三個節點的函數實現。
• 注2：通過分配中間數組（intermediate buffer）并調用相應函數來執行子圖的函數。
• 注3：TVM runtime 兼容的包裝函數。它接收一個輸入張量列表和一個輸出張量（最后一個參數），并將其轉換為正確的數據類型，調用注2 中描述的子圖函數。此外，TVM_DLL_EXPORT_TYPED_FUNC?是一個 TVM 宏，它通過將所有張量打包到?TVMArgs?來生成另一個函數?gcc_0，該函數具有統一的函數參數。因此，TVM runtime 可以直接調用?gcc_0?來執行子圖，無需其他操作。生成上述代碼后，TVM 能夠將其與計算圖的其余部分一起編譯并導出單個庫以進行部署。

在本節的其余部分，我們將逐步創建一個 codegen，來實現上述代碼。你的 codegen 必須位于?src/relay/backend/contrib/<your-codegen-name>/。在這個例子中，我們將 codegen 命名為 “codegen_c”，并將其放在?/src/relay/backend/contrib/codegen_c/?目錄下。你可以隨時查看這個文件，了解完整的實現過程。

具體來說，我們將在這個文件中實現兩個類，兩個類的關系如下：

            subgraph                                subgraph
TVM backend -----------------------------> CSourceCodegen -------------> CodegenC
       ^                                       |    ^                       |
       |                                       |    |                       |
       ----------------------------------------      ------------------------
          generated C source runtime module              generated C code

當 TVM 后端發現 Relay 計算圖中的函數（子圖），用注冊的編譯器標簽（本例中為?ccompiler）進行了注釋時，TVM 后端就會調用?CSourceCodegen?并傳遞子圖。?CSourceCodegen?的成員函數?CreateCSourceModule?將：

1）為子圖生成 C 代碼；

2）將生成的 C 代碼包裝到 C source runtime 模塊中，以便 TVM 后端進行編譯和部署。

特別是，C codegen 對?CodegenC?類是透明的，因為它提供了許多有用的實用程序來簡化 codegen 實現。下面的章節將自下而上實現這兩個類。

實現 CodegenC

在?src/relay/backend/contrib/codegen_c/codegen.cc?中，首先在?tvm.relay.contrib?的命名空間下創建一個 codegen 類骨架：

#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
#include <tvm/relay/type.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/object.h>

#include <fstream>
#include <sstream>

#include "codegen_c.h"

namespace tvm {
namespace relay {
namespace contrib {

class CodegenC : public ExprVisitor, public CodegenCBase {
  public:
    explicit CodegenC(const std::string& id) { this->ext_func_id_ = id; }

    void VisitExpr_(const VarNode* node) { ; }
    void VisitExpr_(const CallNode* call) final { ; }
    std::string JIT() { ; }

  private:
    /*! \brief The function id that represents a C source function. */
    std::string ext_func_id_ = "";
    /*! \brief The index of a wrapped C function. */
    int func_idx = 0;
    /*! \brief The index of allocated buffers. */
    int buf_idx_ = 0;
    /*! \brief The arguments of a C compiler compatible function. */
    std::vector<std::string> ext_func_args_;
    /*! \brief The statements of a C compiler compatible function. */
    std::vector<std::string> ext_func_body;
    /*! \brief The declaration statements of a C compiler compatible function. */
    std::vector<std::string> func_decl_;
    /*! \brief The declaration statements of buffers. */
    std::vector<std::string> buf_decl_;
    /*! \brief The name and index pairs for output. */
    std::vector<std::pair<std::string, int>> out_;
}

CodegenC?類繼承了兩個類：?ExprVisitor?提供遍歷子圖的能力，然后收集所需的信息并生成子圖函數，例如?gcc_0_。

CodegenCBase?提供了生成包裝函數的能力和實用程序，例如上例中的?gcc_0。可以看出，我們只需要在這個 codegen 類中實現三個函數就可以了。

算子的代碼生成

首先實現?VisitExpr_(const CallNode* call)。該函數在遍歷子圖時會訪問所有調用節點。每個調用節點都包含一個我們想要卸載（offload）到硬件中的算子。因此，我們需要按照拓撲順序生成具有正確算子的相應 C 代碼。完整實現過程如下：

1. 生成函數聲明

示例結果：GCC_BINARY_OP_2D(gcc_0_0, *, 10, 10);

要生成函數聲明，如上所示，我們需要：

1）函數名（例如?gcc_0_0）

2）算子的類型（例如?*?）

3）輸入張量 shape（例如?(10, 10)?）

這些信息可以從?CallNode?輕松獲取：

std::ostringstream macro_stream;
std::ostringstream decl_stream;
std::ostringstream buf_stream;

// Generate a unique function name you like.
std::string func_name = ext_func_id_ + "_" + std::to_string(func_idx++);

// Make function declaration string.
macro_stream << "CSOURCE_BINARY_OP_" << call->args.size() << "D(" << func_name << ", ";

// Check the operator type.
if (IsOp(call, "add")) {
  macro_stream << "+";
} else if (IsOp(call, "subtract")) {
  macro_stream << "-";
} else if (IsOp(call, "multiply")) {
  macro_stream << "*";
} else {
  LOG(FATAL) << "Unrecognized op";
}

// Extract the input tensor shape.
auto in_shape = GetShape(call->args[0]->checked_type());
for (size_t i = 0; i < in_shape.size(); ++i) {
  macro_stream << ", " << in_shape[i];
}
macro_stream << ");";
func_decl_.push_back(macro_stream.str());

可以看出，我們將生成的代碼推送到類成員變量?func_decl_?中。這意味著在我們完成遍歷整個子圖之后，我們已經收集了所有必需的函數聲明，我們唯一需要做的就是用 GCC 編譯它們。?VisitExpr_(const CallNode* call)?的其余實現也遵循這個概念。

2. 生成函數調用

示例結果：gcc_0_0(buf_1, gcc_input3, out);

生成函數聲明后，我們需要生成一個具有正確輸入和輸出的函數調用。要想知道調用這個函數時應該放置哪些輸入或數組，必須訪問它的參數：

bool first = true;
decl_stream << func_name << "(";
for (size_t i = 0; i < call->args.size(); ++i) {
  VisitExpr(call->args[i]); // 注 1
  for (auto out : out_) {
    if (!first) {
      decl_stream << ", ";
    }
    first = false;
    decl_stream << out.first;
  }
}
// 注 2

同樣，重點介紹一下上述代碼中的注釋：

注1：VisitExpr(call->args[i])?是訪問當前函數參數的遞歸調用。參數可以是另一個節點的輸出或輸入張量。在該示例中，需要確保每個節點在離開訪問器之前，都更新一個類變量?out_。圖解如下：

 arg_node                 arg_node <- Visit arg (Note 1)       arg_node
     |                        |                                    |
 curr_node <- Process      curr_node                            curr_node <- Put "buf_0" as an input buffer

(a) out_ = {}            (b) out_ = {}                   (c) out_ = {("buf_0", 20)}

從上圖中可以看出，類變量?out_?在訪問參數節點前是空的，它被填充了?arg_node?輸出數組的名稱和大小。因此在完成對參數節點的訪問時，可以通過查看?out_?得知應該放置的正確輸入數組。本節末尾以及下一節中，我們將介紹如何更新?out_。

注2：你可能注意到，我們在這一步沒有關閉函數調用字符串。當前函數調用字符串看起來像：gcc_0_0(buf_1, gcc_input3。這是因為我們沒有將最后一個參數（如 output）放入此調用中。函數調用的輸出可以是分配的臨時數組或子圖輸出張量。簡單起見，在本例中我們為每個調用節點都分配老一個輸出數組（下一步），并將最后一個數組中的結果復制到了輸出張量。

3. 生成輸出數組（output buffer）

示例結果：float buf_0 = (float)malloc(4 * 100);

如上一步所述，除了子圖輸入和輸出張量外，還需要數組來保存中間結果。為了生成數組，我們提取 shape 信息，以確定數組的類型和大小：

// 這個例子僅支持單個輸出。
auto type_node = call->checked_type().as<TensorTypeNode>();
ICHECK(type_node != nullptr && runtime::TypeMatch(type_node->dtype, kDLFloat, 32))
      << "Only support single output tensor with float type";

// 生成一個唯一的數組名字。
std::string out = "buf_" + std::to_string(buf_idx_++);

// 提取 shape 作為數組大小。
auto out_shape = GetShape(call->checked_type());
int out_size = 1;
for (size_t i = 0; i < out_shape.size(); ++i) {
  out_size *= out_shape[i];
}

// 分配數組并推送至數組聲明
buf_stream << "float* " << out << " = (float*)std::malloc(4 * " << out_size << ");";
buf_decl_.push_back(buf_stream.str());

分配了輸出數組之后，現在可以關閉函數調用字符串，并將生成的函數調用推送到類變量?ext_func_body。

decl_stream << ", " << out << ");";
ext_func_body.push_back(decl_stream.str());

4. 更新輸出數組

為了使得下一個節點（接受當前調用節點的輸出，作為其輸入）知道它應該使用哪個數組，我們需要在離開這個訪問函數之前更新類變量?out_：

out_.clear();
out_.push_back({out, out_size});

恭喜！到這一步我們已經完成了這個類中最困難的函數。接下來的兩節中，我們將進一步完善這個函數的功能。

輸入變量的代碼生成

回想一下，我們通過訪問調用節點的參數（上一節中的第 2 步）收集了輸入數組信息，并處理了參數是另一個調用節點的情況（第 4 步）。本節我們將以?VarNode?為例，演示如何處理其他節點。

VarNode?表示模型中的輸入張量。它非常重要的一點就是名稱提示（例如，data、weight?等）。訪問?VarNode?時，只需更新類變量?out_?傳遞名稱提示，后代（descendant）調用節點就可以生成正確的函數調用。

void VisitExpr_(const VarNode* node) {
  ext_func_args_.push_back(node->name_hint());
  out_.clear();
  out_.push_back({node->name_hint(), 0});
}

注意：在這個例子中，我們假設要卸載的子圖只有調用節點和變量節點。如果子圖包含其他類型的節點，如?TupleNode，那么你也需要訪問它們并繞過輸出數組信息。

Code Emitting

Codegen Class 的最后一部分是?JIT?函數，它為子圖 emit 一個 C 函數，并將剛生成的 C 代碼作為函數體。注意，除了在前幾節中生成的子圖函數外，還需要一個具有統一參數的 wrapper 函數，供 TVM runtime 調用和傳遞數據。幸運的是，我們繼承的基類已經提供了一個實現，即?JitImpl，來生成該函數。調用?JitImpl的方式如下：

JitImpl("gcc_0" /* Subgraph symbol (ID) */,
        {"gcc_input0", "gcc_input1", "gcc_input2", "gcc_input3"} /* Input arguments */,
        {"float *buf_0 = (float*)malloc(4 * 20)", ...} /* Buffer allocations */,
        {"gcc_0_2(gcc_input0, gcc_input1, buf_0);"} /* Function body */,
        {"out"} /* Output */);

上述調用將生成三個函數（一個來自 TVM wrapper 宏）：

1. 子圖函數?gcc_0_（函數名末尾多了一個下劃線）以及為執行子圖而生成的所有 C 代碼；
2. 帶有?DLTensor?參數列表的 wrapper 函數?gcc_0__wrapper_?，將數據轉換為正確的類型并調用?gcc_0_
3. TVM runtime 兼容函數?gcc_0?具有 TVM 統一函數參數，可解包 TVM 打包張量并調用?gcc_0__wrapper_

因此，在?JIT?實現中唯一要做的，就是將生成的所有子圖函數代碼傳遞給?JitImpl：

std::string JIT() {
  // Write function macros
  for (auto decl : func_decl_) {
    code_stream_ << decl << "\n";
  }
  return JitImpl(ext_func_id_, ext_func_args_, buf_decl_, ext_func_body, out_);
}

傳遞的所有變量（ext_func_id?等）都是類變量，并在遍歷子圖時被填充。

實現 CSourceCodegen

創建一個類并實現所需功能，注意：需要繼承自?CSourceModuleCodegenBase：

class CSourceCodegen : public CSourceModuleCodegenBase {
 public:
  // 傳遞一個子圖函數, 并生成 C 代碼。
  void GenCFunc(const Function& func) { ; }

  // 使用 GenCFunc 來生成 C 代碼并將它包裝成一個 C 源模塊。
  runtime::Module CreateCSourceModule(const NodeRef& ref) override { ; }

 private:
  std::ostringstream code_stream_;
};

實現 GenCFunc

GenCFunc?只是簡單地使用我們剛剛實現的?CodegenC?來遍歷一個 Relay 函數（子圖），得到生成的 C 代碼。內置函數?GetExtSymbol?在 Relay 函數中檢索唯一的符號名稱（例如?gcc_0），注意：必須將其用作 C 函數名稱，因為該符號將用于 DSO 運行查找。

void GenCFunc(const Function& func) {
  ICHECK(func.defined()) << "Input error: expect a Relay function.";

  // 記錄運行查找的外部符號。
  auto sid = GetExtSymbol(func);

  CodeGenC builder(sid);
  builder.VisitExpr(func->body);
  code_stream_ << builder.JIT();
}

實現 CreateCSourceModule

此函數為外部庫創建了一個 runtime 模塊。本事例中，我們創建了一個可以直接被編譯并與 TVM 生成的 DSOModule 鏈接在一起的 CSourceModule。CodegenC?實現之后，再實現這個功能就比較簡單了：

runtime::Module CreateCSourceModule(const NodeRef& ref) override {
  // 創建頭文件
  code_stream_ << "#include <cstdint>\n";
  code_stream_ << "#include <iostream>\n";
  code_stream_ << "#include <cstdlib>\n";
  code_stream_ << "#include <stdio.h>\n";
  code_stream_ << "#include <cstring>\n";
  code_stream_ << "#include <tvm/runtime/c_runtime_api.h>\n";
  code_stream_ << "#include <dlpack/dlpack.h>\n";

  // 為算子定義添加一些公共宏。
  const char* operator_macro = R"op_macro(
  #define CSOURCE_BINARY_OP_1D(p_ID_, p_OP_, p_DIM1_)       \
    extern "C" void p_ID_(float* a, float* b, float* out) { \
      for (int64_t i = 0; i < p_DIM1_; ++i) {               \
        out[i] = a[i] p_OP_ b[i];                           \
      }                                                     \
    }

  #define CSOURCE_BINARY_OP_2D(p_ID_, p_OP_, p_DIM1_, p_DIM2_)  \
    extern "C" void p_ID_(float* a, float* b, float* out) {     \
      for (int64_t i = 0; i < p_DIM1_; ++i) {                   \
        for (int64_t j = 0; j < p_DIM2_; ++j) {                 \
          int64_t k = i * p_DIM2_ + j;                          \
          out[k] = a[k] p_OP_ b[k];                             \
        }                                                       \
      }                                                         \
    }
  )op_macro";

  code_stream_ << operator_macro << "\n\n";

  // 為子圖生成 C 代碼。
  if (ref->IsInstance<FunctionNode>()) {
    GenCFunc(Downcast<Function>(ref));
  } else if (ref->IsInstance<relay::ModuleNode>()) {
    relay::Module mod = Downcast<relay::Module>(ref);
    for (const auto& it : mod->functions) {
      GenCFunc(Downcast<Function>(it.second));
    }
  } else {
    LOG(FATAL) << "The input ref is expected to be a Relay function or module"
               << "\n";
  }

  // 創建一個 CSourceModule
  const auto* pf = runtime::Registry::Get("module.csource_module_create");
  ICHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
  return (*pf)(code_stream_.str(), "cc");
}

注冊 CodegenC

最后一步是將 codegen 注冊到 TVM 后端。首先實現一個簡單的函數，調用 codegen 并生成一個 runtime 模塊：

runtime::Module CCompiler(const NodeRef& ref) {
  CSourceCodegen csource;
  return csource.CreateCSourceModule(ref);
}

接下來將此函數注冊到 TVM 后端：

TVM_REGISTER_GLOBAL("relay.ext.ccompiler").set_body_typed(CCompiler);

其中?ccompiler?是一個自定義標簽，它告知 TVM 這是用?ccompiler?注釋子圖時，應該用來生成和卸載子圖的 codegen。

最后，設置一個 CMake 配置標志，只包含客戶的編譯器。首先創建一個 cmake 文件：cmake/modules/contrib/CODEGENC.cmake：

if(USE_CODEGENC)
  file(GLOB CSOURCE_RELAY_CONTRIB_SRC src/relay/backend/contrib/codegen_c/codegen.cc)
  list(APPEND COMPILER_SRCS ${CSOURCE_RELAY_CONTRIB_SRC})
endif(USE_CODEGENC)

用戶在使用?config.cmake?配置 TVM 時，可以自行決定是否配置編譯器：

set(USE_CODEGENC ON)

為表征（Representation）實現 Codegen

盡管我們已經演示了如何實現 C codegen，但用戶硬件可能還需要其他形式的計算圖表征（Graph Representation），如 JSON。在這種情況下，用戶可以通過修改?CodegenC?類，生成自己的計算圖表征，并實現一個自定義 runtime 模塊，告訴 TVM runtime 如何執行這個計算圖表征。

簡單起見，本指南中定義了一個名為 “ExampleJSON” 的計算圖表征。 ExampleJSON 并不是 JSON，而是沒有控制流的計算圖的簡單表示。例如，假設有以下名為?subgraph_0?的子圖：

input0
   |
  add <-- input1
   |
subtract <-- input2
   |
multiply <-- input3
   |
  out

那么這個子圖的 ExampleJON 看起來類似：

subgraph_0
  input 0 10 10
  input 1 10 10
  input 2 10 10
  input 3 10 10
  add 4 inputs: 0 1 shape: 10 10
  sub 5 inputs: 4 2 shape: 10 10
  mul 6 inputs: 5 3 shape: 10 10

input?關鍵字聲明一個輸入張量及其 ID 和 shape；其他語句用?<op> <output ID> inputs: [input ID] shape: [shape]?語法描述了其計算過程。

在本節中，我們試圖實現以下自定義 TVM runtime 模塊，來執行 ExampleJSON 計算圖。

runtime::Module ExampleJsonCompiler(const NodeRef& ref) {
    ExampleJsonCodeGen codegen(ref);
    std::string code = codegen.gen(); // 注 1
    const auto* pf = runtime::Registry::Get("module.examplejson_module_create"); // 注 2
    ICHECK(pf != nullptr) << "Cannot find ExampleJson module to create the external runtime module";
    return (*pf)(code);
}
TVM_REGISTER_GLOBAL("relay.ext.examplejsoncompiler").set_body_typed(ExampleJsonCompiler);

注1：稍后我們將實現一個自定義 codegen，通過取一個子圖來生成一個 ExampleJSON 代碼字符串。

注2：此行獲取了一個用于創建自定義 runtime 模塊的函數的指針。可以看到它采用剛剛生成的 ExampleJSON 格式的子圖代碼，并對一個 runtime 模塊進行了初始化。

后續章節中，我們將介紹 1）如何實現?ExampleJsonCodeGen?和 2）如何實現和注冊?examplejson_module_create。

實現 ExampleJsonCodeGen

與 C codegen 類似，從?ExprVisitor?派生?ExampleJsonCodeGen?以訪問器模式進行子圖遍歷。另一方面，因為不會用到 TVM C++ wrapper，所以不必繼承?CodegenCBase。 codegen 類實現如下：

#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
#include <tvm/relay/type.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/object.h>

#include <fstream>
#include <sstream>

namespace tvm {
namespace relay {
namespace contrib {

class ExampleJsonCodeGen : public ExprVisitor {
  public:
    explicit ExampleJsonCodeGen();

    // 注 1
    void VisitExpr_(const VarNode* node) { /* Skip in this example. */ }
    void VisitExpr_(const CallNode* call) final { /* Skip in this example. */ }

    // 注 2
    std::string gen(NodeRef& ref) {
        this->code = "";
        if (ref->IsInstance<FunctionNode>()) {
            this->visit(Downcast<Function>(ref));
        } else if (ref->IsInstance<relay::ModuleNode>()) {
            relay::Module mod = Downcast<relay::Module>(ref);
            for (const auto& it : mod->functions) {
                this->visit(Downcast<Function>(it.second));
            }
        } else {
            LOG(FATAL) << "The input ref is expected to be a Relay function or module";
        }
        return this->code;
    }

  private:
      /*! \brief The function id that represents a C source function. */
     std::string code;
}

注1：再次實現相應的 visitor 函數，以生成 ExampleJSON 代碼，并將其存儲到類變量?code?中（由于與 C codegen 基本一致，這里跳過了 visitor 函數的實現）。完成計算圖訪問后，在?code?中會生成一個 ExampleJSON 計算圖。

注2：定義內部 API?gen?來獲取子圖，并生成 ExampleJSON 代碼。用戶可以依據個人喜好，為這個 API 命名。

接下來，實現一個自定義 runtime，來利用?ExampleJsonCodeGen?的輸出。

實現自定義 runtime

本節將逐步演示如何自定義 TVM runtime，并將其注冊到 TVM runtime 模塊。自定義 runtime 應位于?src/runtime/contrib/<your-runtime-name>/。本示例中，我們將 runtime 命名為 “example_ext_runtime”。

首先，如下所示定義一個自定義 runtime 類。注意：這個類必須由 TVM?ModuleNode?派生，以保證與其他 TVM runtime 模塊兼容。

#include <dmlc/logging.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/memory.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>

#include <fstream>
#include <cmath>
#include <map>
#include <sstream>
#include <string>
#include <vector>

namespace tvm {
namespace runtime {
class ExampleJsonModule : public ModuleNode {
 public:
  explicit ExampleJsonModule(std::string graph_json);

  PackedFunc GetFunction(const std::string& name,
                         const ObjectPtr<Object>& sptr_to_self) final;

  const char* type_key() const { return "examplejson"; }

  void SaveToBinary(dmlc::Stream* stream) final;

  static Module LoadFromBinary(void* strm);

  static Module Create(const std::string& path);

  std::string GetSource(const std::string& format = "");

  void Run(int id, const std::vector<int>& inputs, int output);

  void ParseJson(const std::string& json);

 private:
  /* \brief 代表計算圖的 json 字符串。 */
  std::string graph_json_;
  /* \brief 正在被處理的子圖。 */
  std::string curr_subgraph_;
  /*! \brief 由子圖 id 到節點條目的簡單圖。 */
  std::map<std::string, std::vector<NodeEntry> > graph_;
  /* \brief 包含圖中每一個節點的張量的簡單池。 */
  std::vector<NDArray> data_entry_;
  /* \brief 從節點 id 到算子名字的映射。 */
  std::vector<std::string> op_id_;
};

以下這些從?ModuleNode?派生的函數，必須在?ExampleJsonModule?中實現：

• 構造函數：這個類的構造函數，應該接收一個表征中的子圖，用戶可以自行決定處理和存儲的格式。保存的子圖可以被以下兩個函數使用。
• GetFunction：這是這個類中最重要的函數。當 TVM runtime 要使用編譯器標簽（compiler tag）執行子圖時，它會從自定義 runtime 模塊中調用此函數。它提供函數名及 runtime 參數，GetFunction?會返回一個打包的函數實現，以供 TVM runtime 執行。
• SaveToBinary?和?LoadFromBinary：SaveToBinary?將 runtime 模塊序列化為二進制格式以供后續部署。用戶使用?export_library?API 時，TVM 會調用這個函數。另一方面，由于用戶這時使用的是自己的計算圖表征，因此必須確保?LoadFromBinary?能夠采用SaveToBinary?生成的序列化二進制文件，來構造相同的 runtime 模塊。
• GetSource（可選）：如果想查看生成的 ExampleJSON 代碼，可以實現這個函數來轉存；否則則可以跳過實現。

實現構造函數

explicit ExampleJsonModule(std::string graph_json) {
  this->graph_json_ = graph_json;
  ParseJson(this->graph_json_);
}

接下來，實現?ParseJson?來解析 ExampleJSON 格式的子圖，并在內存中構造一個計算圖供后續使用。由于本示例不支持帶有分支的子圖，因此只需用一個數組，按順序存儲子圖中的每個節點。

void ParseJson(const std::string& json) {
  std::string line;
  std::string curr_subgraph;
  std::stringstream ss(json);

  while (std::getline(ss, line, '\n')) {
    std::stringstream ss2(line);
    std::string token;
    int id = 0;

    ss2 >> token;
    if (token.find("subgraph_") != std::string::npos) {
      curr_subgraph = token;
      continue;
    }

    ss2 >> id;
    if (op_id_.size() <= static_cast<size_t>(id)) {
      op_id_.resize(id + 1);
      data_entry_.resize(id + 1);
    }

    int64_t total_elements = 1;
    std::vector<int64_t> shape;
    if (token == "input") {
      int64_t size = 0;
      while (ss2 >> size) {
        total_elements *= size;
        shape.push_back(size);
      }
    } else {
      op_id_[id] = token; // 注 1
      bool shape_data = false;
      NodeEntry entry;
      while (ss2 >> token) {
        if (token == "shape:") {
          shape_data = true;
        } else if (shape_data) {
          total_elements *= std::stoll(token);
          shape.push_back(std::stoll(token));
        } else if (token != "inputs:") {
          entry.inputs.push_back(std::stoi(token));
        }
      }
      entry.id = id;
      entry.output = id;
      graph_[curr_subgraph].push_back(entry); // 注 2
    }
    DLDevice dev;
    dev.device_type = static_cast<DLDeviceType>(1);
    dev.device_id = 0;
    data_entry_[id] = NDArray::Empty(shape, DLDataType{kDLFloat, 32, 1}, dev); // 注 3
  }
}

注1：使用類變量?op_id_?將子圖節點 ID 映射到算子名稱（例如?add），以便在 runtime 中調用相應的算子函數。

注2：使用類變量?graph_?從子圖名稱映射到節點數組。GetFunction?將在 runtime 通過子圖 ID 查詢計算圖節點。

注3：使用類變量?data_entry_ ?將子圖節點 ID 映射到張量數據占位符。將輸入和輸出放入 runtime 中對應的數據條目中。

實現 GetFunction

構造函數實現后，以上類變量準備就緒。接下來實現?GetFunction?為 TVM runtime 提供可執行的子圖函數：

PackedFunc GetFunction(const std::string& name,
                       const ObjectPtr<Object>& sptr_to_self) final {
  if (this->graph_.find(name) != this->graph_.end()) {
    this->curr_subgraph_ = name;
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {

      // Copy input tensors to corresponding data entries.
      for (auto i = 0; i < args.size(); ++i) {
        ICHECK(args[i].type_code() == kNDArrayContainer || args[i].type_code() == kArrayHandle)
            << "Expect NDArray or DLTensor as inputs\n";
        if (args[i].type_code() == kArrayHandle) {
          DLTensor* arg = args[i];
          this->data_entry_[i].CopyFrom(arg);
        } else {
          NDArray arg = args[i];
          this->data_entry_[i].CopyFrom(arg);
        }
      }

      // Execute the subgraph.
      for (const auto& it : this->graph_[this->curr_subgraph_]) {
        this->Run(it.id, it.inputs, it.output);
      }
      ICHECK_GT(graph_.count(this->curr_subgraph_), 0U);

      // Copy the output from a data entry back to TVM runtime argument.
      auto out_idx = graph_[this->curr_subgraph_].back().output;
      if (args[args.size() - 1].type_code() == kArrayHandle) {
        DLTensor* arg = args[args.size() - 1];
        this->data_entry_[out_idx].CopyTo(arg);
      } else {
        NDArray arg = args[args.size() - 1];
        this->data_entry_[out_idx].CopyTo(arg);
      }
      *rv = data_entry_.back();
    });
  } else {
    LOG(FATAL) << "Unknown subgraph: " << name << "\n";
    return PackedFunc();
  }
}

可以看出，GetFunction?由三個主要部分組成。第一部分將數據從 TVM runtime 參數，復制到構造函數中指定的對應數據條目。第二部分使用?Run?函數執行子圖（并稍后實現），并將結果保存到另一個數據條目。第三部分將輸出數據條目中的結果，復制回對應的 TVM runtime 參數進行輸出。

實現 Run

Run?函數接收 1）子圖 ID，2）輸入數據條目索引列表和 3）輸出數據條目索引。

void Run(int id, const std::vector<int>& inputs, int output) {
  // Make a list data entry indexs.
  std::vector<int> args(inputs.begin(), inputs.end());
  args.push_back(output);

  // Initialize data holders.
  std::vector<TVMValue> values(args.size());
  std::vector<int> type_codes(args.size());

  // Initialize a TVM arg setter with TVMValue and its type code.
  TVMArgsSetter setter(values.data(), type_codes.data());

  // Set each argument to its corresponding data entry.
  if (op_id_[id] == "add" || op_id_[id] == "sub" || op_id_[id] == "mul") {
    for (size_t i = 0; i < args.size(); i++) {
      setter(i, data_entry_[args[i]]);
    }
  }

  // Invoke the corresponding operator function.
  if (op_id_[id] == "add") {
    Add(values.data(), type_codes.data(), args.size());
  } else if (op_id_[id] == "sub") {
    Sub(values.data(), type_codes.data(), args.size());
  } else if (op_id_[id] == "mul") {
    Mul(values.data(), type_codes.data(), args.size());
  } else {
    LOG(FATAL) << "Unknown op: " << op_id_[id] << "\n";
  }
}

Run?函數主要包括兩部分。第一部分負責分配?TVMValue?列表，并映射相應的數據輸入塊。這也會成為算子函數的參數。第二部分調用算子函數。盡管使用的 C 函數與上一個示例相同，但用戶可以將?Add、Sub?和?Mul?替換為自己的引擎。注意，這里需要確保引擎將結果存儲到最后一個參數，從而使得它們可以傳輸回 TVM runtime。

實現上述功能后，用戶自定義的 codegen 和 runtime 就可以執行子圖了。最后一步是注冊一個 API（examplejson_module_create）來創建這個模塊：

TVM_REGISTER_GLOBAL("module.examplejson_module_create")
.set_body_typed([](std::string code){
    auto n = make_object<ExampleJsonModule>(code);
    return runtime::Module(n);
});

實現 SaveToBinary 和 LoadFromBinary

到目前為止，我們已經實現了與其他 TVM runtime 用法一致的自定義 runtime 的主要功能。但是，當用戶想要將構建的 runtime 保存到磁盤以進行部署時，TVM 不知道如何保存。這就是實現?SaveToBinary?和?LoadFromBinary?的原因，它們會告訴 TVM 這個自定義 runtime 如何持久化和復原。

首先實現?SaveToBinary?函數，允許用戶將此模塊保存在磁盤中。

void SaveToBinary(dmlc::Stream* stream) final {
    stream->Write(this->graph_json_);
}

這個函數非常簡單。在構造函數中，我們采取的唯一參數是一個子圖表征（subgraph representation）。也就是說只需一個子圖表征來構造/恢復這個自定義 runtime 模塊。SaveToBinary?只是將子圖寫到一個輸出的 DMLC 流中，當用戶使用?export_library?API 輸出模塊時，自定義模塊將是一個子圖的 ExampleJSON 流。

LoadFromBinary?讀取子圖流并重新構建自定義 runtime 模塊的流程與此類似：

static Module LoadFromBinary(void* strm) {
  dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
  std::string graph_json;
  stream->Read(&graph_json);
  auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);
  return Module(n);
}

此外，還需要注冊以下函數，啟用相應的 Python API：

TVM_REGISTER_GLOBAL("module.loadbinary_examplejson")
.set_body_typed(ExampleJsonModule::LoadFromBinary);

上述注冊意味著當用戶調用?tvm.runtime.load_module(lib_path)?API，并且導出庫有一個 ExampleJSON 流時，LoadFromBinary?將被調用以創建相同的自定義 runtime 模塊。

另外，如果想支持直接從 ExampleJSON 文件創建模塊，還可以實現一個非常簡單的函數，并注冊一個 Python API，如下所示：

static Module Create(const std::string& path) {
    std::ifstream filep;
    filep.open(path, std::ios::in);
    std::string graph_json;
    std::string line;
    while (std::getline(filep, line)) {
        graph_json += line;
        graph_json += "\n";
    }
    filep.close();
    auto n = tvm::runtime::make_object<ExampleJsonModule>(graph_json);
    return Module(n);
}

TVM_REGISTER_GLOBAL("module.loadfile_examplejson")
.set_body([](TVMArgs args, TVMRetValue* rv) {
    *rv = ExampleJsonModule::Create(args[0]);
});

這意味著用戶可以手動編寫/修改 ExampleJSON 文件，并使用 Python API?tvm.runtime.load_module("mysubgraph.examplejson", "examplejson")?構建自定義模塊。

總結

匯總前文重點：

• 從?ExprVisitor?和?CodegenCBase（僅適用于 C codegen）派生的 codegen 類，具有以下功能：

  • VisitExpr_(const CallNode* call)?收集調用節點信息。
  • 收集子圖信息所需的其他 visitor 函數。
  • JIT?生成子圖代碼。
  • 注冊 codegen。

• 創建?CSourceModule?的函數（用于 C codegen）。

• 從?ModuleNode?派生的 runtime 模塊類，具有以下功能（用于計算圖表征）。

  • 構造函數。
  • GetFunction?生成與 TVM runtime 兼容的?PackedFunc。
  • Run?執行子圖。
  • 注冊 runtime creation API。
  • SaveToBinary?和?LoadFromBinary?序列化/反序列化自定義 runtime 模塊。
  • 注冊?LoadFromBinary?API 為tvm.runtime.load_module(your_module_lib_path)提供支持。
  • （可選）Create?支持從表征的子圖文件，構建自定義 runtime 模塊。

• 一個注釋器，用于注釋用戶 Relay 程序，利用編譯器和 runtime（待定）。