在TVM中添加新设备Codegen

First Post:

Blog Link:
[site]

前言

最近想试着往TVM(v0.7)中加入针对某一新的底层设备的Codegen，本以为是个比较容易的事，但源代码看着看着发现TVM中Codegen和其他部分的耦合比预想中的更紧，故写下这篇短文记录一下添加新的Codegen需要涉及的部分。

Codegen所针对的IR

可以针对两种IR进行Codegen：

  • Relay IR（之后简写为Relay）
  • Tensor Level IR（之后简写为TIR）

后者TIR是Relay出现之前TVM使用的IR，相对更加底层（在PL特性的丰富性上）。placeholder等Tensor Expression（之后简写为TE）在C++端都是被转化为TIR（如TE的OperationNode继承自TIR的FunctionBaseNode）。

在Relay出现之前，各种后端的Codegen都是针对TIR进行的。Relay出现之后，可能出于复用性等方面的考虑，Relay的通常的Codegen都是先转化为TIR再进行的，如tvm.relay.build_module.build的调用路径为：

  1. tvm.relay.build_module.build　in python/tvm/relay/build_module.py
  2. tvm.relay.build_module.BuildModule.build in …
  3. tvm::relay::backend::RelayBuildModule::Build in src/relay/backend/build_module.cc
  4. tvm::relay::backend::RelayBuildModule::BuildRelay in …
  5. tvm::build in src/driver/driver_api.cc
  6. tvm::build（前一个build的重载）　in …
  7. tvm::codegen::Build in src/target/codegen.cc，在这里会用tvm::runtime::Registry::Get根据传入的target参数动态派发给具体函数，如target为"cuda"时会派发给注册为codegen.build_cuda的tvm::codegen::BuildCUDA(in src/target/opt/build_cuda_on.cc)。

本文主要介绍针对TIR生成代码所涉及的部分。关于针对Relay做Codegen可以参考官方文档（虽然有点简略）。

接下来将以一个非常简单（代码量不足70行）的后端模块”hello”的构造来展示针对TIR的代码生成。该模块的使用如下：

1
2
3
4
5
6
7
import tvm

A = tvm.placeholder((10, 10))
B = tvm.compute((10, 10), lambda i, j: A[i, j])
s = tvm.create_schedule(B.op)
f = tvm.build(s, [A, B], "hello")
print(f.get_source())

输出：

1
2
3
4
5
6
7
8
9
10
11
12
13
// hello tvm backend!
void default_function( void* args,  void* arg_type_ids, int32_t num_args,  void* out_ret_value,  void* out_ret_tcode) {
  // ......
  float* placeholder = (float*)(((DLTensor*)arg0)[0].data);
  // ......
  float* compute = (float*)(((DLTensor*)arg1)[0].data);
  // ......
  for (int32_t i = 0; i < 10; ++i) {
    for (int32_t j = 0; j < 10; ++j) {
      compute[((i * 10) + j)] = placeholder[((i * 10) + j)];
    }
  }
}

新建模块文件

新增的模块的python代码一般写在python/tvm/contrib/，C++代码一般写在src/runtime/contrib/。我们这里在src/runtime/contrib/下添加文件夹hello/，hello/内添加hello_module.cc。

因为作为示例的hello模块非常简单，因此只用单文件就够了。一般来讲要视情况组织为多个编译单元，具体可参见src/runtime/contrib/中其他模块的组织方式。

hello模块直接利用tvm.build等前端接口进行构建，如果还需要处理一些前端逻辑或是提供其他的前端接口，则可在python/tvm/contrib/中添加对应的文件。

实现后端模块

tvm.build和tvm.build_module.build返回的都是tvm::runtime::Module，而tvm::runtime::Module是tvm::runtime::ModuleNode的运行时引用，因此我们需要先新定义一个ModuleNode的子类。

