如何将您自己的代码生成引入 TVM

2020 年 7 月 15 日 • Zhi Chen 和 Cody Yu，亚马逊网络服务公司

为了将数据科学家从开发新模型时的性能担忧中解放出来，硬件后端提供商（例如，Intel、NVIDIA、ARM 等）要么提供内核库（如 cuBLAS 或 cuDNN），其中包含许多常用的深度学习内核，要么提供框架（如 DNNL 或 TensorRT），其中包含图引擎，让用户以某种方式描述他们的模型以实现高性能。此外，新兴的深度学习加速器也有自己的编译器、内核库或运行时框架。

然而，当用户尝试使用新的内核库或设备时，他们必须学习新的编程接口。因此，对统一编程接口的需求变得越来越重要，以便所有用户和硬件后端提供商能够站在同一起跑线上。

为了与广泛使用的深度学习框架共享编程接口，许多硬件设备提供商已尝试将其设备后端集成到 TensorFlow 中。然而，由于 TensorFlow 没有为新后端提供官方后端接口，您必须破解 TensorFlow 进行注册，这涉及到许多源文件更改，并使未来的维护变得困难。

在这篇文章中，我们将演示作为硬件后端提供商的您，如何轻松利用自带代码生成 (BYOC) 框架，将您的硬件设备的内核库/编译器/框架集成到 TVM 中。利用 BYOC 框架最重要的优势是，您设备的所有相关源文件都是自包含的，因此您设备的代码生成/运行时可以插拔到 TVM 代码库中。 这意味着：1) 带有您的代码生成的 TVM 代码库将是向上游兼容的，并且 2) TVM 用户可以根据他们的需要选择启用代码生成/运行时。

在本文的其余部分，我们首先说明您可能需要带有 BYOC 的 TVM 的场景，然后概述 BYOC 编译和运行时流程。然后，我们通过使用 Intel DNNL（也称为 MKL-DNN、OneDNN）作为运行示例，逐步说明如何使用 BYOC 将供应商库或执行引擎集成到 TVM 中。

将 ASIC 加速器引入 TVM

让我们首先构建一个场景来说明您为什么要将加速器引入 TVM，以及您可以从 BYOC 框架中期望获得哪些功能。如果您不确定您的案例是否适合 BYOC，欢迎在 discuss.tvm.ai 上发起讨论。

假设您刚刚制造了一个边缘设备平台，该平台配备了 ARM CPU 和一个出色的加速器，该加速器在常见的图像分类模型上取得了惊人的性能。换句话说，您的加速器在 Conv2D、ReLU、GEMM 和其他广泛使用的 CNN 算子上表现良好。

不幸的是，目标检测模型也越来越受欢迎，您的客户需要在您的平台上同时运行图像分类和目标检测模型。尽管您的加速器能够执行目标检测模型中的几乎所有算子，但缺少一个算子（例如，非极大值抑制，NMS）。

让 TVM 执行不支持的算子

由于 TVM 具有针对不同后端的多种代码生成器，开源社区可以很容易地在短时间内在 CPU 或 GPU 上实现新的算子。理想情况下，如果您使用 BYOC 将加速器的编译流程集成到 TVM 中，TVM 将执行 Relay 图分区，将图的一部分卸载到您的加速器，同时将其他部分保留在 TVM 上。因此，您可以声称您的平台能够运行所有模型，而无需担心新的算子。

自定义图级优化

您的 ASIC 加速器必须有自己的编译流程。通常，它可能是以下情况之一

生成图表示并将其馈送到图引擎：您可能拥有自己的图引擎，该引擎能够在您的加速器上执行图（或神经网络模型）。例如，Intel DNNL 和 NVIDIA TensorRT 都使用引擎来运行整个图或模型，以便它们能够 1) 减少算子之间的内存事务，以及 2) 通过算子融合优化图执行。

为了实现上述两个优化，您可能需要在编译时处理图。例如，Conv2D 和偏置加法在 TVM 中是两个独立的算子，但在您的加速器上它们可能是一个算子（具有偏置加法功能的 Conv2D）。在这种情况下，您可能希望通过将 conv2d - add 图模式替换为 your_conv2d_with_bias 节点来优化图。

