Micro TVM

library 正常是通过 DSO 的形式保存, 并在运行时通过 dlsym 方式来查找符号, system-lib 机制是为了解决有些系统不支持 dlopen, dlsym 的问题. 如果无法用 dlsym 来查找, 那用什么方式来找到 `tvmgen_default_fused_add`?

有两种方案: c++ runtime 或 c runtime

import tvm
from tvm import relay

x = relay.var("x", shape=(1, 1000), dtype="float32")
y = relay.add(x, x)
func = relay.Function([x], y)
mod = tvm.IRModule.from_expr(func)
print(mod)

with tvm.transform.PassContext(opt_level=1):
    mod = relay.build(mod, target="llvm --system-lib --runtime=c++", params=None)

mod.lib.save("/tmp/a.o")

def @main(%x: Tensor[(1, 1000), float32]) { add(%x, %x) }

readelf -s /tmp/a.o
echo "------"
readelf -a /tmp/a.o|grep "Relocation section '.rela.ctors'" -A 3

Symbol table '.symtab' contains 14 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS TVMMod 2: 0000000000000330 40 FUNC LOCAL DEFAULT 2 _tvm_module_startup 3: 0000000000000210 284 FUNC LOCAL DEFAULT 2 tvmgen_default_fused_add 4: 0000000000000000 0 SECTION LOCAL DEFAULT 2 5: 0000000000000000 0 SECTION LOCAL DEFAULT 5 6: 0000000000000000 0 SECTION LOCAL DEFAULT 8 7: 0000000000000000 0 SECTION LOCAL DEFAULT 9 8: 0000000000000000 0 SECTION LOCAL DEFAULT 10 9: 0000000000000000 0 SECTION LOCAL DEFAULT 19 10: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMAPISetLastError 11: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMBackendRegisterSystemL 12: 0000000000000000 8 OBJECT WEAK DEFAULT 4 __tvm_module_ctx 13: 0000000000000000 523 FUNC GLOBAL DEFAULT 2 tvmgen_default_fused_add

Relocation section '.rela.ctors' at offset 0x1228 contains 1 entry: Offset Info Type Sym. Value Sym. Name + Addend 000000000000 000400000001 R_X86_64_64 0000000000000000 .text + 330

相比没有加 system-lib 的情形, a.o 多了几个 symbol:

1. __tvm_module_startup

   利用 elf 的 ctor 在应用启动时调用 __tvm_module_startup

2. TVMBackendRegisterSystemLibSymbol

   __tvm_module_startup 的代码会通过 TVMBackendRegisterSystemLibSymbol("tvmgen_default_fused_add",tvmgen_default_fused_add) 的方式把 tvmgen_default_fused_add 函数注册到一个 runtime 负责的全局的 `SystemLibrary` 上

当 TVM 需要跑在开发板上时，会有几个问题：

1. 板子可能没有支持 ctor 的环境
2. 板子上关于内存分配，log 等功能需要有特定的实现

为此，TVM 以提出一个 c runtime.

c runtime 的目标场景是那些无法支持 ctor (即 .init/.init_array) 的环境, 例如不支持 rtld 或没有 startfiles 的情况 init_array 如何被调用

c runtime 无法利用 ctor 来注册 mapping, 所以它实现了一个函数 `TVMSystemLibEntryPoint`, 这个函数会返回一个包含 {names, functions} 的结构体,用来初始化 SystemLib

C Runtime 会利用 tvm::codegen::CreateMetadataModule 完成与 c++ 的 __tvm_module_startup 类似的功能.

tvm/src/relay/backend/build_module.cc::ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods, GetTargetHost(),

import tvm
from tvm import relay

x = relay.var("x", shape=(1, 1000), dtype="float32")
y = relay.add(x, x)
func = relay.Function([x], y)
mod = tvm.IRModule.from_expr(func)

with tvm.transform.PassContext(opt_level=1):
    mod = relay.build(mod, target="llvm  --system-lib --runtime=c", params=None)

mod.lib.export_library("/tmp/a.tar")

tar zxvf /tmp/a.tar
echo "------lib0.o"
readelf -s /tmp/lib0.o
echo "------lib1.o"
readelf -s /tmp/lib1.o

lib0.o devc.o lib1.o -–—lib0.o