ModuleNode有以下虚方法：

  • type_key：返回标识该模块的字符串，例如CUDAModuleNode的type_key返回"cuda"。
  • GetFunction：返回用于运行该模块的PackedFunc（关于PackedFunc的介绍参见官方文档），例如CUDAModuleNode的GetFunction返回的PackedFunc利用CUDA Runtime API来加载运行保存在CUDAModuleNode中的PTX代码。
  • SaveToFile：将模块保存进文件。
  • SaveToBinary：将模块输出到二进制流。
  • GetSource：返回模块保存的源代码，例如CUDAModuleNode的GetSource返回其保存的CUDA C代码。

其中type_key和GetFunction为纯虚函数，每个ModuleNode的非抽象子类必须实现，否则实例化处会出现编译期错误。而另外三个虚方法如果子类没有覆写，则在被调用时会抛出运行时错误。

以hello模块的HelloModuleNode为例，该类实现了type_key、GetFunction、GetSource。该类构造时接收代码字符串，GetSource返回代码，GetFunction返回的PackedFunc输出代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class HelloModuleNode : public ModuleNode {
public:
  explicit HelloModuleNode(std::string code) : code_(std::move(code)) {}
  const char *type_key() const final {
    return "hello";
  }
  PackedFunc GetFunction(
      const std::string &name,
      const ObjectPtr<Object> &sptr_to_self) final {
    auto code = code_;
    return PackedFunc([code](TVMArgs args, TVMRetValue *ret) {
      std::cout << code << std::endl;
    });
  }
  std::string GetSource(const std::string &format) final {
    return code_;
  }
private:
  std::string code_;
};

实现翻译单元

因为前端传递过来的是TIR，因此还需要有一个翻译单元来将传递过来的TIR进行转换，以合适的形式传递给后端模块。

一般来说这个单元做的事就是代码生成，如CodeGenCUDA就将传入的TIR翻译为CUDA C代码传递给CUDAModuleNode。

作为示例的hello模块也加一个非常简单的代码生成单元CodeGenHello，该类继承自CodeGenC，所做的仅仅是在CodeGenC输出的前置声明语句中添加一个注释// hello tvm backend!：

1
2
3
4
5
6
7
class CodeGenHello : public CodeGenC {
public:
  std::string Finish()  {
    decl_stream << "// hello tvm backend!" << std::endl;
    return CodeGenC::Finish();
  }
};

注册接口函数

tvm.build和tvm.build_module.build最后一步的派发都发生在tvm::codegen::Build：

1
2
3
4
5
6
7
8
9
std::string build_f_name = "codegen.build_" + mode;
  // the build function.
  const PackedFunc* bf = runtime::Registry::Get(build_f_name);
  CHECK(bf != nullptr)
      << "Target " << target << " is not enabled";
  runtime::Module m = transformed_funcs.empty() ?
                      (*bf)(funcs, target) :
                      (*bf)(transformed_funcs, target);
  return m;

由其代码可见，如果想利用tvm.build等已有的前端接口构建模块，则需将一个函数签名为(tvm::Array<tvm::tir::LoweredFunc>, const std::string &) -> tvm::runtime::Module的后端接口函数注册为"codegen.build_${target}"。

如下为hello模块的后端接口函数及其注册：

1
2
3
4
5
6
7
8
9
Module BulidHello(Array<tir::LoweredFunc> funcs, const std::string &target) {
  codegen::CodeGenHello cg;
  for (auto &f : funcs)
    cg.AddFunction(f);
  return Module(make_object<HelloModuleNode>(cg.Finish()));
}

TVM_REGISTER_GLOBAL("codegen.build_hello")
.set_body_typed(BulidHello);

当然也可以注册其他类型的后端接口，不过这样可能需要自己添加其他函数来处理前端的输入（Relay、Tensor-Expression等），将前端输入转化为后端接口适配的类型。

添加设备类型

因为tvm.build的执行过程会提前进行一些和设备相关的检查（好像会检查这些的代码最终都会流向tvm::codegen::Build，而后者本身就会检查设备后端接口的存在性，所以这种强行增加模块耦合性的行为是为什么呢……），因此需要手动在源码中为我们的hello模块“放行”：