如果您的编译流程属于这种情况，那么我们建议阅读本文的其余部分，但跳过 将 DNNL 引入 TVM：C 源代码生成。

生成汇编代码并将其编译为可执行二进制文件：如果您的平台没有像前一种情况那样的端到端执行框架，您可能有一个编译器来编译 ISA 汇编代码中的程序。为了将汇编代码馈送到您的编译器，您将需要一个代码生成器来从 Relay 图生成和优化汇编代码。

如果您的编译流程属于这种情况，那么我们建议阅读本文的其余部分，但跳过 将 DNNL 引入 TVM：JSON 代码生成/运行时。

BYOC 的工作原理

然后，我们简要解释 BYOC 框架的工作原理。有关底层框架组件及其实现的更详细说明，请参阅 开发者文档。简而言之，给定图 1 中的 Relay 图，BYOC 框架执行以下步骤

1. 图注释

采用用户提供的 Relay 图，我们的第一步是注释图中可能卸载到您的加速器的节点。您需要按照 将 DNNL 引入 TVM：注释规则 来实现受支持算子的白名单，或自定义复合算子的图模式列表。图 2 显示了一个注释结果示例。

2. 图变换

第二步是根据注释变换和优化图。具体来说，BYOC 执行以下变换。

2.1：合并编译器区域：如图 2 所示，我们现在在图中有很多可以卸载到您的加速器的“区域”，但其中一些区域实际上可以合并以减少数据传输和内核启动开销。因此，步骤 2.1 使用贪婪算法来尽可能多地合并这些区域，同时保证功能正确性。结果如图 3 所示。

2.2：分区图：对于上一步中的每个区域，我们创建一个带有属性 Compiler 的 Relay 函数，以指示此 Relay 函数应完全卸载到您的加速器，如图 4 所示。

3. 代码生成

现在我们知道 Relay 图的哪一部分应该卸载。在此步骤中，我们依次将每个带有 Compiler=your_accelerator 的 Relay 函数发送到您的代码生成器。您的代码生成器应将 Relay 函数编译为与您自己的编译流程相匹配的形式。它可以是 C 源代码或任何文本格式。

最后，所有编译后的函数将与其他未卸载的 Relay 函数一起序列化到单个 .so 文件中，通过 TVM export_library Python API。换句话说，用户在运行此流程后只会得到一个 .so 文件。

4. 运行时

您可能还需要实现一个运行时来初始化您的图引擎（如果适用）并执行编译后的函数。在推理期间，TVM 运行时（即，图运行时或 VM）将利用您的运行时来调用卸载的函数，当 TVM 运行时在图 4 中遇到相应的函数调用时。您的运行时负责启动编译后的函数，并使用给定的输入张量数组，并将结果填充到输出张量数组中。

在本文的其余部分，我们使用 DNNL 作为示例来演示如何使用 BYOC 框架实现上述工作流程。请注意，本文中引用的所有代码和行号均基于 TVM 仓库主分支提交 8a0249c。

将 DNNL 引入 TVM：注释规则

BYOC 框架提供了两种方法来描述受支持的算子和模式。您可以同时使用这两种方法。在本节中，我们以 DNNL 为例展示如何使用它们。完整的实现可在 此处 获得。请注意，我们将您的代码生成的注释规则放在 python/tvm/relay/op/contrib/your_codegen_name.py 下。

单算子的规则

您可以直观地使用 BYOC API 指定您的加速器支持哪些 Relay 算子。例如，我们使用以下代码片段来构建一个规则，说明我们的 DNNL 代码生成支持 Conv2D

@tvm.ir.register_op_attr("nn.conv2d", "target.dnnl")
def _dnnl_conv2d_wrapper(attrs, args):
  return True

这会将新的属性 target.dnnl 注册到 Relay nn.conv2d 算子。通过这种方式，BYOC 注释可以为图中的每个算子调用 target.dnnl()，以检查它是否在 DNNL 代码生成中受支持。

另一方面，为每个算子编写上述代码片段可能很乏味。对于 DNNL 实现，我们实现了一个辅助函数 _register_external_op_helper，以使我们的生活更轻松