Symbol table '.symtab' contains 15 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS TVMMetadataMod 2: 0000000000000008 16 OBJECT LOCAL DEFAULT 4 _tvm_crt_func_registry 3: 0000000000000018 8 OBJECT LOCAL DEFAULT 4 _tvm_crt_module 4: 0000000000000000 8 OBJECT LOCAL DEFAULT 4 _tvm_func_registry_ptrs 5: 0000000000000000 0 SECTION LOCAL DEFAULT 2 6: 0000000000000000 0 SECTION LOCAL DEFAULT 4 7: 0000000000000000 0 SECTION LOCAL DEFAULT 6 8: 0000000000000000 0 SECTION LOCAL DEFAULT 7 9: 0000000000000000 0 SECTION LOCAL DEFAULT 8 10: 0000000000000000 0 SECTION LOCAL DEFAULT 9 11: 0000000000000000 0 SECTION LOCAL DEFAULT 16 12: 0000000000000000 0 SECTION LOCAL DEFAULT 18 13: 0000000000000000 8 FUNC GLOBAL DEFAULT 2 TVMSystemLibEntryPoint 14: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND tvmgen_default_fused_add -–—lib1.o

Symbol table '.symtab' contains 11 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS TVMMod 2: 0000000000000210 284 FUNC LOCAL DEFAULT 2 tvmgen_default_fused_add_ 3: 0000000000000000 0 SECTION LOCAL DEFAULT 2 4: 0000000000000000 0 SECTION LOCAL DEFAULT 4 5: 0000000000000000 0 SECTION LOCAL DEFAULT 5 6: 0000000000000000 0 SECTION LOCAL DEFAULT 6 7: 0000000000000000 0 SECTION LOCAL DEFAULT 7 8: 0000000000000000 0 SECTION LOCAL DEFAULT 16 9: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMAPISetLastError 10: 0000000000000000 523 FUNC GLOBAL DEFAULT 2 tvmgen_default_fused_add

• lib1.o 有 tvmgen_default_fused_add 的实现
• lib0.o 中的 TVMSystemLibEntryPoint 是用来返回 {"names",funcs} 的数据的

图中的 unk_28 实际上就是指向 .rodata 上的字符串 "tvmgen_default_fused_add", 所以 lib0.o 中的 TVMSystemLibEntryPoint 最终会返回

{"tvmgen_default_fused_add", tvmgen_default_fused_add}

// crt_runtime_api.c
int SystemLibraryCreate(TVMValue* args, int* type_codes, int num_args, TVMValue* ret_val,
                        int* ret_type_codes) {

  const TVMModule* system_lib;
  system_lib = TVMSystemLibEntryPoint();
  TVMModCreateFromCModule(system_lib, &system_lib_handle);

  ret_val[0].v_handle = system_lib_handle;
  ret_type_codes[0] = kTVMModuleHandle;
  return 0;
}

file:~/Gitbox/code/hello_world/hello_tvm/cpp_runtime

file:~/Gitbox/code/hello_world/hello_tvm/c_runtime

1. 在 target 中使用 mtriple, mcpu 指定, 由 llvm 负责交叉编译.

   mod = relay.build(
       mod,
       target="llvm  --system-lib --runtime=c, -mtriple=armv8l-linux-gnueabihf -mcpu=cortex-m7",
       params=None,
   )
   

2. 使用 c target, 生成 c 代码后自己做交叉编译

   但看起来 c target 没有生成向量化的代码

RPC 的结构为:

                                         |
Session (Flasher/Transport) [python]     |
   -> RPCSession [c++]                   |
     -> MicroTransportChannel [c++]      ->   MicroTVMRpcServerLoop
                                         |        -> Session
                                         |           -> MinRPCServer
                                         |               -> HandleNormalCallFunc, ...
                                         |

2.1 (system-lib & c++ runtime) 是支持注册 __tvm_dev_mblob 的

def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=None, **kwargs):
    # ...
    if self.imported_modules:
        if enabled("llvm") and llvm_target_triple:
            path_obj = os.path.join(workspace_dir, f"devc.{object_format}")
            m = _ffi_api.ModulePackImportsToLLVM(self, is_system_lib, llvm_target_triple)
            m.save(path_obj)
            files.append(path_obj)
        else:
            path_cc = os.path.join(workspace_dir, "devc.c")
            with open(path_cc, "w") as f:
                f.write(_ffi_api.ModulePackImportsToC(self, is_system_lib))
            files.append(path_cc)

