tvm 构建结果调试

背景

使用tvm的irbuilder直接构建了一小段计算程序,运行时有段错误。
tvm对内存边界的检查在较为上层的流程中,直接构建ir无法自动做检查。

尝试的方法

添加调试信息

把target设置为llvm,在tvm build后使用print(build_f.get_source()),可以打印出程序的llvm-ir。
但是,由于tvm是直接发射的llvm-ir,没有上层的文本,所以也没有发射有意义的调试信息。

将LLVM-IR反向翻译回C/C++

由于tvm可以给出llvm-ir的文本形式,如果能将llvm-ir翻译会c/c++代码,那么调试错误(甚至使用sanitizer系列工具)会简单很多。

llvm曾经有一个c/c++ 后端,可以将llvm-ir翻译为c/c++代码。但是后面由于维护问题,主线已经将该功能移除。
经过搜索,发现有一个第三方的项目https://github.com/JuliaComputing/llvm-cbe,可以与llvm-8一起工作。

实际测试发现,使用clang 编译出的hello world llvm-ir能正常翻译会C,但是tvm给出的ir会导致cbe工具出现assert错误。
简单看了一下,可能是tvm用的llvm-ir特性较新,cbe的llvm8还不认识。

从搜索过程来看,cbe的支持和需求都不是很强烈,于是不再考虑投入时间来分析和解决其问题。

使用TVM 的C backend

TVM使用cuda时,其输出的就是文本形式c代码交给cuda编译器。
这样看来,其支持一个c代码的backend应该是顺理成章的事情。

浏览tvm代码,果然其已经存在c backend了。参考https://tvm.apache.org/docs/dev/relay_bring_your_own_codegen.html ,生成c 代码除了帮助调试外,还比较容易与定制的优化库交互,甚至也可以把它当成代码模板,再进行人工修改。

使用c backend也比较简单,只需在tvm.build是传入target=’c’即可。

1
2
build_f  = tvm.build(ls, [], target='c',  name='prune_conv', binds=None)
print(build_f.get_source())

实际使用时,发现其c backend还不支持分配作用域为local的storeage。

1
2
3
4
5
6
void CodeGenC::PrintStorageSync(const CallNode* op) {  // NOLINT(*)
}

void CodeGenC::PrintStorageScope(const std::string& scope, std::ostream& os) { // NOLINT(*)
CHECK_EQ(scope, "global");
}

临时将程序内的 irb.allocate 调用中的 scope=’local’ 改为’global’,可以正常的输出c代码了,如下所示。

1
2
3
4
5
6
7
#include "tvm/runtime/c_runtime_api.h"
#include "tvm/runtime/c_backend_api.h"
void* __tvm_module_ctx = NULL;
#ifdef __cplusplus
extern "C"
#endif
TVM_DLL int32_t main(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode) {

作为调试,这里打印出来已经能满足要求了。
如果需要后续的定制开发,可以参考tvm自带的示例apps/howto_deploy/cpp_deploy.cc 。把tvm生成的代码作为与自己的程序进一步组合。

利用反向调试,直接分析生成的汇编

gdb支持反向调试,虽然其功能不太完善稳定,但是比较适合tvm生成的这类逻辑较为简单的场景。

1
2
3
4
5
6
7
8
b thread_pool.cc中的launch函数
c
停下后,set scheduler-locking on
si进入tvm生成的运算函数
record full
si走到故障处
然后就可以用reverse-stepi等进行反向调试
退出前可以record stop

注意record full是,一定要set scheduler-locking on。
因为full模式需要获取进程的所有信息,而当前gdb还没支持好多线程程序的序列执行功能。不锁住当前线程执行,会导致gdb报下面的assert。

1
2
../../gdb/nat/x86-linux-dregs.c:146: internal-error: void x86_linux_update_debug_registers(lwp_info*): Assertion `lwp_is_stopped (lwp)' failed.
A problem internal to GDB has been detected,

对于简单的程序,知道
可以在python端打印出tvm.nd.array的_tvm_handle。
在gdb中x _tvm_handle的数值,可以找到array对应的数据区地址。
结合反汇编,也可以进行简单的分析。

TVM C backend生成代码直接与python的runtime整合

前面提到可以使用tvm的c backend将tvm的结果输出为c文件,但是对于分析问题来说,还是能直接运行起来更为方便。
最为直接的两个需求就是:使用asan查找内存越界故障, 使用vtune/perf等查找性能瓶颈。
前期本来打算使用c++ runtime来调用c backend的代码,但是工作量稍大,和python这端的配合也比较麻烦。
经过调试分析,发现可以直接将c backend生成的代码放到python运行时中运行。
大致方法如下:

生成c 文件

使用前面介绍的方法,去除local指定后,可以在target=c的情况下,生成出c代码。如下实例命令。

1
2
3
4
build_f = tvm.build(ls, [Input, Offset, Reorder, Index, Stride, Weight, Output],  
target='c', name='prune_conv', binds=None)
with open( './out_csrc.c', 'w' ) as f:
f.write(build_f.get_source())

修改获得的c代码

tvm的runtime装载模块时,希望看到一个名为tvm_main的符号,并且该符号中应该以字符串形式存放模块的实际入口。
所以,需要稍微修改一下前面获得的c代码。将其头部稍微修改一下,将main改为main_t(clang将main视为特殊符号,参数如果和c标准不一致会拒绝编译),并且添加tvm_main符号。

1
2
const char __tvm_main__[] = "main_t";
TVM_DLL int32_t main_t(void* args, void* arg_type_ids, int32_t num_args, void* out_ret_value, void* out_ret_tcode) {

使用clang或者gcc将获得的c文件编译为so

然后就可以使用clang将文件编译为so文件了,此时可以加上调试信息和asan等功能

1
clang ./out_csrc.c  -I ../incubator-tvm/include/ -I ../incubator-tvm/3rdparty/dlpack/include/  -fsanitize=address -save-temps -fPIC -O0 -g3 -shared

在python中运行获得的so

参考面的示例代码,可以将前面获得的so装入python中运行。

1
2
3
4
mod_prune = '/mnt/d/opensource/tvm_files/out_csrc.so'
loaded_prune = tvm.runtime.load_module(mod_prune)
evaluator = loaded_prune.time_evaluator(loaded_prune.entry_name, ctx, number=10)
print('OPt_paper: %f' % evaluator(input_np, Offset_np, Reorder_np, Index_np, Stride_np, Weight_np, output_array).mean)

初步结论

可以考虑优先使用tvm 的c backend进行调试分析。
结合反向调试的直接汇编分析也有帮助。