def _register_external_op_helper(op_name, supported=True):
    @tvm.ir.register_op_attr(op_name, "target.dnnl")
    def _func_wrapper(attrs, args):
        return supported
    return _func_wrapper

_register_external_op_helper("nn.batch_norm")
_register_external_op_helper("nn.conv2d")
_register_external_op_helper("nn.dense")
_register_external_op_helper("nn.relu")
_register_external_op_helper("add")
_register_external_op_helper("subtract")
_register_external_op_helper("multiply")

在上面的示例中，我们指定了一个 DNNL 代码生成可以支持的算子列表。

图模式的规则

您的加速器或编译器可能已经优化了一些模式（例如，Conv2D + add + ReLU）以成为单个指令或 API。在这种情况下，您可以指定从图模式到您的指令/API 的映射。对于 DNNL 的情况，其 Conv2D API 已经包含偏置加法，并且允许附加下一个 ReLU，因此我们可以将 DNNL 称为以下代码片段（完整的实现可以在 此处 找到）

DNNLConv2d(const bool has_bias = false, const bool has_relu = false) {
  // ... skip ...
  auto conv_desc = dnnl::convolution_forward::desc(
    dnnl::prop_kind::forward_inference,
    dnnl::algorithm::convolution_direct,
    conv_src_md, conv_weights_md, conv_bias_md, conv_dst_md,
    strides_dims, padding_dims_l, padding_dims_r);

  // Attach ReLU
  dnnl::primitive_attr attr;
  if (has_relu) {
    dnnl::post_ops ops;
    ops.append_eltwise(1.f, dnnl::algorithm::eltwise_relu, 0.f, 0.f);
    attr.set_post_ops(ops);
  }

  auto conv2d_prim_desc = dnnl::convolution_forward::primitive_desc(
    conv_desc, attr, engine_);
  // ... skip ...

在这种情况下，除了单个 conv2d 之外，我们希望将图模式 conv2d+relu 映射到 DNNLConv2d(false, true)，并将 conv2d+add+relu 映射到 DNNLConv2d(true, true)。我们可以使用以下代码片段来实现它

def make_pattern(with_bias=True):
  data = wildcard()
  weight = wildcard()
  bias = wildcard()
  conv = is_op('nn.conv2d')(data, weight)
  if with_bias:
    conv_out = is_op('add')(conv, bias)
  else:
    conv_out = conv
  return is_op('nn.relu')(conv_out)

@register_pattern_table("dnnl")
def pattern_table():
  conv2d_bias_relu_pat = ("dnnl.conv2d_bias_relu", make_pattern(with_bias=True))
  conv2d_relu_pat = ("dnnl.conv2d_relu", make_pattern(with_bias=False))
  dnnl_patterns = [conv2d_bias_relu_pat, conv2d_relu_pat]
  return dnnl_patterns

在 DNNL 示例中，我们实现了两个具有不同名称的模式，以便我们可以在代码生成中轻松识别它们。请注意，这些模式是在 Relay 模式语言中实现的。您可以按照 本教程 学习如何编写自己的模式。

使用模式表，我们可以使用 Relay pass 来执行从

%1 = nn.conv2d(%data, %weight, ...)
%2 = add(%1, %bias)
%3 = nn.relu(%2)

到

%1 = fn(%input1, %input2, %input3,
        Composite="dnnl.conv2d_bias_relu",
        PartitionedFromPattern="nn.conv2d_add_nn.relu_") {
  %1 = nn.conv2d(%input1, %input2, ...)
  %2 = add(%1, %input3)
  nn.relu(%2)
}
%2 = %1(%data, %weight, %bias)

的转换。因此，DNNL 代码生成可以获取模式名称 conv2d_bias_relu 并将 %1 映射到 DNNLConv2d(true, true)。

您可能已经注意到，我们在复合函数中还有一个名为 “PartitionedFromPattern” 的属性。如果您的模式包含 wildcard 算子，这可能会有所帮助。例如，我们可能有一个模式表 ("conv2d_with_something", conv2d -> *)

def make_pattern(with_bias=True):
  data = wildcard()
  weight = wildcard()
  conv = is_op('nn.conv2d')(data, weight)
  return wildcard()(conv)

在这种情况下，您将获得一个带有 Composite=conv2d_with_something 的复合函数，但您不知道它实际匹配的图是什么。这就是 PartitionedFromPattern 发挥作用的地方。您可以通过查看 PartitionedFromPattern 来了解匹配的图是 conv2d -> add 还是 conv2d -> relu，以查看它是 nn.conv2d_add_ 还是 nn.conv2d_nn.relu_。

将 DNNL 引入 TVM：Relay 图变换

通过上一步的注释规则，我们现在可以应用一系列 BYOC Relay pass 来将 Relay 图从图 1 转换为图 4

mod = create_relay_module_from_model() # Output: Figure 1
mod = transform.MergeComposite(pattern_table)(mod)
mod = transform.AnnotateTarget(["dnnl"])(mod) # Output: Figure 2
mod = transform.MergeCompilerRegions()(mod) # Output: Figure 3
mod = transform.PartitionGraph()(mod) # Output: Figure 4

可以看出，每个 Relay pass 都可以映射到我们在 BYOC 的工作原理 中介绍的步骤。

将 DNNL 引入 TVM：JSON 代码生成/运行时

现在让我们实现 DNNL 代码生成，它将 Relay 图序列化为 JSON 表示，然后实现 DNNL JSON 运行时来反序列化和执行图。请注意，如果您尝试实现代码生成以生成 C 兼容的程序，您可能希望直接进入下一节。

为了使 TVM 中的 DNNL JSON 代码生成/运行时能够在此示例中工作，请确保您的机器上安装了 DNNL，并在 config.cmake 中使用 set(USE_DNNL_CODEGEN ON) 构建 TVM。

DNNL 代码生成在 src/relay/backend/contrib/dnnl/codegen.cc 中实现。由于我们出于说明目的在此文件中以两种形式实现了 DNNL 代码生成，因此在跟踪代码时，您可以专注于 USE_JSON_RUNTIME 宏涵盖的部分。

我们首先使用 TVM 注册 API 注册代码生成器 (L510)。此注册使 TVM 编译引擎将带有 Compiler=<your codegen> 的 Relay 函数分派到 relay.ext.<your codegen>。然后我们实现 DNNL 编译器的入口函数 (L490)。请阅读代码片段中嵌入的注释以了解详细信息

runtime::Module DNNLCompiler(const ObjectRef& ref) {
  // "ref" should be the paritioned Relay function with kCompiler=dnnl.
  CHECK(ref->IsInstance<FunctionNode>());
  auto func = Downcast<Function>(ref);

  // Get the function name as the symbol to match in runtime.
  auto func_name = GetExtSymbol(func);

  // Serialize the function to a JSON string (introduce later).
  DNNLJSONSerializer serializer(func_name, func);
  serializer.serialize();
  std::string graph_json = serializer.GetJSON();

  // The constant tensor names that have been bound to the module.
  // All constant tensors will be serialzied along with the JSON graph
  // when export_library is invoked.
  auto params = serializer.GetParams();

  // The function to create DNNL JSON runtime (introduce later).
  const auto* pf = runtime::Registry::Get("runtime.DNNLJSONRuntimeCreate");
  CHECK(pf != nullptr) << "Cannot find JSON runtime module to create";

  // Create a DNNL runtime module that can run the serialized function.
  auto mod = (*pf)(func_name, graph_json, params);
  return mod;
}
TVM_REGISTER_GLOBAL("relay.ext.dnnl").set_body_typed(DNNLCompiler);

请注意，每个运行时模块仅负责一个 Relay 函数，这意味着在一个 .so 文件中您可能有多个 DNNL 运行时模块。

DNNL JSON 序列化

接下来，我们实现 DNNL JSON 序列化器 (L429)。我们从 BYOC JSON 代码生成器派生它 (src/relay/backend/contrib/codegen_json/codegen_json.h)。DNNL JSON 序列化器中的特殊过程尝试将复合函数调用序列化为 JSON 节点，该节点可以由 DNNL JSON 运行时解释。假设我们有一个复合函数，它匹配模式 dnnl.conv2d_relu，那么 BYOC JSON 代码生成器将生成以下 JSON 节点

{
  op: "kernel",
  name: "dnnl.conv2d_relu",
  inputs: [[0, 0, 0], [1, 0, 0]],
  attrs: {
    PartitionedFromPattern: ["nn.conv2d_nn.relu_"],
    shape: [1, 32, 14, 14]
  }
}

问题是我们在运行时仍然需要 Conv2D 属性，例如 padding 和 strides，但 BYOC JSON 序列化器仅附加复合函数的属性，而不是主体算子的属性。另一方面，自定义的 DNNL JSON 序列化器附加复合函数中第一个也是唯一的 Conv2D 的属性，以生成以下 JSON 节点

{
  op: "kernel",
  name: "dnnl.conv2d_relu",
  inputs: [[0, 0, 0], [1, 0, 0]],
  attrs: {
    shape: [1, 32, 14, 14],
    data_layout: ["NCHW"],
    kernel_layout: ["OIHW"],
    strides: [1, 1],
    padding: [1, 1, 1, 1]
  }
}

从 DNNL JSON 序列化器可以看出，您可以自定义序列化器以生成您喜欢的任何 JSON 形式，只要您的 JSON 运行时可以解释它们。

DNNL JSON 运行时

然后，我们实现一个 DNNL JSON 运行时来解释和执行序列化的 JSON 图。我们将其放在 src/runtime/contrib/dnnl/dnnl_json_runtime.cc 下。

同样，我们首先注册两个 API 来创建运行时，以便我们可以在任何地方使用它们。runtime.DNNLJSONRuntimeCreate 在上一部分序列化后使用，而 runtime.module.loadbinary_dnnl_json 可以在加载 .so 文件时使用。

// Create a DNNL JSON runtime to interpret and execute the given JSON graph.
runtime::Module DNNLJSONRuntimeCreate(String symbol_name, String graph_json,
                                      const Array<String>& const_names) {
  auto n = make_object<DNNLJSONRuntime>(symbol_name, graph_json, const_names);
  return runtime::Module(n);
}
TVM_REGISTER_GLOBAL("runtime.DNNLJSONRuntimeCreate")
    .set_body_typed(DNNLJSONRuntimeCreate);

TVM_REGISTER_GLOBAL("runtime.module.loadbinary_dnnl_json")
    .set_body_typed(JSONRuntimeBase::LoadFromBinary<DNNLJSONRuntime>);

现在我们解释 DNNL JSON 运行时的实现。基本类结构是

class DNNLJSONRuntime : public JSONRuntimeBase {
  const  char* type_key() const { return  "dnnl_json"; } 
  void Init(const Array<NDArray>& consts) override {
    // Initialize the DNNL graph engine.
    BuildEngine();
    
    // Setup constants entries for weights.
    CHECK_EQ(consts.size(), const_idx_.size())
      << "The number of input constants must match the number of required.";
    SetupConstants(consts);
  }

  void Run() override {
   // 1. Fill in the input buffers.
   // 2. Invoke the engine through intepreting the stream.
   // 3. Read and fill output buffers.
  }
}

Init 函数负责通过解释 JSON 图字符串来构建 DNNL 引擎（参见 L93 中的 BuildEngine），并将常量权重填充到相应的数据条目缓冲区（SetupConstant 在 JSON 运行时基类中实现，因此您只需要在 Init 中调用它）。请注意，即使我们多次运行推理，此函数也只会被调用一次。

接下来，Run 函数 (L64) 首先将输入张量（可能来自用户输入或常量权重）写入我们在构建 DNNL 引擎时初始化的相应 DNNL 内存缓冲区。然后启动 DNNL 引擎以执行 JSON 图。最后，它将 DNNL 输出内存缓冲区写回相应的输出张量。

由于 DNNL JSON 运行时中的其余实现过于 DNNL 特定，无法在本文中深入探讨细节，我们将在此处停止。我们想强调的是，虽然 DNNL JSON 运行时是一个很好的入门参考，但您的 JSON 运行时可以完全自定义以满足您的要求。

将 DNNL 引入 TVM：C 源代码生成

现在让我们实现 DNNL 代码生成，它生成调用 DNNL API 以执行 Relay 图的 C 源代码。请注意，如果您尝试实现代码生成以生成其他图表示形式（如 JSON 格式），您可能希望阅读 将 DNNL 引入 TVM：JSON 代码生成/运行时 并跳过本节。

为了使 TVM 中的 DNNL C 源代码生成能够在此示例中工作，请确保您的机器上安装了 DNNL，并在 config.cmake 中使用 set(USE_DNNL_CODEGEN C_SRC) 构建 TVM。

DNNL 代码生成在 src/relay/backend/contrib/dnnl/codegen.cc 中实现。由于我们出于说明目的在此文件中以两种形式实现了 DNNL 代码生成，因此在跟踪代码时，您可以专注于 未 被 USE_JSON_RUNTIME 宏涵盖的部分。

我们首先使用 TVM 注册 API 注册代码生成器 (L510)。此注册使 TVM 编译引擎将带有 Compiler=<your codegen> 的 Relay 函数分派到 relay.ext.<your codegen>。然后我们实现 DNNL 编译器的入口函数 (L490)

runtime::Module DNNLCompiler(const ObjectRef& ref) {
  DNNLModuleCodegen dnnl;
  return dnnl.CreateCSourceModule(ref);
}
TVM_REGISTER_GLOBAL("relay.ext.dnnl").set_body_typed(DNNLCompiler);

请注意，每个运行时模块仅负责一个 Relay 函数，这意味着在一个 .so 文件中您可能有多个 DNNL 运行时模块。

然后，我们派生 CSourceModuleCodegenBase 以在 L362 中实现 DNNLModuleCodegen。虽然 CSourceModuleCodegenBase 负责其他模块级流程（如序列化），但我们只需要在 CreateCSourceModule 函数中实现 DNNL 代码生成 (L389)

runtime::Module CreateCSourceModule(const ObjectRef& ref) override {
    // Include headers
    // ...skip...
    code_stream_ << "#include <dnnl/dnnl_kernel.h>\n";
    // ...skip...

    // "ref" should be the paritioned Relay function with kCompiler=dnnl.
    CHECK(ref->IsInstance<FunctionNode>());
    auto res = GenDNNLFunc(Downcast<Function>(ref));

    // "code" is the generated C code with DNNL APIs.
    std::string code = code_stream_.str();

    // "res" is a tuple of constant weights (symbols, values).
    // All constant tensors will be serialzied along with the generated C code
    // when export_library is invoked.
    String sym = std::get<0>(res);
    Array<String> variables = std::get<1>(res);

    // Create a CSource module with all above artifacts.
    const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate");
    CHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
    return (*pf)(code, "c", sym, variables);
  }

接下来，我们实现 GenDNNLFunc (L365) 以生成使用 DNNL API 的可编译 C 代码，如下所示。请参阅嵌入的注释以了解 TVM C 源代码运行时模块兼容函数接口的说明。

// The example Relay graph: conv2d -> add -> relu.
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/container.h>
#include <tvm/runtime/packed_func.h>
#include <dlpack/dlpack.h>
#include <dnnl/dnnl_kernel.h>
using namespace tvm::runtime;
using namespace tvm::runtime::contrib;

// Execute the conv2d->add->relu graph with DNNL.
extern "C" void dnnl_0_(float* dnnl_0_i0, float* dnnl_0_i1,
                        float* dnnl_0_i2, float* out0) {
  // Allocate intermediate buffers.
  float* buf_0 = (float*)std::malloc(4 * 4608);
  float* buf_1 = (float*)std::malloc(4 * 4608);
  float* buf_2 = (float*)std::malloc(4 * 4608);

  // Pre-implemented op-based DNNL functions.
  dnnl_conv2d(dnnl_0_i0, dnnl_0_i1, buf_0, 1, 32, 14, 14, 32, 1, 0, 0, 3, 3, 1, 1);
  dnnl_add(buf_0, dnnl_0_i2, buf_1, 1, 32, 12, 12);
  dnnl_relu(buf_1, buf_2, 1, 32, 12, 12);

  // Copy the final output to the corresponding buffer.
  std::memcpy(out0, buf_2, 4 * 4608);
  std::free(buf_0);
  std::free(buf_1);
  std::free(buf_2);
}

// The wrapper function with all arguments in DLTensor type.
extern "C" int dnnl_0_wrapper_(DLTensor* arg0,
        DLTensor* arg1,
        DLTensor* arg2,
        DLTensor* out0) {

  // Cast all DLTensor to primitive type buffers and invoke the above
  // execution function.
  dnnl_0_(static_cast<float*>(arg0->data),
  static_cast<float*>(arg1->data),
  static_cast<float*>(arg2->data),
  static_cast<float*>(out0->data));
  return 0;
}

// The TVM macro to generate TVM runtime compatible function "dnnl_0"
// from our generated "dnnl_0_wrapper_".
TVM_DLL_EXPORT_TYPED_FUNC(dnnl_0, dnnl_0_wrapper_);

请注意，预先实现的基于算子的 DNNL 函数位于 src/runtime/contrib/dnnl/dnnl.cc 中。

由于 src/relay/backend/contrib/dnnl/codegen.cc 中的其余实现过于 DNNL 特定，无法在本文中深入探讨细节，我们将在此处停止。主要思想是实现一个 Relay 图访问器 (L138) 以访问给定的 Relay 函数并生成上述 C 代码。只要您的代码生成器能够生成 TVM 运行时兼容的 C 代码，您就可以完全自定义代码生成器以满足您的要求。

C 源代码编译

您可能已经注意到，DNNLCompiler 的输出是一个模块，其中包含文本格式的生成的 C 代码，该代码尚未由 gcc 编译为可执行二进制文件。实际上，生成的 C 代码将在用户调用 export_libray(mod) 时编译，如下面的代码片段

def update_lib(lib):
    # Include the path of src/runtime/contrib/dnnl/dnnl.cc
    test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
    source_dir = os.path.join(test_dir, "..", "..", "..")
    contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")

    # Setup the gcc flag to compile DNNL code.
    kwargs = {}
    kwargs["options"] = ["-O2", "-std=c++14", "-I" + contrib_path]
    tmp_path = util.tempdir()
    lib_name = 'lib.so'
    lib_path = tmp_path.relpath(lib_name)

    # The generated C code with DNNL APIs is compiled to a binary lib.so.
    lib.export_library(lib_path, fcompile=False, **kwargs)

    # Load the lib.so back to a runtime module.
    lib = runtime.load_module(lib_path)
    return lib

with tvm.transform.PassContext(opt_level=3):
    json, lib, param = relay.build(mod, target=target, params=params)
lib = update_lib(lib)
rt_mod = tvm.contrib.graph_runtime.create(json, lib, ctx)

将 DNNL 引入 TVM：使用 DNNL 代码生成/运行时构建 TVM

最后，我们创建 cmake/modules/contrib/DNNL.cmake 以在构建 TVM 时包含 DNNL 代码生成。出于演示目的，我们的 DNNL 代码生成在同一个 cmake 文件中有两个实现。您可以根据需要只关注其中一个。

准备好 cmake 文件后，用户现在可以在他们的 build/config.cmake 中指定 set(USE_DNNL_CODEGEN ON) 以启用 DNNL 代码生成。

• Zhi Chen 是 TVM PMC 成员以及 SageMaker Neo、Amazon AI、AWS 的高级工程师。

• Cody Yu 是 TVM 审阅者以及 Amazon AI、AWS 的应用科学家。

致谢

我们要感谢我们的同事 Animesh Jain 在框架设计中进行的宝贵讨论；OctoML 的 Tianqi Chen 和 Jared Roesch 进行了系统设计讨论和原型设计；TVM 社区的 Masahiro Masuda 帮助代码审查并改进了 DNNL 集成。我们还要感谢 ARM, U.K. 的 Ramana Radhakrishnan、Matthew Barrett、Manupa Karunaratne 和 Luke Hutton，他们贡献了几个有用的想法、相关的 Relay pass 以及 Arm Compute Library (ACL) 与 BYOC 的集成。