ModulePackImportsToLLVM 和 ModulePackImportsToC 会生成类似下面的代码:

extern int TVMBackendRegisterSystemLibSymbol(const char*, void*);
static int __tvm_dev_mblob_reg_ = TVMBackendRegisterSystemLibSymbol(
    "__tvm_dev_mblob", (void*)__tvm_dev_mblob);

通过 c++ 的 ctor 完成对 TVMBackendRegisterSystemLibSymbol 的调用.

从当前的代码来看:

1. 在 export_library 时不会针对 c_runtime 生成注册 __tvm_dev_mblob 的代码 (例如 ModulePackImportsToLLVM)
2. 运行时 c runtime 也不会解析 __tvm_dev_mblob
3. 以及一些其它的问题例如: tvm/src/target/metadata_module.cc::if (is_targeting_crt)

所以 microtvm 无法支持 BYOC.

但是如果 c runtime 只是为了解决 c++ runtime ctor 的问题, 我们其实可以手动调用 c++ ctor:

#include <stdio.h>
class A {
   public:
    A() { printf("hello\n"); }
};

A a, b, c;

typedef void (*FP)();

// __init_array_start, __init_array_end 是通过 linker script 定义的符号
extern FP __init_array_start;
extern FP __init_array_end;

int main(int argc, char *argv[]) {
    printf("----------main----------\n");
    for (FP *p = &__init_array_start; p < &__init_array_end; p++) {
        (*p)();
    }
    return 0;
}

hello hello hello -----–—main-----–— hello hello hello

ModulePackImportsToLLVM 时如果 is_system_lib 为 False, 则不会生成注册 __tvm_dev_mblob 的代码, 但 is_system_lib 的判断条件是:

is_system_lib = (
    module.type_key == "llvm" and module.get_function("__tvm_is_system_module")()
)

PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
    if (name == "__tvm_is_system_module") {
      // 判断有没有 __tvm_module_startup        
      bool flag = (mptr_->getFunction("__tvm_module_startup") != nullptr);
      return PackedFunc([flag](TVMArgs args, TVMRetValue* rv) { *rv = flag; });
    }
    // ...
}

// 添加 __tvm_module_startup 的代码
void CodeGenCPU::AddStartupFunction() {
  if (!target_c_runtime_) {
    llvm::FunctionType* ftype = llvm::FunctionType::get(t_void_, {}, false);
    function_ = llvm::Function::Create(ftype, llvm::Function::InternalLinkage,
                                       "__tvm_module_startup", module_.get());
    for (const auto& kv : export_system_symbols_) {
      llvm::Value* name = GetConstString(kv.first);
      builder_->CreateCall(f_tvm_register_system_symbol_,
                           {name, builder_->CreateBitCast(kv.second, t_void_p_)});
    }    
    // ...
  }
}

std::unique_ptr<llvm::Module> CodeGenLLVM::Finish() {
  this->AddStartupFunction();
  // ...
  // optimize
  this->Optimize();
  return std::move(module_);
}

问题出在最后那个 llvm::Optimize

当 relay IR 时没有任何函数需要在 CPU 执行时 (所有代码都交给 dnnl 执行), export_system_symbols_ 会为空, 这时虽然已经在 llvm module 里生成了一个 __tvm_module_startup 函数, 但因为它的函数体为空且函数是 InternalLinkage, llvm::Optimize 会把它删除…

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-08-03 11:11
import tvm
import numpy as np
from tvm import te
from tvm import relay
import os


def get_good_mod():
    d1 = relay.var("d1", shape=(1, 32, 56, 56), dtype="float32")
    w1 = relay.var("w1", shape=(32, 32, 3, 3), dtype="float32")
    b1 = relay.var("b1", shape=(32,), dtype="float32")
    conv = relay.nn.conv2d(d1, w1, strides=(1, 1), padding=(1, 1))
    bias = relay.nn.bias_add(conv, b1)
    relu = relay.nn.relu(bias)

    func = relay.Function([d1, w1, b1], relu)
    mod = tvm.IRModule.from_expr(func)
    mod = relay.transform.InferType()(mod)
    return mod


