如何在TVM上集成Codegen（下）

如何在TVM上集成Codegen（下）

Bring DNNL to TVM: JSON Codegen/Runtime

现在实现将中继图序列化为JSON表示的DNNL codegen，然后实现DNNL JSON runtime来反序列化和执行该图。请注意，如果尝试实现codegen来生成C兼容的程序，可能需要直接进入下一节。             

要使TVM中的DNNL JSON codegen/runtime在本例中工作，请确保DNNL在计算机上可用，并在中使用set(USE_DNNL_CODEGEN ON)构建TVM配置文件制作。           

DNNL codegen在src/relay/backend/contrib/dnnl/codegen.cc。 因为在这个文件中的两个表单中都实现了DNNL codegen，所以在跟踪代码时，可以将注意力集中在USE_JSON_RUNTIME宏所涵盖的部分。             

首先用TVM注册API（L510）注册codegen。此注册使TVM编译引擎将Compiler=<your codegen>的中继函数分派到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);

每个 runtime模块只负责一个中继函数，这意味着您可能在一个single .so文件中有多个DNNL runtime模块。             

DNNL JSON序列化             

接下来，实现dnnljson序列化器（L429）。

我们从BYOC JSON codegen (src/relay/backend/contrib/codegen_json/codegen_json.h)派生而来。DNNL JSON serializer中的特殊进程尝试序列化对可由DNNL JSON runtime解释的JSON节点的复合函数调用。假设我们有一个与模式匹配的复合函数dnnl.conv2d_relu公司，则BYOC JSON codegen将生成以下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]

  }

}

问题是在runtime仍然需要Conv2D属性，比如padding和stripes，但是BYOC JSON序列化程序只附加复合函数的属性，而不附加body算子。另一方面，定制的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 runtime能够解释，就可以定制序列化程序以生成JSON格式的任何表单。

DNNL JSON Runtime

实现一个DNNL JSON runtime来解释和执行序列化的JSON图。把它放在src/runtime/contrib/dnnl/dnnl_json_runtime.cc。             

同样，首先注册两个api来创建 runtime，这样就可以在任何地方使用。这个runtime.DNNLJSONRuntimeCreate序列化后在上一部分中使用，并且runtime.module.loadbinary_dnnl_json可以在加载.so back时使用。

// 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>);

Now we explain DNNL JSON runtime implementation. The basic class structure is:

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引擎（BuildEngine请参阅L93），并将常量权重填充到相应的数据输入缓冲区（SetupConstant在JSON runtime基类中实现，只需在Init中调用它）。

即使我们运行多次推断，这个函数也只会被调用一次。             

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

由于DNNL JSON runtime中的rest实现太过DNNL特定，因此在本文中我们将停止讨论。要强调的是，虽然DNNL JSON runtime是一个很好的参考，但是JSON runtime可以完全定制以满足需求。

Bring DNNL to TVM: C Source Codegen

现在让我们实现DNNL codegen，它生成C源代码，它调用dnnlapi来执行中继图表。注释如果试图实现一个codegen来生成JSON格式的其他图形表示，那么可能需要阅读DNNL to TVM: JSON Codegen/Runtime并跳过这一节。             

要使TVM中的DNNL C源代码生成在本例中工作，确保DNNL在计算机上可用，并在中使用set(USE_DNNL_CODEGEN C_SRC)构建TVM配置文件制作.             

DNNL codegen在src/relay/backend/contrib/dnnl/codegen.cc。由于在这个文件中的两个表单中都实现了DNNL codegen，所以在跟踪代码时，可以将注意力集中在USE_JSON_RUNTIME runtime宏未涵盖的部分。             

首先用TVM注册API（L510）注册codegen。此注册使TVM编译引擎将Compiler=<your codegen>的中继函数分派到relay.ext.<your codegen>。然后实现DNNL编译器（L490）的entry函数：

runtime::Module DNNLCompiler(const ObjectRef& ref) {

  DNNLModuleCodegen dnnl;

  return dnnl.CreateCSourceModule(ref);

}

TVM_REGISTER_GLOBAL("relay.ext.dnnl").set_body_typed(DNNLCompiler);

每个 runtime模块只负责一个中继函数，这意味着您可能在single .so文件中有多个DNNL runtime模块。             

然后，推导了CSourceModuleCodegenBase，在L362中实现了DNNLModuleCodegen。虽然CSourceModuleCodegenBase负责序列化等其他模块级流程，只需要在CreateCSourceModule函数（L389）中实现DNNL代码生成：

runtime::Module CreateCSourceModule(const ObjectRef& ref) override {

    // Include headers

    // ...skip...

    code_stream_ << "#include <dnnl/dnnl_kernel.h> ";

    // ...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），用DNN API生成可编译的C代码，如下所示。有关TVM C source runtime模块兼容函数接口的说明，请参阅嵌入的注释。

// 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。

因为rest实现在src/relay/backend/contrib/dnnl/codegen.cc太DNNL的具体细节，本文就到此为止。其主要思想是实现一个中继图访问者（L138）来访问给定的中继函数并生成上面的C代码。只要codegen能够生成与TVM运行时兼容的C代码，就可以完全定制codegen以满足您的需求。

C Source Compilation

DNNLCompiler的输出是一个模块，其中生成的C代码是文本格式的，GCC尚未将其编译为可执行的二进制文件。实际上，当用户调用export_libray(mod)时，会编译生成的C代码，如下面的代码片段：

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)

Bring DNNL to TVM: Build TVM with DNNL Codegen/Runtime

最后，创建cmake/modules/contrib/DNNL.cmake在构建TVM时包括DNNL codegen。为了演示，DNNL codegen在同一个cmake文件中有两个实现。只能根据需要专注于其中的一个。             

cmake文件就绪后，现在用户可以在其构建中指定set(USE_DNNL_CODEGEN ON)的build/config.cmake配置文件制作启用DNNL codegen。