在python/tvm/_ffi/runtime_ctypes.py中的TVMContext.STR2MASK中添加一项：

1
2
3
4
5
6
7
8
9
10
11
STR2MASK = {
        'llvm': 1,
        'stackvm': 1,
        'cpu': 1,
        'c': 1,
        'hello': 1, # here
        'gpu': 2,
        'cuda': 2,
        # ......
        'micro_dev': 13,
    }

在src/target/target.cc中的tvm::CreateTarget里添加一项判断：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
if (target_name == "c" && t->device_name == "micro_dev") {
    t->device_type = kDLMicroDev;
  } else if (target_name == "c" || target_name == "llvm") {
    t->keys_array.push_back(tir::StringImmNode::make("cpu"));
  } else if (target_name == "cuda" || target_name == "nvptx") {
    t->device_type = kDLGPU;
    t->keys_array.push_back(tir::StringImmNode::make("cuda"));
    t->keys_array.push_back(tir::StringImmNode::make("gpu"));
    t->max_num_threads = 1024;
    t->thread_warp_size = 32;
  // ......
  } else if (target_name == "hybrid") {
    t->device_type = kDLCPU;
  } else if (target_name == "hello") { // here
    t->device_type = kDLCPU;
  } else {
    LOG(ERROR) << "Unknown target name " << target_name;
    return target::stackvm();
  }

改写cmake文件

最后一步就是将我们新增的模块集成到原本的项目构建系统中。

首先在cmake/modules/contrib/中添加Hello.cmake文件，将我们的源文件加入编译列表。

1
2
3
list(APPEND RUNTIME_SRCS src/runtime/contrib/hello/hello_module.cc)
list(APPEND COMPILER_SRCS src/runtime/contrib/hello/hello_module.cc)
message(STATUS "Build with Hello support")

在CMakeList.txt中添加一项：

1
include(cmake/modules/contrib/Hello.cmake)

至此大功告成。

总结

  • 可以针对Relay和TIR两种IR进行代码生成。推荐针对TIR生成代码，这样可以复用大量的前端/后端代码。

  • 对TIR生成代码主要需要以下步骤：

  • 实现后端模块，用于运行代码。

  • 实现翻译单元，用于代码生成。

  • 注册接口函数，给TVM其他部分调用。

  • 适当修改其他部分的代码。

  • 改写cmake文件。

附上hello_module.cc的内容：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
#include <string>
#include <tvm/runtime/module.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/node/container.h>
#include <tvm/runtime/registry.h>
#include "../../../target/source/codegen_c.h"
#include <iostream>

namespace tvm {

namespace codegen {

class CodeGenHello : public CodeGenC {
public:
  std::string Finish() {
    decl_stream << "// hello tvm backend!" << std::endl;
    return CodeGenC::Finish();
  }
};

}

namespace runtime {

class HelloModuleNode : public ModuleNode {
public:
  explicit HelloModuleNode(std::string code) : code_(std::move(code)) {}
  const char *type_key() const final {
    return "hello";
  }
  PackedFunc GetFunction(
      const std::string &name,
      const ObjectPtr<Object> &sptr_to_self) final {
    auto code = code_;
    return PackedFunc([code](TVMArgs args, TVMRetValue *ret) {
      std::cout << code << std::endl;
    });
  }
  std::string GetSource(const std::string &format) final {
    return code_;
  }
private:
  std::string code_;
};

Module BulidHello(Array<tir::LoweredFunc> funcs, const std::string &target) {
  codegen::CodeGenHello cg;
  for (auto &f : funcs)
    cg.AddFunction(f);
  return Module(make_object<HelloModuleNode>(cg.Finish()));
}

TVM_REGISTER_GLOBAL("codegen.build_hello")
.set_body_typed(BulidHello);

}

}
≡