def get_bug_mod():
    x = relay.var("x", shape=(10,), dtype="float32")
    y = relay.add(x, x)
    func = relay.Function([x], y)
    mod = tvm.IRModule.from_expr(func)
    return mod


def prepare_relay_libs():
    mod = get_bug_mod()
    mod = relay.transform.AnnotateTarget("dnnl")(mod)
    mod = relay.transform.PartitionGraph()(mod)
    with tvm.transform.PassContext(opt_level=1):
        mod = relay.build(mod, target="llvm --system-lib --runtime=c++", params=None)
    mod.lib.export_library("/tmp/libdouble_bug.tar")

    mod = get_good_mod()
    mod = relay.transform.AnnotateTarget("dnnl")(mod)
    mod = relay.transform.PartitionGraph()(mod)
    with tvm.transform.PassContext(opt_level=1):
        mod = relay.build(mod, target="llvm --system-lib --runtime=c++", params=None)
    mod.lib.export_library("/tmp/libdouble_good.tar")


if __name__ == "__main__":
    prepare_relay_libs()

echo "----------good----------"
tar xvf /tmp/libdouble_good.tar -C /tmp
readelf -s /tmp/lib0.o
readelf -s /tmp/devc.o
echo "----------bug----------"
tar xvf /tmp/libdouble_bug.tar -C /tmp
readelf -s /tmp/lib0.o
readelf -s /tmp/devc.o

-----–—good-----–— lib0.o devc.o

Symbol table '.symtab' contains 16 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS TVMMod 2: 0000000000000000 0 NOTYPE LOCAL DEFAULT 4 .LCPI0_0 3: 00000000000005f0 40 FUNC LOCAL DEFAULT 2 __tvm_module_startup 4: 0000000000000370 40 FUNC LOCAL DEFAULT 2 tvmgen_default_fused_nn_b 5: 0000000000000000 0 SECTION LOCAL DEFAULT 2 6: 0000000000000000 0 SECTION LOCAL DEFAULT 6 7: 0000000000000000 0 SECTION LOCAL DEFAULT 9 8: 0000000000000000 0 SECTION LOCAL DEFAULT 10 9: 0000000000000000 0 SECTION LOCAL DEFAULT 11 10: 0000000000000000 0 SECTION LOCAL DEFAULT 20 11: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMAPISetLastError 12: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMBackendParallelLaunch 13: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMBackendRegisterSystemL 14: 0000000000000000 8 OBJECT WEAK DEFAULT 5 __tvm_module_ctx 15: 0000000000000000 869 FUNC GLOBAL DEFAULT 2 tvmgen_default_fused_nn_b

Symbol table '.symtab' contains 10 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS devc 2: 0000000000000000 8 FUNC LOCAL DEFAULT 3 GLOBAL__sub_I_devc 3: 0000000000000010 28 FUNC LOCAL DEFAULT 3 __cxx_global_var_init 4: 0000000000000000 4 OBJECT LOCAL DEFAULT 6 __tvm_dev_mblob_reg 5: 0000000000000000 0 SECTION LOCAL DEFAULT 3 6: 0000000000000000 0 SECTION LOCAL DEFAULT 5 7: 0000000000000000 0 SECTION LOCAL DEFAULT 6 8: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND TVMBackendRegisterSystemL 9: 0000000000000000 2990 OBJECT GLOBAL DEFAULT 5 __tvm_dev_mblob -----–—bug-----–— lib0.o devc.o

Symbol table '.symtab' contains 6 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS TVMMod 2: 0000000000000000 0 SECTION LOCAL DEFAULT 3 3: 0000000000000000 0 SECTION LOCAL DEFAULT 4 4: 0000000000000000 0 SECTION LOCAL DEFAULT 5 5: 0000000000000000 0 SECTION LOCAL DEFAULT 12

Symbol table '.symtab' contains 3 entries: Num: Value Size Type Bind Vis Ndx Name 0: 0000000000000000 0 NOTYPE LOCAL DEFAULT UND 1: 0000000000000000 0 FILE LOCAL DEFAULT ABS devc 2: 0000000000000000 1015 OBJECT GLOBAL DEFAULT 3 __tvm_dev_mblob