在macOS上用LLDB调试TensorFlow源码

TensorFlow架构

TensorFlow由两个部分组成,一部分是面向使用者的操作接口,支持多种语言,目前主要是C++和Python;另一部分是核心运行时,由C++编写,负责执行使用者通过操作接口定义的计算图。

对于一般需求,使用者基本上不需要考虑底层C++代码是怎么执行的。但如果需要添加自定义操作,或者仅仅出于好奇,想了解TensorFlow的内部运行机制,就需要一种方式来调试TensorFlow的核心C++部分。

调试环境准备

调试TensorFlow源码本身其实并不是多么苦难的事情,但是环境准备比较麻烦,对于不熟悉的人来说,还是要耗费很多时间。调试环境主要包扩:

  • Xcode Command Line Tools
  • Debug版本的Python
  • Debug版本的TensorFlow

Xcode Command Line Tools

Xcode Command Line Tools提供了macOS上所需的几乎所有命令行工具,包括LLVM系列工具,gcc,m4,yacc等开发工具,当然也包括后面要用到的lldb。要安装Xcode Command Line Tools,需要在命令行执行:

1
~ xcode-select --install

在弹出框中选择“安装”,然后就等着进度条走完就好了。

获取Debug版本Python

尽管可以通过运行C++代码的测试用例来调试,但是常用而且方便的做法,是通过Python运行TensorFlow,对运行的程序进行调试。这里调试Python代码并不是调试Python逻辑,而是调试底层运行时的运行情况,具体来说,本文的Python环境是CPython,所以需要CPython的调试信息,才能在调试器中看到代码,而不是汇编代码。可以通过pyenv安装Python。

1
2
3
4
5
6
7
8
9
10
11
~ pyenv install -l # 列出所有支持的版本
...
3.5.1
3.5.2
...
~ pyenv install -g 3.5.2 # -g表示安装Debug版本
~ pyenv versions
* system (set by /Users/dtong/.pyenv/version)
3.5.1
3.5.2
3.5.2-debug

可以通过pyenv选择某个目录下的Python版本。但更好的办法,是通过virtualenv创建Python虚拟环境,在创建的时候通过--python=<path to python binary>参数指定要使用的Python,比如:

1
~ mkvirtualenv tensorflow --python=/Users/dtong/.pyenv/versions/3.5.2-debug/bin/python

其中mkvirtualenv是虚拟环境管理工具virtualenvwrapper工具中提供的命令。

Debug版本的TensorFlow

带有带有Debug信息的TensorFlow需要通过源码进行编译。TensorFlow的构建管理没有使用传统的Makefile,而是用了Google自己开发的bazel,可以通过brew install bazel来安装。

先用git将TensorFlow的源码clone下来:

1
~ git clone https://github.com/tensorflow/tensorflow

进入目录,切换Python虚拟环境,安装编译所需的依赖:

1
2
3
~ cd tensorflow
~ workon tensorflow
(tensorflow) ~ pip install numpy

执行./configure进行配置,如果你也是使用CPU版本,则一路回车就可以了。配置的最后,会下载编译需要的依赖库,所以需要等一会儿。然后使用bazel正式开始编译:

1
~ bazel build -c dbg //tensorflow/tools/pip_package:build_pip_package

需要注意的是,根据官方文档的介绍,bazel build-c参数传的是opt,也就是优化编译,对应的是g++的-O2 -DNDEBUG选项,不仅开启优化选项,还禁用了所有的assert(),以提高性能。而我们需要完整的调试信息,所以要使用-c dbg,即g++的-g选项,这样就可以使用调试器了。

编译大概需要十几分钟的时间,这个编译出来的是Python形式的包,需要使用pip进行安装:

