TVM/VTA代码生成流程



# TVM/VTA代码生成流程 

    1.         [TVM](https://krantz-xrf.github.io/archive.html?tag=TVM)
    1.         [VTA](https://krantz-xrf.github.io/archive.html?tag=VTA)
    1.         [编程](https://krantz-xrf.github.io/archive.html?tag=%E7%BC%96%E7%A8%8B)
    1.         [机器学习](https://krantz-xrf.github.io/archive.html?tag=%E6%9C%BA%E5%99%A8%E5%AD%A6%E4%B9%A0)

  •  2019年 10月24日

最近看了很多TVM/VTA后端代码生成的代码，现在就把近日所得总结一下，以备有需求的朋友参考。

关于TVM/VTA

TVM是一个深度学习描述框架，通过Python代码描述算子（输入、输出、运算方法等）形成抽象语法树（Abstract Syntax Tree，AST），然后在TVM内部转换为中间表示（Intermediate Representation，IR），最终转换成目标平台的机器代码，以作为算子用于构成更复杂的神经网络。
VTA（Versatile Tensor Accelerator，多功能张量加速器）是TVM框架的一个扩展，可以简单理解成一个深度神经网络的底层硬件实现。
这篇文章就是TVM从IR生成后端机器代码过程的一个概览。
由于手头并没有VTA的硬件，因此使用了TVM提供的模拟器tsim；以下所述的过程都是针对tsim的，如果针对其他硬件后端，大体思想应当是一致的，但细节肯定颇有不同。

分析的代码样例

分析代码生成流程时，使用了官方教程提供的测试代码，详见Get Started with VTA。
以下描述的流程从vta.build的调用开始。

代码生成

vta.build首先判断出我们host端（宿主平台，主程序所运行的平台）使用llvm做代码生成（target_host='llvm'）；随后，它直接将调用转发给tvm.build。

整理输入

由于Python的特殊性，函数参数的类型是不定的；TVM允许传入tvm.build函数的参数有如下四种：

  • Schedule
  • LoweredFunc
  • [LoweredFunc]
  • {target: [LoweredFunc]}

转换经历下图的过程： Schedulelower−−−→LoweredFunc[∗]−→[LoweredFunc]{target:∗}−−−−−−→{target : [LoweredFunc]}(1)Schedule→lowerLoweredFunc→[∗][LoweredFunc]→{target:∗}{target : [LoweredFunc]}最终所有的输入都被整理成如下形式：

  1. target_flist = {'ext_dev': [LoweredFunc]}

降级代码表示

这一部分对应于函数lower，总的流程参看下图：

  1. ![](https://www.yuque.com/attachments/yuque/0/2020/svg+xml/644279/1594191727903-702a00a7-65f4-4734-8b30-b5ebc3eca873.svg+xml)

以我们的测试代码为例，每一趟后代码发生的变化如下表：

| 阶段 | 处理阶段 | 是否变化 | 发生的变化 | | —- | —- | —- | —- |

  1. | 
  2.   0
  3.    |  
  4.    |  
  5.    | 初始状态
  6.  |
  8. | 
  9.   1.1
  10.    | `StorageFlatten`
  11.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191727944-6320e00b-9b7f-4a57-ad1d-3ccd6ee5aaf5.png ":heavy_check_mark:")
  12.    |

realize -> allocate，指标的表示形式（多维转化为一维） |

  1. | 
  2.   1.2
  3.    | `CanonicalSimplify`
  4.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191727999-dcef415f-def1-461b-b25d-e17edd6f55ed.png ":heavy_check_mark:")
  5.    | 双层`for`循环 -> `TAStoreBuffer2D`
  7.  |
  9. | 
  10.   1.3
  11.    | （外部过程）
  12.    |  
  13.    |  
  14.  |
  16. | 
  17.   1.4
  18.    | （外部过程）
  19.    | ![:white_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728091-0d4964a9-2948-46de-9551-40ba43126eb5.png ":white_check_mark:")
  20.    | 增加了一些新属性
  21.  |
  23. | 
  24.   1.5
  25.    | （外部过程）
  26.    | ![:white_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728143-8f01d01b-b0e2-47f1-aa0a-b6e07d66560f.png ":white_check_mark:")
  27.    | 移除了一些属性
  28.  |
  30. | 
  31.   1.6
  32.    | （外部过程）
  33.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728218-e1ddbe9f-d532-456b-974f-64c3c7f88491.png ":heavy_check_mark:")
  34.    | 缓冲区内存分配从`produce`块中移出
  35.  |
  37. | 
  38.   1.7
  39.    | （外部过程）
  40.    |  
  41.    | 增加同步属性
  42.  |
  44. | 
  45.   1.8
  46.    | （外部过程）
  47.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728264-deffd343-d8dc-4c2b-9b81-8df4eec14db2.png ":heavy_check_mark:")
  48.    |

A、B、C的分配合并成A的分配 |

  1. | 
  2.   1.9
  3.    | （外部过程）
  4.    |  
  5.    |  
  6.  |
  8. | 
  9.   2.1
  10.    | `LoopPartition`
  11.    |  
  12.    |  
  13.  |
  15. | 
  16.   2.2
  17.    | `VectorizeLoop`
  18.    |  
  19.    |  
  20.  |
  22. | 
  23.   2.3
  24.    | `InjectVirtualThread`
  25.    |  
  26.    |  
  27.  |
  29. | 
  30.   2.4
  31.    | `InjectDoubleBuffer`
  32.    |  
  33.    |  
  34.  |
  36. | 
  37.   2.5
  38.    | `StorageRewrite`
  39.    |  
  40.    |  
  41.  |
  43. | 
  44.   2.6
  45.    | `UnrollLoop`
  46.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728325-4570bb38-41d6-42df-9771-5564df9d1ce0.png ":heavy_check_mark:")
  47.    | 循环转化为`VTAUopLoopBegin`、`VTAUopPush`和`VTAUopLoopEnd`
  49.  |
  51. | 
  52.   2.7
  53.    | （外部过程）
  54.    |  
  55.    |  
  56.  |
  58. | 
  59.   3.1
  60.    | `Simplify`
  61.    | ![:heavy_check_mark:](https://cdn.nlark.com/yuque/0/2020/png/644279/1594191728385-8b14554e-7550-4038-b79d-31b9c3bff193.png ":heavy_check_mark:")
  62.    | 缓冲区内存分配完全移除
  63.  |
  65. | 
  66.   3.2
  67.    | `LowerStorageAccessInfo`
  68.    |  
  69.    |  
  70.  |
  72. | 
  73.   3.3
  74.    | `RemoveNoOp`
  75.    |  
  76.    |  
  77.  |
  79. | 
  80.   3.4
  81.    | `RewriteUnsafeSelect`
  82.    |  
  83.    |  
  84.  |
  86. | 
  87.   3.5
  88.    | （外部过程）
  89.    |  
  90.    |  
  91.  |
  93. | 
  94.   3.6
  95.    | （外部过程）
  96.    |  
  97.    |  
  98.  |
  100. | 
  101.   4
  102.    |  
  103.    |  
  104.    | 最终状态
  105.  |

标明“（外部过程）”是从C++注册的处理过程，在Python的跟踪过程中无法看到。
这部分的中间结果文件可以在这里下载。

遍历target_list检查

对所有的目标⟨target,flist⟩∈target_flist⟨target,flist⟩∈target_flist：

  • 函数名查重：存在相同的函数名字就报错
  • 验证目标targettarget是str或者_target.Target

下面，flist被传入函数_build_for_device处理。

为特定设备生成目标代码

总的思想是这样的：将flist分离为宿主代码（fhost）和设备代码（mdev），然后分别生成机器代码；其中设备端的模块会导入到宿主模块中，最终的结果是宿主代码模块mhost。 flist_build_for_device−−−−−−−−−−→⎧⎪⎨⎪⎩fhostcodegen.build_module−−−−−−−−−−−−−→mdevimport_module−−−−−−−−−→⎫⎪⎬⎪⎭→mhost(2)flist→_build_for_device{fhost→codegen.build_modulemdev→import_module}→mhost下图是IR的多趟（pass）处理流程：

  1. ![](https://www.yuque.com/attachments/yuque/0/2020/svg+xml/644279/1594191728429-4e9b8933-8ee9-43e2-964a-081c968e6cfd.svg+xml)

加载生成好的目标代码

这部分对应的Python代码如下：

  1. remote.upload("vadd.o")
  2. f = remote.load_module("vadd.o")

流程概览

  1. ![](https://www.yuque.com/attachments/yuque/0/2020/svg+xml/644279/1594191728484-220a69e6-be8e-4a65-a7dd-c957ee72deb2.svg+xml)

下面代码分tsim和真实硬件两种情况。
相关分析只列出被执行的关键路径，依照代码中的注释应当很容易理解。代码块的缩进表示嵌套的函数调用。
贴出的代码有Python也有C++，由于每一段都有注释，应该很好分辨（Python是#，C++是//）。

使用模拟器

此时代码中的远端设备remote是一个LocalSession。
这一部分的关键就是拼接命令，调用系统编译器g++来把对象文件（.o文件）链接成动态库。

  1. # ..., then, in LocalSession.load_module
  2. #   with path = "vadd.o"
  3. _load_module(self._temp.relpath(path))
  4. # _load_module is module.load
  1. # in module.load, with path = (full path for "vadd.o"), fmt = ""
  2. if path.endswith(".o"): # true
  3.     _cc.create_shared(path + ".so", path)
  • ```

    in create_shared

    with output = “vadd.o.so”

    objects = “vadd.o”

    options = None

    cc = “g++”

    if sys.platform == “darwin” or sys.platform.startswith(“linux”): _linux_compile(output, objects, options, cc)
  3.     -

in _linux_compile

with output = “vadd.o.so”

objects = “vadd.o”

options = None

compile_cmd = “g++”

cmd = [compile_cmd] # cmd: g++ if output.endswith(“.so”): # true cmd += [“-shared”, “-fPIC”] if sys.platform == “darwin”: # true cmd += [“-undefined”, “dynamic_lookup”] else: # false, …

cmd: g++ -shared -fPIC -undefined dynamic_lookup

cmd += [“-o”, output]

cmd: g++ -shared -fPIC -undefined dynamic_lookup -o vadd.o.so

if isinstance(objects, str): # true cmd += [objects] else: # false, …

cmd: g++ -shared -fPIC -undefined dynamic_lookup -o vadd.o.so vadd.o

if options: # false, …

run cmd

proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) (out, _) = proc.communicate() if proc.returncode != 0: msg = “Compilation error:\n” msg += py_str(out) raise RuntimeError(msg)

back in create_shared

else: # false, …

back in module.load

  1. path += ".so"

else: # false

  1.    # ...

return _LoadFromFile(path, fmt)

  1. `_LoadFromFile`是Python端封装的C++函数，在C++端对应`Module::LoadFromFile`，见[最终加载动态库](https://krantz-xrf.github.io/2019/10/24/tvm-workflow.html#%E6%9C%80%E7%BB%88%E5%8A%A0%E8%BD%BD%E5%8A%A8%E6%80%81%E5%BA%93)。
  2. ### 使用真实的VTA设备[]()
  3. 这里，由于没有真实设备，执行流程是~~静态分析~~瞪眼猜测的结果。<br />此时远端设备`remote`是一个`RPCSession`。<br />在`LoadRemoteModule`函数执行前也应该有一些额外的操作，把`.o`对象文件生成为动态链接库。

// in “rpc._LoadRemoteModule” sess->CallRemote(RPCCode::kModuleLoad, args[1]);

  1. ```
  2. // in RPCSession::HandlePackedCall
  3. switch (code_)
  4. {
  5.     // ...
  6.     case RPCCode::kModuleLoad: CallHandler(RPCModuleLoad); break;
  7.     // ...
  8. }
  1. // in RPCModuleLoad
  2. fsys_load_ = runtime::Registry::Get("tvm.rpc.server.load_module");
  3. /* ... */ (*fsys_load_)(file_name);
  1. // in "tvm.rpc.server.load_module"
  2. // Below is Objective-C++:
  3. //   - not quite familiar
  4. //   - might misinterpret
  5. // in "tvm.rpc.server.load_module", with name = "vadd.o.so"s
  6. std::string fmt = GetFileFormat(name, "");
  • ``` // in tvm::runtime::GetFileFormat // with file_name = “vadd.o.so”s // format = “”s if (format.length() == 0) { // true // … size_t pos = file_name.find_last_of(“.”); // 6 if (pos != std::string::npos) { // true
    1.     return file_name.substr(pos + 1, file_name.length() - pos - 1);
    2.     // "vadd.o.so"s.substr[from: 6, length: 2] = "so"s
    } // … } // …

// fmt = “so”s // … converting name to path // not quite sure because of use of Obj-C++ NSString path = [base stringByAppendingPathComponent: [NSString stringWithUTF8String:name.c_str()]]; // … and again back to name name = [path UTF8String]; // finally! loading from file? // - no! yet another propagation / … */ Module::LoadFromFile(name, fmt);

  1. ### 最终加载动态库[]()
  2. 这一部分核心就是转化为系统调用`dlopen`（POSIX系统）或`LoadLibraryW`（Windows系统）。

// in Module::LoadFromFile // with filename = … (“vadd.o.so” with full path) // format = “so”s std::string fmt = GetFileFormat(file_name, format); // “so”s.length() != 0, should just return “so”s // fmt = “so”s if (fmt == “dll” || fmt == “dylib” || fmt == “dso”) { / … / } // false std::string load_f_name = “module.loadfile“ + fmt; // load_f_name = “module.loadfile_so”s f = Registry::Get(load_f_name); / … / (*f)(file_name, format);

  1. ```
  2. // in "module.loadfile_so"
  3. n = std::make_shared<DSOModuleNode>();
  4. n->Init(args[0]);
  1. // in DSOModuleNode::Init
  2. DSOModuleNode::Load(name); // propagate to LoadLibraryW/dlopen
  3. // ...
  4. InitContextFunctions([this](const char* fname) { return GetSymbol(fname); });
  5. // Load the imported modules
  6. const char* dev_mblob = GetSymbol(runtime::symbol::tvm_dev_mblob);
  7. if (dev_mblob != nullptr) { /* ... */ }

订阅本文遵守 Attribution-NonCommercial-ShareAlike 4.0 International 许可协议。 TVM/VTA代码生成流程 - 图1

上篇开始下篇PyBind11：基本用法和底层实现 0 条评论未登录用户TVM/VTA代码生成流程 - 图2 TVM/VTA代码生成流程 - 图3 TVM/VTA代码生成流程 - 图4 支持 Markdown 语法 来做第一个留言的人吧！