1
2
3
4
5
~ mkdir _python_build
~ cd _python_build
~ ln -s ../bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/* .
~ ln -s ../tensorflow/tools/pip_package/* .
~ python setup.py develop

到此,TensorFlow的调试环境基本就准备好了。

使用LLDB进行调试

说到调试C/C++代码,几乎第一个想到的一定是GDB。然而GDB在macOS上有很多问题,比如处于安全性考虑,如果一个进程要访问另一个进程的内存空间,需要提供codesign签名。解决了签名问题,又出现无法加载动态库、找不到调试信息的问题,对此众说纷纭,总而言之,GDB对新版本的macOS支持的不好,而苹果也不准备提供定制版本,因为有更好的替代者,LLDB

对于LLDB的介绍,推荐看WWDC 2012上的视频,里面介绍了作为重新设计的LLDB,都有哪些比GDB强大的特性。这里简单的列出比较有特点的:

  • 语法更加一致,所有的命令都是<noun> <verb> [-options [option-value]] [argument [argument...]]这样的形式
  • 概念层次更加清楚,target->process->thread->frame->variable,同一个回话支持多个target,同一回话下的不同target,可以共享系统动态库,节省内存
  • 和编译器相结合,能够提供更准确的类型信息,而GDB是在内部自己管理一套类型,在查看类型信息时可能不准确
  • 具有更好的变量格式化能力,不仅可以修改变量的输出格式,如二进制、八进制等,还支持用Python脚本自定义打印格式,苹果提供了C++标准库类型的打印格式作为示例
  • 支持在调试回话过程中,执行表达式,这里的表达式会通过集成进来的编译器,编译成语法树进行计算,相比GDB要更加准确

总结起来就是,LLDB是未来,GDB已经过时。

使用LLDB调试TensorFlow很简单,首先,运行Python解释器,然后加载TensorFlow,并获取进程ID。

1
2
3
4
5
~ python
>>> import tensorflow as tf
>>> import os
>>> os.getpid()
65685

TensorFlow的Python代码通过加载_pywrap_tensorflow.so,来和TensorFlow的C++执行器进行交互,TensorFlow的C++代码和Python代码运行在同一个进程内。因此,可以使用LLDB对这个进程进行调试:

1
2
3
4
5
6
7
8
9
10
11
12
13
~ lldb -p 65685
Process 65685 stopped
* thread #1: tid = 0x31e966, 0x00007fffc4d83f4e libsystem_kernel.dylib`__select + 10, queue = 'com.apple.main-thread', stop reason = signal SIGSTOP
frame #0: 0x00007fffc4d83f4e libsystem_kernel.dylib`__select + 10
libsystem_kernel.dylib`__select:
-> 0x7fffc4d83f4e <+10>: jae 0x7fffc4d83f58 ; <+20>
0x7fffc4d83f50 <+12>: movq %rax, %rdi
0x7fffc4d83f53 <+15>: jmp 0x7fffc4d7cd94 ; cerror
0x7fffc4d83f58 <+20>: retq
Executable module set to "/Users/dtong/.virtualenvs/tensorflow/bin/python".
Architecture set to: x86_64h-apple-macosx.
(lldb)

设定断点:

1
2
3
4
(lldb) breakpoint set --name TF_NewSession
Breakpoint 1: where = _pywrap_tensorflow.so`::TF_NewSession(const TF_SessionOptions *, TF_Status *) + 19 at c_api.cc:262, address = 0x0000000123c66303
(lldb) continue
Process 65685 resuming

TF_NewSession会在TensorFlow创建Session时触发。

1
>>> sess = tf.Session()

在LLDB中断点被触发:

1
2
3
4
5
6
7
8
9
10
11
Process 65685 stopped
* thread #1: tid = 0x31e966, 0x0000000123c66303 _pywrap_tensorflow.so`::TF_NewSession(opt=0x00007ffcf56f8e20, status=0x00007ffcf510a800) + 19 at c_api.cc:262, queue = 'com.apple.main-thread', stop reason = breakpoint 2.1
frame #0: 0x0000000123c66303 _pywrap_tensorflow.so`::TF_NewSession(opt=0x00007ffcf56f8e20, status=0x00007ffcf510a800) + 19 at c_api.cc:262
259
260 TF_Session* TF_NewSession(const TF_SessionOptions* opt, TF_Status* status) {
261 Session* session;
-> 262 status->status = NewSession(opt->options, &session);
263 if (status->status.ok()) {
264 return new TF_Session({session});
265 } else {
(lldb)

通过这样的方式,我们就可以一步一步的看TensorFlow是怎么运行的了。由于TensorFlow内部封装了很多自定义类型,即使可以单步调试,看起来也是非常苦难。如果可以借助于LLDB的自定义格式化打印的功能,很可能给TensorFlow的调试带来极大的便利。

其他

如果是在macOS上调试自定义的操作时,需要注意,在编译时,要给g++或者clang++加上-D_GLIBCXX_USE_CXX11_ABI=0参数。比如这样:

1
g++ -v -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC -I $TF_INC -O2 -undefined dynamic_lookup -D_GLIBCXX_USE_CXX11_ABI=0

如果是使用bazel进行构建:

1
bazel build -s --copt="-D_GLIBCXX_USE_CXX11_ABI=0" -c opt //tensorflow/core/user_ops:zero_out.so

如果不这样做,在加载这个库的时候,会出现内存分配错误。